fix bug: fastervlm for llavahf

[zhangbilang]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @SmudgedWings, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses several bugs within the fastervlm module to improve its compatibility and correctness when used with llavahf (Llava Hugging Face) models. The changes focus on ensuring proper handling of model inputs, including image attentions, attention masks, and position IDs, and correcting the registration of a critical forward hook to ensure the model processes inputs as expected.

Highlights

• Input Handling for LLaVA-HF: Introduced a specific check to ensure image_attentions are correctly populated from image_attentions_list for llavahf models, particularly when processing with a batch size of 1. This ensures the visual attention data is properly available for subsequent operations within the update_attentions_hook.
• Accurate Attention and Position ID Propagation: Modified the prepare_inputs_for_llm_hook to correctly slice and reuse the original attention_mask and position_ids based on the keep_indexs. Previously, these were generically re-generated, which could lead to incorrect contextual information for the model. This change ensures that the precise attention and positional information is maintained after token reduction.
• Correct Hook Registration Path: Adjusted the module path for registering the prepare_inputs_for_llm_hook by adding an extra .model attribute. This structural fix ensures the hook is correctly applied to the intended part of the llavahf model's forward pass, resolving an issue with hook placement.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This PR fixes a bug for FasterVLM with LlavaHf models. The changes adjust how attention masks and position IDs are handled, and where a forward hook is registered. The review identifies a potential error with attention handling and a maintainability issue with deeply nested attribute access.

[gemini-code-assist]

This logic introduces a potential IndexError. Accessing pruning_paras['image_attentions_list'][0] will raise an exception if the list is empty, which can happen with text-only inputs where no image attentions are generated. A more robust implementation should handle this issue.

if 'image_attentions' not in pruning_paras and pruning_paras.get('image_attentions_list'):
    pruning_paras['image_attentions'] = pruning_paras['image_attentions_list'][-1]

[gemini-code-assist]

Consider using the original attention_mask and position_ids tensors with the keep_indexs to maintain the correct attention patterns and positional information after pruning. This ensures that the model attends to the correct tokens and positions after the visual tokens have been reduced.

new_attention_mask = kwargs['attention_mask'][:, keep_indexs]
new_position_ids = kwargs['position_ids'][:, keep_indexs]

[gemini-code-assist]

Consider simplifying the nested attribute access self.model.model.model to improve code readability and maintainability. Accessing nested attributes can be fragile and tightly couples the code to the specific internal structure of the LlavaHf and Hugging Face Llama models.