MLX & CUDA examples with Vision encoder for MultiModal Model like LLaVA to perform as Visual Instruct Assistant

LLaVA — Large Language and Vision Assistant is an end-to-end trained large multimodal model that connects a vision encoder and a LLM for general-purpose vision and language understanding.

Requirements:

• Apple M1 and above or Nvidia CUDA chipset
• Hugging Face login
• Python 3.10 and above
• Mac OS mlx package. (coremltools mlx)
• Transformer pip package(Hugging Face)

On a MLX Platform these are the requirement…

mlx>=0.8.0
numpy
transformers
torch
huggingface_hub
Pillow

On a CUDA based environment these are the requiremen…

transformers
torch
huggingface_hub
Pillow
jinja2==3.1.0

The LLaVA model consists of the following components…

• Uses open-set visual encoder of CLIP and vision to text decoder of Vicuna and fine tuning end-to-end on the generated instructional vision-language data.
• Large Language Model like GPT-4.
• It representa data reformative and perpective to convert image-text pairs into appropriate instruction-following format.

The result is a Large Multimodal Model(LMM) encompassing all of the above components.

It has achived SoTA on ScienceQA mutlimodal reasoning dataset.

Architecture:

For an input image Xv we use CLIP to get a feature Zv. Then apply a trainable projection matrix W to convert Zv to language embedding tokens Hv, which has the same dimentionality as the word embeddings space in Large Language models.

This is the sequence of Visual Tokens..

Traning:

It is two staged

1. Pre-training for feature allignment : Filter CC3M to 595K image text pairs, these pairs are converted in instruction-following data, thereby we have image features Hv can be aligned with pre-trained LLM word embedding. Training a compatible visual tokenizer for a frozen LLM.
2. Fine-tuning end-to-end : keep the visual encoder weights frozen and continue to update both the pre-trained weights of project-layer and LLM in LLaVA. The 2 use cases are (i) Multimodal Chat Bot and (ii) Scrience QA .

More details can be found in the link in the Reference section.

To test your image you can use a visual-chat mode.

Note: In this article we will test only the Hugging face model which is PyTorch based.

Let us look at the code..

#!/usr/bin/env python3
#author:rangapv@yahoo.com
#27-04-26

from PIL import Image
from transformers import AutoProcessor, LlavaForConditionalGeneration

model = LlavaForConditionalGeneration.from_pretrained("llava-hf/llava-1.5-7b-hf")
processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
file1 = Image.open("/Users/rangaswamypv/rangapv/macOD-ML/mlx/tensorflow/daisy.jpg")

def mesg1():
   input1 = input("USER:")
   m2 = [
       {
        "role": "user",
        "content": [
            {"type": "image", "image": file1},
            {"type": "text", "text": input1}
        ]   
       }
   ]   
   m1 = m2
   return m1

while True:
 message = mesg1()
 inputs = processor.apply_chat_template(
    message,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
    add_generation_prompt=True
    #continue_final_message=True
 )
# Generate
 generate_ids = model.generate(**inputs,max_new_tokens=625)
 result = processor.decode(generate_ids, skip_special_tokens=True)
 result1 = result[0]
 w1 = "ASSISTANT:"
 indx = result1.find(w1)
 result2 = result1[indx:]
 print(result2)

Note: In the above code we are using the generate() method form Hugging Face model, which will take care of all the sampling and decoding the tokens for us, to give the text output. In the subsequent examples we will test the Tensor forward_Pass() methods, sampling() method and do the token_decode() ourselves to generate the output.

The image that I want to decode is this…

Lets run the Chat-code , the output …

rangaswamypv@Rangaswamys-MacBook-Pro llava % ./four.py 
Fetching 3 files: 100%|█████████████████████████████████████████████████████| 3/3 [00:00<00:00, 3102.30it/s]
Download complete: : 0.00B [00:00, ?B/s]                                              | 0/3 [00:00<?, ?it/s]
Loading weights: 100%|██████████████████████████████████████████████████| 686/686 [00:00<00:00, 7391.77it/s]

