Llava Onevision: add model
Author: zucchini-nlp
Pull request: yes
Labels: name: run-slow

# What does this PR do?

Adds llava onevision models (https://arxiv.org/abs/2408.03326) which are basically same as LLaVa-NeXT with have a different vision tower and image grid pinpoints

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_32673). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Thanks for adding this model @zucchini-nlp! 

If the model has its own paper - we should add it as its own model and leverage the diff converter for the model addition if changes are needed 

Agree with @amyeroberts, which is more in line with the [philosophy](https://huggingface.co/blog/transformers-design-philosophy)

@zucchini-nlp You've hardcoded a path in modeling_llava_next.py and I suppose it's for debugging?
```
        if input_ids.shape[1] > 1:
            torch.save(inputs_embeds, "/raid/raushan/inputs_embeds_our.pt")
```

Yes, will remove that. I am making a completely new modeling folder for Onevision currently

@amyeroberts Ready for review. Added a new model LLaVa-Onevision. I didn't use diff because I realized that copying everything from llava-next is not a good decision. The current code incorporates all the recent refactors we've been working on and supports `GenerationTesterMixin` 🎉  Apart from that it can support static cache, though not compile due to dynamic control.

I added two processing files, one for image and another for video, since I've been planning to go over video-image LLMs and separate processing into two files. And later add more utility functions to the `BaseVideoProcessor` class. But in general, there's no difference between them except for that videos can be list of 4D tensors and we need one more loop over frame dimension.



@amyeroberts should be ready! I left the `config vs text_config` when applicable until @gante merged his PR, after that those can be replaced by a call to something like `self.get_config()`

Failing tests are OOM or not related

I tried the script `convert_llava_onevision_weights_to_hf.py`, and somehow even though it explicitly provides images as input to the processor, it calls the `preprocess` method of the `video_processor`, as this error points out:

```
Traceback (most recent call last):
  File "/capstor/scratch/cscs/ndeperr/code/health_mm_llm/health_mm_llm_evaluation/convert_llava_onevision_weights_to_hf.py", line 361, in <module>
    convert_llava_to_hf(args.model_id, args.model_path, args.pytorch_dump_folder_path, args.push_to_hub)
  File "/capstor/scratch/cscs/ndeperr/code/health_mm_llm/health_mm_llm_evaluation/convert_llava_onevision_weights_to_hf.py", line 198, in convert_llava_to_hf
    inputs = processor(images=image, text=prompt, return_tensors="pt").to(torch.float16)
  File "/capstor/scratch/cscs/ndeperr/virtual_env/llava-env/lib/python3.10/site-packages/transformers/models/llava_onevision/processing_llava_onevision.py", line 161, in __call__
    image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"])
  File "/capstor/scratch/cscs/ndeperr/virtual_env/llava-env/lib/python3.10/site-packages/transformers/image_processing_utils.py", line 41, in __call__
    return self.preprocess(images, **kwargs)
  File "/capstor/scratch/cscs/ndeperr/virtual_env/llava-env/lib/python3.10/site-packages/transformers/models/llava_onevision/video_processing_llava_onevision.py", line 626, in preprocess
    videos = make_batched_videos(videos)
  File "/capstor/scratch/cscs/ndeperr/virtual_env/llava-env/lib/python3.10/site-packages/transformers/models/llava_onevision/video_processing_llava_onevision.py", line 65, in make_batched_videos
    elif is_valid_image(videos) and len(videos.shape) == 4:
AttributeError: 'JpegImageFile' object has no attribute 'shape'. Did you mean: 'save'?
```
Is it because of the overlapping name of the `preprocess` method in both image and video processors?
The error is solved when replacing the object `self.video_processor` by `self.image_processor` 


> Failing tests are OOM or not related

If the tests are failing with OOM, this is still an issue which means we can't merge them into main, unfortunately 


Oke, I'll load them in lower precision. Wondering if the GPU memory is cleared after each slow test as the model is only 0.8B parameters and the GPU apparently has 14GB memory

@zucchini-nlp hmmm - that's funny. @ydshieh will know more than me about the mem management here! 

May I know if the video frames are correctly processed (they need to be bilinear interpolated to 196 tokens each frame). I quickly checked the code and didnt see such operations, may wondering if that will cause the OOM problem?

Hey @Luodian , thanks for reviewing the PR! Yes, tye implementation doesn't do interpolation. I saw you guys mentioned it in the paper but the demo notebook didn't lead to the path with interpolation, so I thought it's from llava-next-video. 

Prob it was a typo in demo notebook, passing "modality" arg triggers interpolation. I'll add it then, thanks ;)

> Hey @Luodian , thanks for reviewing the PR! Yes, tye implementation doesn't do interpolation. I saw you guys mentioned it in the paper but the demo notebook didn't lead to the path with interpolation, so I thought it's from llava-next-video.
> 
> Prob it was a typo in demo notebook, passing "modality" arg triggers interpolation. I'll add it then, thanks ;)

Yes, there's mistake on `modality` in our previous notebook demo. Really nice and also thanks to see you made such great contribution to the community.

@amyeroberts done! The tests are passing now, including slow + added the extra pooling layer for videos

> Oke, I'll load them in lower precision. Wondering if the GPU memory is cleared after each slow test as the model is only 0.8B parameters and the GPU apparently has 14GB memory

No, we don't do `torch.cuda.empty_cache` after each call. (And it's not really a magic: it's not for torch itself, it's for other application to get some GPU memory.


Will need some changes in video processing as discussed here (https://github.com/LLaVA-VL/LLaVA-NeXT/issues/144). Comes out the demo notebook was buggy and videos should be processed in a different way. Will work on it when I get reply from authors on what should be the correct way 

@zucchini-nlp Thanks for the update! I'll mark myself a re-requesting for review then as this will be a significant change 

@amyeroberts yes, the tests should be green now. Merging `main` somehow rewinded my local changes

oke, will merge it now as the model works. The only thing I am struggling now is ONNX conversion which is not related to the PR

@zucchini-nlp 

hello any idea how to get a llama model working with this pr. mylesgoose/Llama-3.1-Minitron-4B-Llava-Nvidia-siglip-ov regards myles I have tried adding the similar files to the huggingface repo as per the qwen model which works. there is small issue I think because  the 4b model is a stripped down 8b model with weight sizes. 


@mylesgoose I am not sure if you got the model from tuning using original repo or in another way. If you used the original repo, feel free to convert to HF style using `conver_weights.py` script. Otherwise, you would need to write your own conversion script

We try to reserve GH for issues/bugs and feature requests, feel free to open a thread in the [forum](https://discuss.huggingface.co/latest) for any further questions 🤗 

@zucchini-nlp  no I made the model using your git repo repo llava one vision. it's a llama 4b model. and this is a feature request. I trained it.  because the model is only supported of qwen type models. at present this transformers can't understand  my weight are smaller. I am not sure exactly  how to modify this one vision  transformers to load the llama model type. can you have a look at the huggingface repo.and you will see what I mean.

@mylesgoose I am not sure I understand you. The repo shows a trained model with a different LM backbone and has inference script with llava-vl repo code. So I guess you did training and can do generation with llava-vl repo

In transformers we already support any LM backbone, and all you need is to convert the weights, which has several conversion scripts for different llava models. So this is not a feature request. In case you still believe we can help you by adding/modifying llava code, please open an issue and provide minimal script to see what you want to achieve and what output is expected from running it 🤗 

Okay thanks @zucchini-nlp  I understand now. I found that script and ran one to generate the bin file and that worked. however the script does not work on the llama model. so I will modify it tomorrow and get it to work and then return working adjusted one. I think it also does not add tokens for video etc also. or check if pad token exists prior to creating one. and created meta tensors on devices and could not use in latter step. will figure it out. thanks allot. 


This seems to have worked. @zucchini-nlp

```python
import argparse
import torch
from transformers import (
    AddedToken,
    AutoConfig,
    AutoImageProcessor,
    AutoTokenizer,
    LlavaConfig,
    LlavaForConditionalGeneration,
    LlavaProcessor,
    SiglipVisionConfig,
)

EPILOG_TXT = """Example:
    python convert_llava_weights_to_hf.py \\
      --text_model_id mylesgoose/Llava-Llama-3.1-Minitron-4B-Nvidia-siglip-ov \\
      --vision_model_id google/siglip-so400m-patch14-384 \\
      --output_hub_path mylesgoose/Llama-3.1-Minitron-4B-Llava-Nvidia-siglip-ov-hf \\
      --old_state_dict_id /home/myles/LLaVA-NeXT/tmp/hf_models/llava-Llama-3.1-Minitron-4B/model_state_dict.bin

Example for creating the old state dict file with Python:

    import torch
    from llava.model.language_model.llava_llama import LlavaLlamaForCausalLM

    # Load model
    kwargs = {"device_map": "auto", "torch_dtype": torch.float16}
    model = LlavaLlamaForCausalLM.from_pretrained("mylesgoose/Llava-Llama-3.1-Minitron-4B-Nvidia-siglip-ov", low_cpu_mem_usage=True, **kwargs)

    # Load vision tower
    model.get_vision_tower().load_model()

    # Save state dict
    torch.save(model.state_dict(), "/home/myles/LLaVA-NeXT/tmp/hf_models/llava-Llama-3.1-Minitron-4B/model_state_dict.bin")
"""

KEYS_TO_MODIFY_MAPPING = {
    "model.vision_tower.vision_tower.": "vision_tower.",
    "model.vision_tower.": "vision_tower.",
    ".vision_resampler": "",  # Not used in the original implementation
    "model.mm_projector": "multi_modal_projector",
    "model": "model.model",
    "vision_model.model": "vision_model",
    "lm_head": "language_model.lm_head",
    "model.model": "language_model.model",
    "multi_modal_projector.0": "multi_modal_projector.linear_1",
    "multi_modal_projector.2": "multi_modal_projector.linear_2",
}

def convert_state_dict_to_hf(state_dict):
    new_state_dict = {}
    for key, value in state_dict.items():
        # Skip keys ending with .inv_freq (if present)
        if key.endswith(".inv_freq"):
            continue
        # Modify keys based on the mapping
        for key_to_modify, new_key in KEYS_TO_MODIFY_MAPPING.items():
            if key_to_modify in key:
                key = key.replace(key_to_modify, new_key)
        new_state_dict[key] = value
    return new_state_dict

def convert_llava_llama_to_hf(text_model_id, vision_model_id, output_hub_path, old_state_dict_id):
    torch.set_default_dtype(torch.bfloat16)
    text_config = AutoConfig.from_pretrained(text_model_id)

    tokenizer = AutoTokenizer.from_pretrained(text_model_id)
    tokenizer.add_tokens(AddedToken("<image>", special=True, normalized=False), special_tokens=True)
    tokenizer.add_tokens(AddedToken("<video>", special=True, normalized=False), special_tokens=True)
    if tokenizer.pad_token_id is None:
        tokenizer.add_special_tokens({"pad_token": "
"})

    image_processor = AutoImageProcessor.from_pretrained(vision_model_id)
    processor = LlavaProcessor(tokenizer=tokenizer, image_processor=image_processor)

    # Set vision_config based on the vision model
    vision_config = SiglipVisionConfig.from_pretrained(vision_model_id)

    # Adjust the vision_config to match your model's parameters
    vision_config.hidden_size = 1152
    vision_config.num_attention_heads = 16
    vision_config.num_hidden_layers = 24
    vision_config.intermediate_size = 4304
    vision_config.patch_size = 14
    vision_config.image_size = 384
    vision_config.vision_use_head = False

    config = LlavaConfig(
        text_config=text_config,
        vision_config=vision_config.to_dict(),
    )

    # Set other configuration parameters as needed
    config.pad_token_id = tokenizer.pad_token_id
    config.image_token_index = tokenizer.convert_tokens_to_ids("<image>")
    config.video_token_index = tokenizer.convert_tokens_to_ids("<video>")
    config.vision_feature_layer = -2  # Adjust if necessary
    config.vision_tower_pretrained = vision_model_id

    # Instantiate the model
    model = LlavaForConditionalGeneration(config)

    # Load the state dictionary from the provided path
    state_dict_path = old_state_dict_id
    state_dict = torch.load(state_dict_path, map_location="cpu")

    # Convert state dict to match Hugging Face model's expected keys
    state_dict = convert_state_dict_to_hf(state_dict)

    # Load the state dictionary into the model
    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
    print(f"Missing keys: {missing_keys}")
    print(f"Unexpected keys: {unexpected_keys}")

    # Resize token embeddings to accommodate new tokens
    model.resize_token_embeddings(len(tokenizer))

    # Get the updated sizes
    pre_expansion_embeddings = model.language_model.model.embed_tokens.weight.data
    vocab_size = pre_expansion_embeddings.size(0)
    new_vocab_size = model.language_model.model.embed_tokens.weight.size(0)
    num_new_tokens = new_vocab_size - vocab_size

    print(f"Original vocab size: {vocab_size}")
    print(f"New vocab size: {new_vocab_size}")
    print(f"Number of new tokens: {num_new_tokens}")

    # Initialize new token embeddings only if there are new tokens
    if num_new_tokens > 0:
        mu = torch.mean(pre_expansion_embeddings, dim=0).float()
        n = pre_expansion_embeddings.size(0)
        sigma = ((pre_expansion_embeddings - mu).T @ (pre_expansion_embeddings - mu)) / n
        dist = torch.distributions.multivariate_normal.MultivariateNormal(
            mu.cpu(), covariance_matrix=1e-5 * sigma.cpu()
        )

        new_embeddings = torch.stack(
            [dist.sample() for _ in range(num_new_tokens)], dim=0
        ).to(pre_expansion_embeddings.device)
        model.language_model.model.embed_tokens.weight.data[vocab_size:] = new_embeddings
        model.language_model.lm_head.weight.data[vocab_size:] = new_embeddings
    else:
        print("No new tokens to add to embeddings.")

    # Save the converted model and processor
    model.push_to_hub(output_hub_path)
    processor.push_to_hub(output_hub_path)

def main():
    parser = argparse.ArgumentParser(
        epilog=EPILOG_TXT,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--text_model_id",
        default="mylesgoose/Llava-Llama-3.1-Minitron-4B-Nvidia-siglip-ov",
        help="Hub location of the text model",
    )
    parser.add_argument(
        "--vision_model_id",
        default="google/siglip-so400m-patch14-384",
        help="Hub location of the vision model",
    )
    parser.add_argument(
        "--output_hub_path",
        default="mylesgoose/Llama-3.1-Minitron-4B-Llava-Nvidia-siglip-ov-hf",
        help="Location on the hub of the converted model",
    )
    parser.add_argument(
        "--old_state_dict_id",
        default="/home/myles/LLaVA-NeXT/tmp/hf_models/llava-Llama-3.1-Minitron-4B/model_state_dict.bin",
        help="Path to the raw state dict of the original model. The filename needs to be `model_state_dict.bin`",
    )
    args = parser.parse_args()
    convert_llava_llama_to_hf(
        args.text_model_id,
        args.vision_model_id,
        args.output_hub_path,
        args.old_state_dict_id
    )

if __name__ == "__main__":
    main()
```