fix: unboundlocalerror when trying to retrieve attention feature maps using Qwen2FlashAttention2().forward()
Author: EigenTom
Pull request: yes

# What does this PR do?

Fix "unboundlocalerror" occurrs when trying to retrieve attention layer feature maps for models using `Qwen2FlashAttention2().forward()` to inference such as LLaVA-NeXT:

<img width="906" alt="image" src="https://github.com/user-attachments/assets/da91f90d-60f4-4f23-9d65-f571c3169fdf">


The main issue is in this method, the logic for computing `attn_weights` when flag `output_attentions` is set to `True` was not implemented, therefore leaving variable `atten_weights` unbounded.



Fixes # (issue)

Following code demonstrated in Qwen2Attention(nn.Module).forward(), several lines were added to compute the value of variable `atten_weights` when `output_attentions` is set to `True`:

<img width="1387" alt="image" src="https://github.com/user-attachments/assets/d85dea7d-6188-4f46-a38d-0896b0abbf09">



## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts, @qubvel
- speech models: @ylacombe, @eustlb
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
Replication Code:

~~~python3
import sys
sys.path.append(<path_to_LLaVA-NeXT>)
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

from llava.model.builder import load_pretrained_model
from llava.mm_utils import get_model_name_from_path, process_images, tokenizer_image_token
from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, IGNORE_INDEX
from llava.conversation import conv_templates, SeparatorStyle
from PIL import Image
import requests
import copy
import torch

import sys
import warnings

class LLAVANeXTService:
    def __init__(self, pretrained_model="lmms-lab/LLaVA-Video-7B-Qwen2", device="cuda:0", device_map="auto"):
        # Define the model loading arguments
        self.llava_model_args = {
            "multimodal": True,
        }
        overwrite_config = {}
        overwrite_config["image_aspect_ratio"] = "pad"
        self.llava_model_args["overwrite_config"] = overwrite_config

        # Load the tokenizer, model, and image processor
        self.tokenizer, self.model, self.image_processor, self.max_length = load_pretrained_model(
            pretrained_model, 
            None, 
            "llava_qwen", 
            device_map=device_map, 
            **self.llava_model_args,
            # attn_implementation="sdpa"
        )

        # Set model to evaluation mode
        self.model.eval()

        # Define the device (CUDA or CPU)
        self.device = torch.device(device)

    def predict(self, prompt, img_list,):
        # Load and process multiple images
        images = [Image.open(img_path) for img_path in img_list]
        image_tensors = process_images(images, self.image_processor, self.model.config)
        image_tensors = [_image.to(dtype=torch.float16, device=self.device) for _image in image_tensors]
        image_sizes = [image.size for image in images]

        # Prepare interleaved text-image input
        conv_template = "qwen_1_5"
        conv = copy.deepcopy(conv_templates[conv_template])
        
        prompt = f"Provided image is: {DEFAULT_IMAGE_TOKEN} \n" + prompt
        
        conv.append_message(conv.roles[0], prompt)
        conv.append_message(conv.roles[1], None)
        prompt_question = conv.get_prompt()

        # Tokenize the text prompt with image tokens
        inputs = tokenizer_image_token(prompt_question, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt")
        input_ids = inputs.unsqueeze(0).to(self.device)

        # Generate the response from the model
        cont = self.model.generate(
            input_ids,
            images=image_tensors,
            image_sizes=image_sizes,
            do_sample=False,
            temperature=0,
            max_new_tokens=4096,
            return_dict_in_generate=True,
            output_attentions=True,
        )

        # Decode the output
        text_outputs = self.tokenizer.batch_decode(cont["sequences"], skip_special_tokens=True)
        # text_outputs = self.tokenizer.batch_decode(cont, skip_special_tokens=True)
        return text_outputs[0], cont

service = LLAVANeXTService()
text_output, raw_output = service.predict("what is the most salient object in the image", [<path_to_img>])
~~~