USER:Describe the image ?
ASSISTANT: The image features a close-up of a white flower with a yellow center, possibly a daisy. The flower 
is positioned in the center of the frame, and its petals are spread out, creating a visually appealing composition. 
The flower appears to be in a dark setting, possibly a black background, which emphasizes its beauty and contrasts 
with the white petals.

USER:What is the color of the stem ?
ASSISTANT: The color of the stem is green.

USER:How many questions have I asked you so far ?
ASSISTANT: You have asked me 1 question so far.

USER:Are there more than 8 petals in the flower ?
ASSISTANT: Yes, there are more than 8 petals in the flower. The image shows a white flower with a yellow center, 
and it is described as having a lot of petals.

USER:Which part of the world can I find these flowers ?
ASSISTANT: You can find these white flowers with yellow centers in various parts of the world, particularly in 
temperate and colder climates. They are commonly found in gardens, fields, and meadows, and they can be found in 
different regions, such as North America, Europe, and Asia. The specific type of flower may vary depending on the region, 
but the general appearance and characteristics of the flower remain consistent.

USER:What are the other colors for this flower ?
ASSISTANT: The other colors for this flower are yellow and white. In the image, the flower is predominantly white 
with a yellow center, which makes it visually appealing and unique.

USER:Does this flower also come in other colors ?
ASSISTANT: Yes, besides the white color seen in the image, daisies also come in various other colors, such as pink, yellow, orange, and red. 
The different colors are a result of the presence of different pigments in the flower's petals and other parts. 
The white daisy in the image is a classic example of the flower's natural beauty, but it is important to note that there are many other 
colorful daisies that can be found in nature.

USER:What is the total number of questions I have asked so far ?
ASSISTANT: You have asked a total of 10 questions so far.

As you can see the model(LLaVA) is not good at math (like counting..etc)

Otherwise it is very good at classification/segmentation/decription/Caption etc…

As a real time image I uploaded an image as I was writing this article. Here I will make it run on GPUs instead of the earlier CPU only. Also we will do the entire Forward-Pass to prompt decode in the code itself instead of relying on the model.generate() like earlier. This is on a MLX (Apple Silicon) platform.

#!/usr/bin/env python3
import torch, requests, io
from PIL import Image
from transformers import AutoProcessor, LlavaForConditionalGeneration

# -------------------------------------------------
# 1️⃣  Model / device
# -------------------------------------------------
model_name = "llava-hf/llava-1.5-7b-hf"
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")

model = LlavaForConditionalGeneration.from_pretrained(
    model_name, torch_dtype=torch.float16
).to(device)
processor = AutoProcessor.from_pretrained(model_name)

# -------------------------------------------------
# 2️⃣  Load image once
# -------------------------------------------------
url = "https://static01.nyt.com/images/2023/07/21/multimedia/21baguettesrex-hbkc/21baguettesrex-hbkc-videoSixteenByNineJumbo1600.jpg"
img = Image.open(io.BytesIO(requests.get(url).content)).convert("RGB")

# -------------------------------------------------
# 3️⃣  Helper functions
# -------------------------------------------------
def mesg1():
    txt = input("USER: ")
    if txt.lower() in {"quit", "exit"}:
        raise KeyboardInterrupt
    return f"USER: <image>\n{txt}\nASSISTANT:"

def sample(logits, temperature=0.0):
    """Pure PyTorch sampling — greedy or temperature."""
    if temperature == 0.0:
        return torch.argmax(logits, dim=-1, keepdim=True)  # (B, 1)
    probs = torch.softmax(logits / temperature, dim=-1)
    return torch.multinomial(probs, num_samples=1)          # (B, 1)

# -------------------------------------------------
# 4️⃣  Interactive loop
# -------------------------------------------------
try:
    while True:
        prompt = mesg1()

        # Processor → PyTorch tensors on device
        inputs = processor(
            text=prompt,
            images=img,
            return_tensors="pt"
        ).to(device)

        input_ids     = inputs["input_ids"]        # (1, seq_len)
        pixel_values  = inputs["pixel_values"]     # (1, C, H, W)
        attention_mask = inputs["attention_mask"]  # (1, seq_len)

        # First forward pass — prompt + image
        with torch.inference_mode():
            out = model(
                input_ids=input_ids,
                pixel_values=pixel_values,
                attention_mask=attention_mask,
                use_cache=True,
            )

        logits = out.logits[:, -1, :]              # (1, vocab)
        cache  = out.past_key_values
        y      = sample(logits, temperature=0.0)   # (1, 1)
        tokens = [int(y.item())]

        # Token-by-token sampling loop
        with torch.inference_mode():
            for _ in range(199):                   # max_new - 1
                # Grow attention mask by 1
                attention_mask = torch.cat(
                    [attention_mask,
                     torch.ones((1, 1), dtype=attention_mask.dtype, device=device)],
                    dim=1
                )
                out = model(
                    input_ids=y,                   # (1, 1)
                    past_key_values=cache,
                    attention_mask=attention_mask,
                    use_cache=True,
                )
                cache  = out.past_key_values
                logits = out.logits[:, -1, :]
                y      = sample(logits, temperature=0.0)
                token  = int(y.item())
                if token == processor.tokenizer.eos_token_id:
                    break
                tokens.append(token)

        result = processor.tokenizer.decode(tokens, skip_special_tokens=True)
        print("\nASSISTANT:", result.strip(), "\n")

except KeyboardInterrupt:
    print("\nBye!")

The image..

And ran the Python code…

And the PROMPT: “USER: <image>\nDescribe the image .. \nASSISTANT:”

rangaswamypv@Rangaswamys-MacBook-Pro llava % ./device.py 
Fetching 3 files: 100%|███████████████████████████████████████████████| 3/3 [00:00<00:00, 4894.17it/s]
Download complete: : 0.00B [00:00, ?B/s]                                        | 0/3 [00:00<?, ?it/s]
Loading weights: 100%|████████████████████████████████████████████| 686/686 [00:00<00:00, 7818.73it/s]

USER:Describe the image ?
ASSISTANT: The image features a white laptop computer sitting on a wooden desk. The laptop is open and 
ready for use. Beside the laptop, there is a cup, possibly containing coffee or tea, and a cell phone. 
The scene appears to be a typical workspace or study area.

USER:How many objects are in the image ?
ASSISTANT: There are three objects in the image: a laptop, a cup, and a cell phone.

USER:What is the total number of questions that I have asked ?
ASSISTANT: You have asked a total of 1 question.

USER:If have to buy a new Apple latop which specification would you recommend under $4000 dollars?
ASSISTANT: If you have to buy a new Apple laptop under $4000, I would recommend the 
MacBook Pro 13-inch (2020) model. This laptop features a powerful 
10th-generation Intel Core processor, up to 16GB of memory, and up to 2TB of 
storage. It also has a Retina display, Touch ID, and a Thunderbolt 3 port for 
connectivity. Additionally, the MacBook Pro 13-inch is known for its 
portability and durability, making it a great choice for both personal and 
professional use.

USER:Is the cup to the left or right of the laptop ?
ASSISTANT: The cup is to the right of the laptop.

Now lets run this code on a CUDA device. Here as well we are not using mode.generate() instead we are coding the entire Forward_pass, Sampling and Promp decode().

#!/usr/bin/env python3

import torch
import requests, io
from PIL import Image
from transformers import AutoProcessor, LlavaForConditionalGeneration

# -------------------------------------------------
# 1️⃣  Model / device setup
# -------------------------------------------------
model_name = "llava-hf/llava-1.5-7b-hf"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = LlavaForConditionalGeneration.from_pretrained(
    model_name,
    torch_dtype=torch.float16,          # half‑precision saves VRAM
    low_cpu_mem_usage=True,
).to(device)

processor = AutoProcessor.from_pretrained(model_name)

# -------------------------------------------------
# 2️⃣  Load the image once (the same image for every turn)
# -------------------------------------------------
url1 = "https://static01.nyt.com/images/2023/07/21/multimedia/21baguettesrex-hbkc/21baguettesrex-hbkc-videoSixteenByNineJumbo1600.jpg"

# -------------------------------------------------
# 3️⃣  Helper functions
# -------------------------------------------------
def mesg1():
    """Read user text and build the multimodal message."""
    txt = input("USER: ")
    if txt.lower() in {"quit", "exit"}:
        raise KeyboardInterrupt
    return [
        {
            "role": "user",
            "content": [
                {"type": "image", "url": url1},
                {"type": "text", "text": txt},
            ],
        }
    ]

def sample(logits, temperature=0.0):
    """Deterministic argmax when temperature==0, otherwise multinomial."""
    if temperature == 0.0:
        return torch.argmax(logits, dim=-1)
    probs = torch.softmax(logits / temperature, dim=-1)
    return torch.multinomial(probs, num_samples=1).squeeze(-1)

# -------------------------------------------------
# 4️⃣  Main interactive loop
# -------------------------------------------------
try:
    while True:
        # ----- Build the prompt -----
        message = mesg1()
        inputs = processor.apply_chat_template(
            message,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
            add_generation_prompt=True,
        ).to(device)

        input_ids = inputs["input_ids"]          # (1, seq_len)
        pixel_values = inputs["pixel_values"]    # (1, 3, H, W)
        attention_mask = inputs["attention_mask"]# (1, seq_len)

        # ----- First forward pass (image + text) -----
        out = model(
            input_ids=input_ids,
            pixel_values=pixel_values,
            attention_mask=attention_mask,
            use_cache=True,
        )
        logits = out.logits[:, -1, :]            # (1, vocab)
        cache = out.past_key_values
        temperature=0.0
        y = sample(logits, temperature)        # (1,)
        tokens = [y.item()]

        # ----- Generation loop -----
        max_new = 200
        for _ in range(max_new - 1):
            # extend attention mask for the newly generated token
            attention_mask = torch.cat(
                [attention_mask, torch.ones_like(y[:, None])], dim=1
            )

            out = model(
                input_ids=y[:, None],            # (1, 1)
                past_key_values=cache,
                attention_mask=attention_mask,
                use_cache=True,
            )
            cache = out.past_key_values
            logits = out.logits[:, -1, :]
            y = sample(logits, temperature)
            token = y.item()
            if token == processor.tokenizer.eos_token_id:
                break
            tokens.append(token)

        # ----- Decode & clean up the answer -----
        result = processor.decode(tokens, skip_special_tokens=True)
        # Optional: strip a leading "ASSISTANT:" prefix
        prefix = "ASSISTANT:"
        if prefix in result:
            result = result[result.find(prefix) + len(prefix):].strip()
        print("\nASSISTANT:", result, "\n")
except KeyboardInterrupt:
    print("\nSession ended.")

The image I used is .. https://static01.nyt.com/images/2023/07/21/multimedia/21baguettesrex-hbkc/21baguettesrex-hbkc-videoSixteenByNineJumbo1600.jpg

Successfully installed jinja2-3.1.0
ubuntu@ip-172-31-38-0:~/mac/mlx/more-examples/llava$ ./device2.py
Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
Fetching 3 files: 100%|██████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:00<00:00, 45425.68it/s]
Download complete: : 0.00B [00:00, ?B/s]                                                                                | 0/3 [00:00<?, ?it/s]
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████| 686/686 [00:00<00:00, 5681.05it/s]

USER:Describe the image ?
ASSISTANT: The image features a large stack of freshly baked bread, with each loaf wrapped in a paper bag. The bread is arranged in a neat pile, 
showcasing the variety of loaves. The loaves are of different sizes and shapes, indicating a diverse selection of breads. The scene is set against 
a backdrop of a dining table, emphasizing the abundance of bread.

USER:How many breads are there in the image ?
ASSISTANT: There are three breads in the image.

USER:What are the ingredients needed to make this bread ?
ASSISTANT: To make this bread, the ingredients needed include flour, water, yeast, salt, and sometimes sugar and oil. The dough is mixed and 
kneaded until it becomes smooth and elastic. It is then allowed to rise until it doubles in size, and then it is shaped into the desired form, 
such as the long, twisted loaves seen in the image. The bread is then baked in an oven until it reaches the desired color and texture. The combination 
of these ingredients and the proper baking process results in the delicious, freshly baked bread seen in the image.

USER:How many questions have I asked ?
ASSISTANT: You have asked three questions.

USER:Which part of the world eat this bread ?
ASSISTANT: The bread in the image is a baguette, which is a type of French bread. Baguettes are typically associated with French cuisine and are 
commonly consumed in France and other French-speaking countries.

USER:What are some of the pairing for this bread ?
ASSISTANT: The bread in the image is a long, crusty loaf, which can be paired with various dishes and ingredients. Some potential pairings include:

1. Cheese: The crusty texture of the bread can complement the richness of cheese, making it an ideal choice for a cheese platter or a cheese-based appetizer.

2. Meats: The crusty bread can be used to make sandwiches with meats like ham, salami, or turkey, providing a contrasting texture and flavor to the meat.

3. Vegetables: The bread can be paired with a variety of vegetables, such as roasted or grilled vegetables, to create a flavorful and satisfying dish.

4. Dips: The crusty bread can be served with dips like hummus, tapenade, or pesto, adding a del

USER:To make the bread in the image for 2 people give me the ingredients list
ASSISTANT: To make the bread in the image for 2 people, you would need the following ingredients:

* 2 cups of flour
* 1 teaspoon of salt
* 1 teaspoon of sugar
* 1 cup of warm water
* 1/2 cup of olive oil
* 1 package of active dry yeast (2 1/4 teaspoons)
* 1/2 cup of sugar (optional, for the yeast)

Follow the recipe instructions to prepare the dough, shape it into loaves, and let it rise before baking. This will result in two loaves of bread, 
which can be enjoyed by two people.

BINGO!

If you have any issues, you can open an issue in the github repo page or alternatively contact me at rangapv@yahoo.com and you can also find me on X.com(twitter) @rangapv . My linkedin ***https://linkedin.com/in/rangapv***

References:

[embed]Visual Instruction Tuning Instruction tuning large language models (LLMs) using machine-generated instruction-following data has improved…arxiv.org

[embed]Paper page - Visual Instruction Tuning Join the discussion on this paper pagehuggingface.co

Some of my other AI/ML related works

[embed]Source Install of TensorRT_LLM In this article we will Build tensorRT_LLM from source. In the previous article we covered the Build dependencies…rangapv.medium.com

[embed]TensorRT install for Deep Neural Network(DNN) inference In this article we will look at the setup of the TensorRT for inferencing a deep neural network. The goal is to have a…rangapv.medium.com

[embed]GPU based Decision Optimization to solve Linear Programming using cuOpt: Calculate CVaR and VaR GPU based Decision Optimization to solve Linear Programming using cuOpt: Calculate CVaR and VaR In this article we will…rangapv.medium.com

[embed]Deploying Ray Cluster for AI/ML workloads on a Kubernetes Cluster In this article we are going to install a Ray operator in a Kubernetes cluster. We have looked into setting up a…rangapv.medium.com