    """
    Img2Prompt_VQA model consists of three submodels for zero-shot VQA:
        1. Image-questioning matching model
        2. Image captioning model
        3. Large Language model

    Supported model types:
        - base: BLIPITM, BLIPCaption, PNPUnifiedQAv2FiD (t5-base)
        - large: BLIPITM, BLIPCaption, PNPUnifiedQAv2FiD (t5-large)
        - 3b: BLIPITM, BLIPCaption, PNPUnifiedQAv2FiD (t5-3b)

    Usage:
        >>> from lavis.models import load_model
        >>> model = load_model("img2prompt_vqa", "base", is_eval=True)
    """

    PRETRAINED_MODEL_CONFIG_DICT = {
        "base": "configs/models/img2prompt-vqa/img2prompt_vqa_base.yaml",
    }

    def __init__(
        self,
        image_question_matching_model,
        image_captioning_model,
        question_generation_model,
        question_generation_tokenizer,
        offload_model=False,
    ):
        super().__init__()

        self.image_question_matching_model = image_question_matching_model
        self.image_captioning_model = image_captioning_model
        self.question_generation_model = question_generation_model
        self.question_generation_tokenizer = question_generation_tokenizer
        self.offload_model = offload_model
        self.nlp = spacy.load("en_core_web_sm")

    def forward_itm(self, samples, block_num=7):
        """
        Args:
            samples (dict): A dictionary containing the following keys:
                - image (torch.Tensor): A tensor of shape (batch_size, 3, H, W)
                - text_input (list): A list of strings of length batch_size
            block_num (int): The index of cross-attention block for gradcam computation.

        Returns:
            samples (dict): A dictionary containing the following keys:
                - image (torch.Tensor): A tensor of shape (batch_size, 3, H, W)
                - text_input (list): A list of strings of length batch_size
                - gradcams (torch.Tensor): A tensor of shape (batch_size, H*W)
        """
        image = samples["image"]
        question = [text.strip("?") for text in samples["text_input"]]
        tokenized_text = self.image_question_matching_model.tokenizer(
            question, padding="longest", truncation=True, return_tensors="pt"
        ).to(self.image_question_matching_model.device)


        "base": "configs/models/img2prompt-vqa/img2prompt_vqa_base.yaml",
    }

    def __init__(
        self,
        image_question_matching_model,
        image_captioning_model,
        question_generation_model,
        question_generation_tokenizer,
        offload_model=False,
    ):
        super().__init__()

        self.image_question_matching_model = image_question_matching_model
        self.image_captioning_model = image_captioning_model
        self.question_generation_model = question_generation_model
        self.question_generation_tokenizer = question_generation_tokenizer
        self.offload_model = offload_model
        self.nlp = spacy.load("en_core_web_sm")

    def forward_itm(self, samples, block_num=7):
        """
        Args:
            samples (dict): A dictionary containing the following keys:
                - image (torch.Tensor): A tensor of shape (batch_size, 3, H, W)
                - text_input (list): A list of strings of length batch_size
            block_num (int): The index of cross-attention block for gradcam computation.

        Returns:
            samples (dict): A dictionary containing the following keys:
                - image (torch.Tensor): A tensor of shape (batch_size, 3, H, W)
                - text_input (list): A list of strings of length batch_size
                - gradcams (torch.Tensor): A tensor of shape (batch_size, H*W)
        """
        image = samples["image"]
        question = [text.strip("?") for text in samples["text_input"]]
        tokenized_text = self.image_question_matching_model.tokenizer(
            question, padding="longest", truncation=True, return_tensors="pt"
        ).to(self.image_question_matching_model.device)
        with torch.set_grad_enabled(True):
            gradcams, _ = compute_gradcam(
                model=self.image_question_matching_model,
                visual_input=image,
                text_input=question,
                tokenized_text=tokenized_text,
                block_num=block_num,
            )

        gradcams = [gradcam_[1] for gradcam_ in gradcams]
        samples["gradcams"] = torch.stack(gradcams).reshape(
            samples["image"].size(0), -1
        )

        return samples

    def itm_rank(self, image_embeds, image_atts, encoder_input_ids, match_head="itm"):
        # breakpoint()
        encoder_input_ids = encoder_input_ids.clone()
        encoder_input_ids = encoder_input_ids[:, self.prompt_length - 1 :]
        text_attention_mask = (encoder_input_ids != self.tokenizer.pad_token_id).long()

        if match_head == "itm":
            # encoder_input_ids = encoder_input_ids.clone()
            encoder_input_ids[:, 0] = self.tokenizer.enc_token_id
            output = self.text_encoder(
                encoder_input_ids,
                attention_mask=text_attention_mask,
                encoder_hidden_states=image_embeds,
                encoder_attention_mask=image_atts,
                return_dict=True,
            )
            itm_output = self.itm_head(output.last_hidden_state[:, 0, :])
            return itm_output  # , mask, token_length


        question = [text.strip("?") for text in samples["text_input"]]
        tokenized_text = self.image_question_matching_model.tokenizer(
            question, padding="longest", truncation=True, return_tensors="pt"
        ).to(self.image_question_matching_model.device)
        with torch.set_grad_enabled(True):
            gradcams, _ = compute_gradcam(
                model=self.image_question_matching_model,
                visual_input=image,
                text_input=question,
                tokenized_text=tokenized_text,
                block_num=block_num,
            )

        gradcams = [gradcam_[1] for gradcam_ in gradcams]
        samples["gradcams"] = torch.stack(gradcams).reshape(
            samples["image"].size(0), -1
        )

        return samples

    def itm_rank(self, image_embeds, image_atts, encoder_input_ids, match_head="itm"):
        # breakpoint()
        encoder_input_ids = encoder_input_ids.clone()
        encoder_input_ids = encoder_input_ids[:, self.prompt_length - 1 :]
        text_attention_mask = (encoder_input_ids != self.tokenizer.pad_token_id).long()

        if match_head == "itm":
            # encoder_input_ids = encoder_input_ids.clone()
            encoder_input_ids[:, 0] = self.tokenizer.enc_token_id
            output = self.text_encoder(
                encoder_input_ids,
                attention_mask=text_attention_mask,
                encoder_hidden_states=image_embeds,
                encoder_attention_mask=image_atts,
                return_dict=True,
            )
            itm_output = self.itm_head(output.last_hidden_state[:, 0, :])
            return itm_output  # , mask, token_length

        elif match_head == "itc":
            encoder_input_ids[:, 0] = self.tokenizer.cls_token_id
            text_output = self.text_encoder(
                encoder_input_ids,
                attention_mask=text_attention_mask,
                return_dict=True,
                mode="text",
            )
            image_feat = F.normalize(self.vision_proj(image_embeds[:, 0, :]), dim=-1)
            text_feat = F.normalize(
                self.text_proj(text_output.last_hidden_state[:, 0, :]), dim=-1
            )

            sim = image_feat @ text_feat.t()
            return sim

    def forward_cap(
        self,
        samples,
        cap_max_length=20,
        cap_min_length=0,
        top_p=1,
        top_k=50,
        repetition_penalty=1.0,
        num_captions=100,
        num_patches=20,
    ):
        """
        Args:
            samples (dict): A dictionary containing the following keys:
                - image (torch.Tensor): A tensor of shape (batch_size, 3, H, W)
                - text_input (list): A list of strings of length batch_size
                - gradcams (torch.Tensor): A tensor of shape (batch_size, H*W)
            cap_max_length (int): The maximum length of the caption to be generated.
            cap_min_length (int): The minimum length of the caption to be generated.


            )
            itm_output = self.itm_head(output.last_hidden_state[:, 0, :])
            return itm_output  # , mask, token_length

        elif match_head == "itc":
            encoder_input_ids[:, 0] = self.tokenizer.cls_token_id
            text_output = self.text_encoder(
                encoder_input_ids,
                attention_mask=text_attention_mask,
                return_dict=True,
                mode="text",
            )
            image_feat = F.normalize(self.vision_proj(image_embeds[:, 0, :]), dim=-1)
            text_feat = F.normalize(
                self.text_proj(text_output.last_hidden_state[:, 0, :]), dim=-1
            )

            sim = image_feat @ text_feat.t()
            return sim

    def forward_cap(
        self,
        samples,
        cap_max_length=20,
        cap_min_length=0,
        top_p=1,
        top_k=50,
        repetition_penalty=1.0,
        num_captions=100,
        num_patches=20,
    ):
        """
        Args:
            samples (dict): A dictionary containing the following keys:
                - image (torch.Tensor): A tensor of shape (batch_size, 3, H, W)
                - text_input (list): A list of strings of length batch_size
                - gradcams (torch.Tensor): A tensor of shape (batch_size, H*W)
            cap_max_length (int): The maximum length of the caption to be generated.
            cap_min_length (int): The minimum length of the caption to be generated.
            top_p (float): The cumulative probability for nucleus sampling.
            top_k (float): The number of the highest probability tokens for top-k sampling.
            repetition_penalty (float): The parameter for repetition penalty. 1.0 means no penalty.
            num_captions (int): Number of captions generated for each image.
            num_patches (int): Number of patches sampled for each image.

        Returns:
            samples (dict): A dictionary containing the following keys:
                - image (torch.Tensor): A tensor of shape (batch_size, 3, H, W)
                - text_input (list): A list of strings of length batch_size
                - gradcams (torch.Tensor): A tensor of shape (batch_size, H*W)
                - captions (nested list): A nested list of strings of total length batch_size * num_captions
        """
        encoder_out = self.image_captioning_model.forward_encoder(samples)
        captions = [[] for _ in range(encoder_out.size(0))]

        min_num_captions = 0

        while min_num_captions < num_captions:
            encoder_out_samples = []
            for i in range(num_captions):
                patch_id = (
                    torch.multinomial(
                        samples["gradcams"].to(self.image_captioning_model.device),
                        num_patches,
                    ).reshape(encoder_out.size(0), -1)
                    + 1
                )
                patch_id = (
                    patch_id.sort(dim=1)
                    .values.unsqueeze(-1)
                    .expand(-1, -1, encoder_out.size(2))
                )
                encoder_out_sample = torch.gather(encoder_out, 1, patch_id)
                encoder_out_samples.append(encoder_out_sample)

            stacked = torch.stack(encoder_out_samples, dim=1)
            image_embeds = torch.flatten(
                stacked, start_dim=0, end_dim=1
            )  # (bsz*num_seq, num_patch, dim)

            image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(
                self.image_captioning_model.device
            )
            model_kwargs = {
                "encoder_hidden_states": image_embeds,
                "encoder_attention_mask": image_atts,
            }

            prompt = [self.image_captioning_model.prompt] * image_embeds.size(0)
            prompt = self.image_captioning_model.tokenizer(
                prompt, return_tensors="pt"
            ).to(self.image_captioning_model.device)
            prompt.input_ids[:, 0] = self.image_captioning_model.tokenizer.bos_token_id
            prompt.input_ids = prompt.input_ids[:, :-1]

            decoder_out = self.image_captioning_model.text_decoder.generate(
                input_ids=prompt.input_ids,
                max_length=cap_max_length,
                min_length=cap_min_length,
                do_sample=True,
                top_p=top_p,
                top_k=top_k,
                num_return_sequences=1,
                eos_token_id=self.image_captioning_model.tokenizer.sep_token_id,
                pad_token_id=self.image_captioning_model.tokenizer.pad_token_id,
                repetition_penalty=repetition_penalty,
                **model_kwargs
            )

            itm_outputs = self.image_question_matching_model.itm_rank(
                image_embeds, image_atts, encoder_input_ids=decoder_out
            )  # caption filter

            outputs = self.image_captioning_model.tokenizer.batch_decode(
                decoder_out, skip_special_tokens=True
            )

            for counter, output in enumerate(outputs):
                ind = counter // num_captions
                if len(captions[ind]) < num_captions:
                    caption = output[len(self.image_captioning_model.prompt) :]
                    overlap_caption = [1 for caps in captions[ind] if caption in caps]
                    # print(itm_outputs)
                    if (
                        len(overlap_caption) == 0 and itm_outputs[counter] >= 0.5
                    ):  # image filter
                        captions[ind].append(caption)

            min_num_captions = min([len(i) for i in captions])

        samples["captions"] = captions

        return samples

    def answer_extraction(self, caption, num_question_generation=30):
        cap_use = ""
        # print(caption)
        caption = caption
        ans_to_cap_dict = {}
        answers = []
        for cap_idx, cap in enumerate(caption):
            # print(cap)
            cap_use += cap
            cap = cap.strip().strip(".")
            # print(cap)
            cap = self.nlp(cap)
            for token in cap:  # Noun /Verb/Adj//NUM
                if token.pos_ in open_pos:
                    if token.text.lower() not in ans_to_cap_dict:
                        ans_to_cap_dict[token.text.lower()] = [cap_idx]
                    else:
                        if cap_idx not in ans_to_cap_dict[token.text.lower()]:
                            ans_to_cap_dict[token.text.lower()].append(cap_idx)


                decoder_out, skip_special_tokens=True
            )

            for counter, output in enumerate(outputs):
                ind = counter // num_captions
                if len(captions[ind]) < num_captions:
                    caption = output[len(self.image_captioning_model.prompt) :]
                    overlap_caption = [1 for caps in captions[ind] if caption in caps]
                    # print(itm_outputs)
                    if (
                        len(overlap_caption) == 0 and itm_outputs[counter] >= 0.5
                    ):  # image filter
                        captions[ind].append(caption)

            min_num_captions = min([len(i) for i in captions])

        samples["captions"] = captions

        return samples

    def answer_extraction(self, caption, num_question_generation=30):
        cap_use = ""
        # print(caption)
        caption = caption
        ans_to_cap_dict = {}
        answers = []
        for cap_idx, cap in enumerate(caption):
            # print(cap)
            cap_use += cap
            cap = cap.strip().strip(".")
            # print(cap)
            cap = self.nlp(cap)
            for token in cap:  # Noun /Verb/Adj//NUM
                if token.pos_ in open_pos:
                    if token.text.lower() not in ans_to_cap_dict:
                        ans_to_cap_dict[token.text.lower()] = [cap_idx]
                    else:
                        if cap_idx not in ans_to_cap_dict[token.text.lower()]:
                            ans_to_cap_dict[token.text.lower()].append(cap_idx)
                    answers.append(token.text)
            for ent in cap.ents:

                if ent.text not in answers:
                    if ent.text.lower() not in ans_to_cap_dict:
                        ans_to_cap_dict[ent.text.lower()] = [cap_idx]
                    else:
                        if cap_idx not in ans_to_cap_dict[ent.text.lower()]:
                            ans_to_cap_dict[ent.text.lower()].append(cap_idx)
                    answers.append(ent.text)
            for chunk in cap.noun_chunks:
                if len(chunk.text.split()) < 4:
                    if chunk.text.lower() not in ans_to_cap_dict:
                        ans_to_cap_dict[chunk.text.lower()] = [cap_idx]
                    else:
                        if cap_idx not in ans_to_cap_dict[chunk.text.lower()]:
                            ans_to_cap_dict[chunk.text.lower()].append(cap_idx)
                    #                 print(chunk.text)
                    answers.append(chunk.text)
        answers = sorted(answers, key=answers.count, reverse=True)
        real_answers = []
        for i in answers:
            i = i + "."
            if i not in real_answers:
                real_answers.append(i)

        contexts_for_question_generation = []
        answers = []
        for ans in real_answers[
            :num_question_generation
        ]:  # Generate questions for 30 answers with max frequencies.
            contexts_for_question_generation.append(
                "answer: %s  context: %s." % (ans, cap_use)
            )
            answers.append(ans)
        contexts_for_question_generation.append(
            "answer: %s  context: %s." % ("yes.", cap_use)
        )
        answers.append("yes.")
        return contexts_for_question_generation, answers, ans_to_cap_dict

    def forward_qa_generation(self, samples):
        caption = samples["captions"][0]
        (
            contexts_for_question_generation,
            answers,
            ans_to_cap_dict,
        ) = self.answer_extraction(caption)
        inputs = self.question_generation_tokenizer(
            contexts_for_question_generation,
            padding="longest",
            truncation=True,
            max_length=2048,
            return_tensors="pt",
        ).to(self.device)
        question_size = inputs.input_ids.shape[0]
        cur_b = 0
        true_input_size = 10
        outputs_list = []
        while cur_b < question_size:


        for i in answers:
            i = i + "."
            if i not in real_answers:
                real_answers.append(i)

        contexts_for_question_generation = []
        answers = []
        for ans in real_answers[
            :num_question_generation
        ]:  # Generate questions for 30 answers with max frequencies.
            contexts_for_question_generation.append(
                "answer: %s  context: %s." % (ans, cap_use)
            )
            answers.append(ans)
        contexts_for_question_generation.append(
            "answer: %s  context: %s." % ("yes.", cap_use)
        )
        answers.append("yes.")
        return contexts_for_question_generation, answers, ans_to_cap_dict

    def forward_qa_generation(self, samples):
        caption = samples["captions"][0]
        (
            contexts_for_question_generation,
            answers,
            ans_to_cap_dict,
        ) = self.answer_extraction(caption)
        inputs = self.question_generation_tokenizer(
            contexts_for_question_generation,
            padding="longest",
            truncation=True,
            max_length=2048,
            return_tensors="pt",
        ).to(self.device)
        question_size = inputs.input_ids.shape[0]
        cur_b = 0
        true_input_size = 10
        outputs_list = []
        while cur_b < question_size:
            outputs = self.question_generation_model.generate(
                input_ids=inputs.input_ids[cur_b : cur_b + true_input_size],
                attention_mask=inputs.attention_mask[cur_b : cur_b + true_input_size],
                num_beams=3,
                max_length=30,
            )
            questions = self.question_generation_tokenizer.batch_decode(
                outputs, skip_special_tokens=True
            )
            outputs_list += questions
            cur_b += true_input_size
        questions = outputs_list
        samples["questions"] = questions
        samples["answers"] = answers
        samples["ans_to_cap_dict"] = ans_to_cap_dict
        # results.append({"question_id": ques_id, "question":questions,"answer":answers})
        return samples

    def create_context_prompt(self, samples, num_caps_per_img=30):
        ans_dict_queid = samples["ans_to_cap_dict"]
        # print(ans_dict_queid)
        caption = samples["captions"][0]
        answers = samples["answers"]
        Context_Prompt = ""
        mycontexts_id = []
        for idx in range(num_caps_per_img):
            cap_id_list = ans_dict_queid.get(
                answers[(len(answers) - 1 - idx) % len(answers)][:-1].lower(), [0]
            )
            for cap_id in cap_id_list:
                if cap_id not in mycontexts_id:
                    Context_Prompt += caption[cap_id]
                    mycontexts_id.append(cap_id)
                    break  # We just take one cap for each answer
        samples["Context_Prompt"] = Context_Prompt
        return Context_Prompt


        outputs_list = []
        while cur_b < question_size:
            outputs = self.question_generation_model.generate(
                input_ids=inputs.input_ids[cur_b : cur_b + true_input_size],
                attention_mask=inputs.attention_mask[cur_b : cur_b + true_input_size],
                num_beams=3,
                max_length=30,
            )
            questions = self.question_generation_tokenizer.batch_decode(
                outputs, skip_special_tokens=True
            )
            outputs_list += questions
            cur_b += true_input_size
        questions = outputs_list
        samples["questions"] = questions
        samples["answers"] = answers
        samples["ans_to_cap_dict"] = ans_to_cap_dict
        # results.append({"question_id": ques_id, "question":questions,"answer":answers})
        return samples

    def create_context_prompt(self, samples, num_caps_per_img=30):
        ans_dict_queid = samples["ans_to_cap_dict"]
        # print(ans_dict_queid)
        caption = samples["captions"][0]
        answers = samples["answers"]
        Context_Prompt = ""
        mycontexts_id = []
        for idx in range(num_caps_per_img):
            cap_id_list = ans_dict_queid.get(
                answers[(len(answers) - 1 - idx) % len(answers)][:-1].lower(), [0]
            )
            for cap_id in cap_id_list:
                if cap_id not in mycontexts_id:
                    Context_Prompt += caption[cap_id]
                    mycontexts_id.append(cap_id)
                    break  # We just take one cap for each answer
        samples["Context_Prompt"] = Context_Prompt
        return Context_Prompt

    def create_task_prompt(
        self, samples, question_type="neural", num_question_per_img=30
    ):
        syn_question_queid = samples["questions"]
        syn_ans_queid = samples["answers"]
        Task_Prompt = ""
        for idx in range(num_question_per_img):
            # if config['random_question']:
            #     qa_idx = random.randint(0, len(syn_question_queid) - 1)
            # else:
            qa_idx = idx
            if (
                question_type != "rule" and num_question_per_img > 0 and idx < 1
            ):  ## yes and no questions for vqav2
                # Task_Prompt += "Question:"
                # Task_Prompt += syn_question_queid_next[-1]
                # Task_Prompt += '\n'
                # Task_Prompt += "Answer:no\n"
                Task_Prompt += "Question:"


    def create_context_prompt(self, samples, num_caps_per_img=30):
        ans_dict_queid = samples["ans_to_cap_dict"]
        # print(ans_dict_queid)
        caption = samples["captions"][0]
        answers = samples["answers"]
        Context_Prompt = ""
        mycontexts_id = []
        for idx in range(num_caps_per_img):
            cap_id_list = ans_dict_queid.get(
                answers[(len(answers) - 1 - idx) % len(answers)][:-1].lower(), [0]
            )
            for cap_id in cap_id_list:
                if cap_id not in mycontexts_id:
                    Context_Prompt += caption[cap_id]
                    mycontexts_id.append(cap_id)
                    break  # We just take one cap for each answer
        samples["Context_Prompt"] = Context_Prompt
        return Context_Prompt

    def create_task_prompt(
        self, samples, question_type="neural", num_question_per_img=30
    ):
        syn_question_queid = samples["questions"]
        syn_ans_queid = samples["answers"]
        Task_Prompt = ""
        for idx in range(num_question_per_img):
            # if config['random_question']:
            #     qa_idx = random.randint(0, len(syn_question_queid) - 1)
            # else:
            qa_idx = idx
            if (
                question_type != "rule" and num_question_per_img > 0 and idx < 1
            ):  ## yes and no questions for vqav2
                # Task_Prompt += "Question:"
                # Task_Prompt += syn_question_queid_next[-1]
                # Task_Prompt += '\n'
                # Task_Prompt += "Answer:no\n"
                Task_Prompt += "Question:"
                Task_Prompt += syn_question_queid[-1]
                Task_Prompt += "\n"
                Task_Prompt += "Answer:"
                Task_Prompt += "yes\n"
                Task_Prompt += "Question:Is this a toilet?\n"
                Task_Prompt += "Answer:no\n"
            if "question_type" == "rule":  # Rule-Based Question Generation
                Noun_Questions = [
                    "What item is this in this picture?",
                    "What item is that in this picture?",
                ]

                Verb_Questions = [
                    "What action is being done in this picture?",
                    "Why is this item doing in this picture?",
                    "Which action is being taken in this picture?",
                    "What action is item doing in this picture?",
                    "What action is item performing in this picture?",
                ]

                Adj_Questions = [
                    "How to describe one item in this picture?",
                    "What is item's ADJ TYPE in this picture?",
                    "What is the ADJ TYPE in this picture?",
                ]

                Task_Prompt += "Question:"
                doc = self.nlp(syn_ans_queid[(qa_idx) % len(syn_ans_queid)][:-1].lower())
                if doc[-1].pos_ == "NOUN":
                    Task_Prompt += Noun_Questions[
                        random.randint(0, len(Noun_Questions) - 1)
                    ]
                elif doc[-1].pos_ == "VERB":
                    Task_Prompt += Verb_Questions[
                        random.randint(0, len(Verb_Questions) - 1)
                    ]
                elif doc[-1].pos_ == "ADJ":
                    Task_Prompt += Adj_Questions[
                        random.randint(0, len(Adj_Questions) - 1)
                    ]

                Task_Prompt += "\n"

                Task_Prompt += "Answer:"
                Task_Prompt += syn_ans_queid[(qa_idx) % len(syn_ans_queid)][:-1].lower()
                Task_Prompt += "\n"
        samples["Task_Prompt"] = Task_Prompt
        # print(Task_Prompt)
        return Task_Prompt

    def prompts_construction(
        self,
        samples,
        question_type="neural",
        num_caps_per_img=30,
        num_question_per_img=30,
    ):
        Prompt = "Please reason the answer of the questions according to the given contexts.\n"

        Context_Prompt = self.create_context_prompt(samples, num_caps_per_img)

        Task_Prompt = self.create_task_prompt(
            samples, question_type, num_question_per_img
        )

        Img2Prompt = (
            Prompt
            + "Contexts:"
            + Context_Prompt


                        random.randint(0, len(Noun_Questions) - 1)
                    ]
                elif doc[-1].pos_ == "VERB":
                    Task_Prompt += Verb_Questions[
                        random.randint(0, len(Verb_Questions) - 1)
                    ]
                elif doc[-1].pos_ == "ADJ":
                    Task_Prompt += Adj_Questions[
                        random.randint(0, len(Adj_Questions) - 1)
                    ]

                Task_Prompt += "\n"

                Task_Prompt += "Answer:"
                Task_Prompt += syn_ans_queid[(qa_idx) % len(syn_ans_queid)][:-1].lower()
                Task_Prompt += "\n"
        samples["Task_Prompt"] = Task_Prompt
        # print(Task_Prompt)
        return Task_Prompt

    def prompts_construction(
        self,
        samples,
        question_type="neural",
        num_caps_per_img=30,
        num_question_per_img=30,
    ):
        Prompt = "Please reason the answer of the questions according to the given contexts.\n"

        Context_Prompt = self.create_context_prompt(samples, num_caps_per_img)

        Task_Prompt = self.create_task_prompt(
            samples, question_type, num_question_per_img
        )

        Img2Prompt = (
            Prompt
            + "Contexts:"
            + Context_Prompt
            + "\n"
            + Task_Prompt
            + "Question:"
            + samples["text_input"][0]
            + "\nAnswer:"
        )
        return Img2Prompt

    def prepare_LLM_input(
        self,
        samples,
        num_beams=1,
        inference_method="generate",
        max_len=20,
        min_len=0,
        internal_bsz_fid=1,
        num_captions=50,
        num_captions_fid=1,
        cap_max_length=20,
        cap_min_length=10,
        top_k=50,
        top_p=1,
        repetition_penalty=1,
        num_patches=20,
        block_num=7,
    ):
        """


        Prompt = "Please reason the answer of the questions according to the given contexts.\n"

        Context_Prompt = self.create_context_prompt(samples, num_caps_per_img)

        Task_Prompt = self.create_task_prompt(
            samples, question_type, num_question_per_img
        )

        Img2Prompt = (
            Prompt
            + "Contexts:"
            + Context_Prompt
            + "\n"
            + Task_Prompt
            + "Question:"
            + samples["text_input"][0]
            + "\nAnswer:"
        )
        return Img2Prompt

    def prepare_LLM_input(
        self,
        samples,
        num_beams=1,
        inference_method="generate",
        max_len=20,
        min_len=0,
        internal_bsz_fid=1,
        num_captions=50,
        num_captions_fid=1,
        cap_max_length=20,
        cap_min_length=10,
        top_k=50,
        top_p=1,
        repetition_penalty=1,
        num_patches=20,
        block_num=7,
    ):
        """
        Args:
            samples (dict): A dictionary containing the following keys:
                - image (torch.Tensor): A tensor of shape (batch_size, 3, H, W). Default H=480, W=480.
                - text_input (str or [str]): String or a list of strings, each string is a question.
                                             The number of questions must be equal to the batch size. If a single string, will be converted to a list of string, with length 1 first.
            num_beams (int): Number of beams for beam search. 1 means no beam search.
            inference_method (str): Inference method. Must be "generate". The model will generate answers.
            max_len (int): Maximum length of generated answers.
            min_len (int): Minimum length of generated answers.
            internal_bsz_fid (int): Internal batch size when using FiD decoding.
            num_captions (int): Number of captions generated for each image.
            num_captions_fid (int): Number of captions concatenated with a question during FiD decoding.
            cap_max_length (int): The maximum length of the caption to be generated.
            cap_min_length (int): The minimum length of the caption to be generated.
            top_k (float): The number of the highest probability tokens for top-k sampling.
            top_p (float): The cumulative probability for nucleus sampling.
            repetition_penalty (float): The parameter for repetition penalty. 1.0 means no penalty.
            num_patches (int): Number of patches sampled for each image.
            block_num (int): The index of cross-attention block for gradcam computation.

        Returns:
            List: A list of strings, each string is an answer.
            gradcams (torch.Tensor): A tensor of shape (batch_size, H*W)
            captions (nested list): A nested list of strings of total length batch_size * num_captions
        """
        assert inference_method in [
            "generate",
        ], "Inference method must be 'generate', got {}.".format(inference_method)

        if isinstance(samples["text_input"], str):
            samples["text_input"] = [samples["text_input"]]

        assert len(samples["text_input"]) == samples["image"].size(
            0
        ), "The number of questions must be equal to the batch size."

        samples = self.forward_itm(samples, block_num=block_num)

        samples = self.forward_cap(
            samples,
            cap_max_length=cap_max_length,
            cap_min_length=cap_min_length,
            top_k=top_k,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            num_captions=num_captions,
            num_patches=num_patches,
        )

        if self.offload_model:
            samples["image"] = samples["image"].to("cpu")
            self.image_question_matching_model.to("cpu")
            self.image_captioning_model.to("cpu")
        torch.cuda.empty_cache()

        pred_answers = self.forward_qa(
            samples,
            num_beams=num_beams,
            max_len=max_len,
            min_len=min_len,
            internal_bsz_fid=internal_bsz_fid,
            num_captions=num_captions,
            num_captions_fid=num_captions_fid,
        )

        if self.offload_model:
            self.image_question_matching_model.to(self.question_answering_model.device)
            self.image_captioning_model.to(self.question_answering_model.device)

        return pred_answers, samples["captions"], samples["gradcams"]

    @classmethod
    def from_config(cls, model_config):
        itm_config = model_config.image_question_matching_model
        cap_config = model_config.image_captioning_model

        itm_cls = registry.get_model_class(itm_config.arch)
        cap_cls = registry.get_model_class(cap_config.arch)

        image_question_matching_model = itm_cls.from_config(itm_config)
        image_captioning_model = cap_cls.from_config(cap_config)

        question_generation_tokenizer = T5Tokenizer.from_pretrained(
            "google/t5-large-lm-adapt"
        )
        question_generation_model = T5ForConditionalGeneration.from_pretrained(
            "google/t5-large-lm-adapt"
        )
        cached_file = download_cached_file(
            "https://storage.googleapis.com/sfr-vision-language-research/LAVIS/projects/img2prompt/T5_large_QG.pth",


            self.image_captioning_model.to("cpu")
        torch.cuda.empty_cache()

        pred_answers = self.forward_qa(
            samples,
            num_beams=num_beams,
            max_len=max_len,
            min_len=min_len,
            internal_bsz_fid=internal_bsz_fid,
            num_captions=num_captions,
            num_captions_fid=num_captions_fid,
        )

        if self.offload_model:
            self.image_question_matching_model.to(self.question_answering_model.device)
            self.image_captioning_model.to(self.question_answering_model.device)

        return pred_answers, samples["captions"], samples["gradcams"]

    @classmethod
    def from_config(cls, model_config):
        itm_config = model_config.image_question_matching_model
        cap_config = model_config.image_captioning_model

        itm_cls = registry.get_model_class(itm_config.arch)
        cap_cls = registry.get_model_class(cap_config.arch)

        image_question_matching_model = itm_cls.from_config(itm_config)
        image_captioning_model = cap_cls.from_config(cap_config)

        question_generation_tokenizer = T5Tokenizer.from_pretrained(
            "google/t5-large-lm-adapt"
        )
        question_generation_model = T5ForConditionalGeneration.from_pretrained(
            "google/t5-large-lm-adapt"
        )
        cached_file = download_cached_file(
            "https://storage.googleapis.com/sfr-vision-language-research/LAVIS/projects/img2prompt/T5_large_QG.pth",
            check_hash=False,
            progress=True,
        )
        checkpoint = torch.load(cached_file, map_location="cpu")
        state_dict = checkpoint["model"]
        question_generation_model.load_state_dict(state_dict)
        model = cls(
            image_question_matching_model=image_question_matching_model,
            image_captioning_model=image_captioning_model,
            question_generation_model=question_generation_model,
            question_generation_tokenizer=question_generation_tokenizer,
            offload_model=False,
        )

        return model


"""
 Copyright (c) 2023, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""
import logging

import copy
import torch
import torch.nn as nn
from torch.cuda.amp import autocast as autocast
from transformers import T5TokenizerFast, BertTokenizer

from lavis.common.registry import registry
from lavis.models.blip2_models.Qformer import BertConfig, BertLMHeadModel
from lavis.models.blip2_models.blip2 import Blip2Base, disabled_train
from lavis.models.blip2_models.modeling_t5 import T5Config, T5ForConditionalGeneration
from positional_encodings.torch_encodings import PositionalEncoding1D

def print_trainable_parameters(model):
    trainable_params = 0
    all_param = 0
    for _, param in model.named_parameters():
        all_param += param.numel()
        if param.requires_grad:
            trainable_params += param.numel()
    print(
        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param:.2f}"
    )

@registry.register_model("crema")
class CREMA(Blip2Base):
    """
    BLIP2 T5 model.
    Supported model types:
        - pretrain_flant5xl: pretrained model with FlanT5-XL
        - pretrain_flant5xxl: pretrained model with FlanT5-XXL
        - caption_coco_flant5xl: fintuned image captioning model with FlanT5-XL
    Usage:
        >>> from lavis.models import load_model
        >>> model = load_model("blip2_t5", "pretrain_flant5xl")
    """
    PRETRAINED_MODEL_CONFIG_DICT = {
        "pretrain_flant5xl": "configs/models/blip2/blip2_pretrain_flant5xl.yaml",
        "pretrain_flant5xxl": "configs/models/blip2/blip2_pretrain_flant5xxl.yaml",
        "caption_coco_flant5xl": "configs/models/blip2/blip2_caption_flant5xl.yaml",
    }

    def __init__( self, img_size=224, drop_path_rate=0,
        use_grad_checkpoint=False, vit_precision="fp16", freeze_vit=True,
        num_query_token=32, t5_model="google/flan-t5-xl", prompt="",
        max_txt_len=32, frame_num=8, answer_num=5, apply_lemmatizer=False,
        task='concate',
        modalities='rgb',
        downstream_task='mcqa', # caption / oeqa / mcqa
        lora_rank=64,
        lora_layer=None,
        lora_dropout=0.1,
        fuse_with_base_modality=False):

        """
        apply_lemmatizer: when set to True, postprocess predict_answers() result with lemmas.
        """
        super().__init__()

        self.task = task #.split('_')
        self.modalities = modalities.split('_')
        self.fuse_with_base_modality=fuse_with_base_modality

        print(self.modalities)
        num_features = 1408
        # ========= init vision encoder ============
        # init vision backbone for vision experts
        if 'rgb' in self.modalities or 'depth' in self.modalities or 'flow' in self.modalities or 'norm' in self.modalities:
            self.visual_encoder = self.init_vision_encoder_only(
                img_size, drop_path_rate, use_grad_checkpoint, vit_precision)
            for name, param in self.visual_encoder.named_parameters():
                param.requires_grad = False
            self.visual_encoder = self.visual_encoder.eval()
            self.visual_encoder.train = disabled_train
            logging.info("freeze vision encoder")
        if 'audio' in self.modalities:
            self.audio_encoder, self.ln_audio = self.init_audio_encoder('beats', cached_audio=False)
            for name, param in self.audio_encoder.named_parameters():
                param.requires_grad = False
            self.audio_encoder = self.audio_encoder.eval()
            self.audio_encoder.train = disabled_train
            logging.info("freeze audio encoder")

        if 'pc' in self.modalities:
            # pre-extracted features
            pass

        # print('num_features', self.visual_encoder.num_features) 1408
        # ========= init LLM ============
        # text backbone
        self.t5_tokenizer = T5TokenizerFast.from_pretrained(t5_model)
        t5_config = T5Config.from_pretrained(t5_model)
        t5_config.dense_act_fn = "gelu"
        self.t5_model = T5ForConditionalGeneration.from_pretrained(t5_model, config=t5_config)
        # freeze T5
        for name, param in self.t5_model.named_parameters():
            param.requires_grad = False
            param.data = param.data.bfloat16()

        # ========= init Qformer ============
       # if 'rgb' in self.modalities or 'depth' in self.modalities or 'flow' in self.modalities or 'norm' in self.modalities:
        self.Qformer, encoder_config = self.init_Multimodal_Qformer(
            num_query_token, num_features, #self.visual_encoder.num_features,
            modulars=self.modalities,
            r=lora_rank, lora_layer=lora_layer, lora_dropout=lora_dropout)

        self.Qformer.cls = None
        self.Qformer.bert.embeddings.word_embeddings = None
        self.Qformer.bert.embeddings.position_embeddings = None
        for layer in self.Qformer.bert.encoder.layer:
            layer.output = None
            layer.intermediate = None
        self.num_query_token = num_query_token

        if 'rgb' in self.modalities:
            self.query_tokens_rgb = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_rgb.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_rgb = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)
            self.ln_rgb = nn.LayerNorm(self.visual_encoder.num_features)

        if 'flow' in self.modalities:
            self.query_tokens_flow = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_flow.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_flow = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)
            self.ln_flow = nn.LayerNorm(self.visual_encoder.num_features)

        if 'norm' in self.modalities:
            self.query_tokens_norm = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_norm.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_norm = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)
            self.ln_norm = nn.LayerNorm(self.visual_encoder.num_features)

        if 'depth' in self.modalities:
            self.query_tokens_depth = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_depth.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_depth = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)
            self.ln_depth = nn.LayerNorm(self.visual_encoder.num_features)

        if 'audio' in self.modalities:
            self.query_tokens_audio = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_audio.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_audio = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)

            self.projection_audio = nn.Linear(self.audio_encoder.num_features, num_features)
            self.ln_audio = nn.LayerNorm(num_features)

        if 'pc' in self.modalities:
            self.query_tokens_pc = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_pc.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_pc = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)
            self.ln_pc = nn.LayerNorm(num_features)
            pos_model = PositionalEncoding1D(1408 // 3)
            x = torch.zeros(1, 256, 1408 // 3)
            self.pos_embedding = pos_model(x).squeeze().cuda()

        if 'espresso' in self.task:

            self.fusion = nn.Sequential(
                nn.Linear(2048*(len(self.modalities)-1), 2048),
                )
            self.sigmoid = nn.Sigmoid()

        self.downstream_task = downstream_task
        self.max_txt_len = 77
        answer_id = [71, 272, 205, 309, 262] # A B C D E
        self.answer_id = answer_id[:answer_num]

        self._apply_lemmatizer = apply_lemmatizer
        self._lemmatizer = None
        self.frame_num = frame_num
        self.ANS_MAP = {'A':0, 'B':1, 'C':2, 'D':3, 'E':4}

        if frame_num == 1:
            self.vid_prefix = ['Frame: ']
            self.depth_prefix = ['Depth Map: ']
            self.flow_prefix = ['Optical Flow: ']
            self.norm_prefix = ['Surface Normalization: ']
        else:
            self.vid_prefix = ['Frame {}: '.format(str(i+1)) for i in range(frame_num)]
            self.depth_prefix = ['Depth Map {}: '.format(str(i+1)) for i in range(frame_num)]
            self.flow_prefix = ['Optical Flow {}: '.format(str(i+1)) for i in range(frame_num)]
            self.norm_prefix = ['Surface Normalization {}: '.format(str(i+1)) for i in range(frame_num)]

        self.audio_prefix = ['Audio: ']
        self.pc_prefix = ['3D Model: ']

    def forward(self, samples):

        # rgb visual embedding
        qa_text, answer = samples['qa_input'], samples['qa_output']
        b = len(qa_text)

        input_embed_dict, input_att_dict = {}, {}

        for modal in self.modalities:
            input = samples[modal]
            if modal in ['rgb', 'depth', 'norm', 'flow']:
                # fix some loading issue
                if input.shape[1] == 3:
                    input = input.permute(0, 2, 1, 3, 4)

            # following 3D-LLM
            if modal == 'pc':
                with torch.cuda.amp.autocast(dtype=torch.float32):
                    pc_embeds = samples["pc_feat"]
                    pc = samples["pc"].long()
                    all_pcs = torch.zeros((pc_embeds.shape))
                    for j in range(pc.shape[0]):
                        pcs = []
                        for i in range(3):
                            pc_i = pc[j][:, i]
                            pcs.append(self.pos_embedding[pc_i])
                        pcs = torch.cat(pcs, -1)
                        all_pcs[j][:, :1407] = pcs
                    all_pcs = all_pcs.cuda()
                pc_embeds = pc_embeds + 0.01 * all_pcs
                atts = torch.ones(pc_embeds.size()[:-1], dtype=torch.long).to(pc_embeds.device)
                input_embed_dict[modal], input_att_dict[modal] = pc_embeds, atts
            else:
                input_embed_dict[modal], input_att_dict[modal] = self.encode_input(input, modal)

        device = input_embed_dict[list(input_embed_dict.keys())[0]].device

        input_text= self.t5_tokenizer(
            qa_text, padding="longest", truncation=True,
            max_length=self.max_txt_len, return_tensors="pt").to(device)
        input_text_embeds = self.t5_model.encoder.embed_tokens(input_text.input_ids)

        fusion_modal = []
        t5_inputs, t5_atts, t5_query = {}, {}, {}
        for modal in self.modalities:
            t5_inputs[modal], t5_atts[modal], t5_query[modal] = self.get_qformer_embedding(input_embed_dict[modal], input_att_dict[modal], device, modal, b)

        if 'rgb' in self.modalities:
            inputs_t5_rgb = t5_inputs['rgb']
            atts_t5_rgb = t5_atts['rgb']
            vid_prefix_embed, vid_prefix_mask = self.get_prefix_embedding(self.vid_prefix, b, device)
            inputs_t5 = torch.cat([vid_prefix_embed, inputs_t5_rgb], dim=2) # b, t, n_word + m, c
            atts_t5 = torch.cat([vid_prefix_mask, atts_t5_rgb], dim=2) # b, t, n_word + m

            for modal in self.modalities:
                if modal == 'rgb':
                    continue
                if modal in ['depth', 'norm', 'flow']:
                    if 'espresso' in self.task:
                        fusion_modal.append(t5_inputs[modal])
                    else:
                        inputs_t5 = torch.cat([inputs_t5, t5_inputs[modal]], dim=2)
                        atts_t5 = torch.cat([atts_t5, t5_atts[modal]], dim=2)

                if modal in ['pc']:
                    if 'espresso' in self.task:
                        pc = t5_inputs[modal]
                        pc = pc.unsqueeze(1)
                        pc = torch.repeat_interleave(pc, self.frame_num, 1) # [16, 4, 32, 2048]
                        fusion_modal.append(pc)

                if modal in ['audio']:
                    if 'espresso' in self.task:
                        audio = t5_inputs[modal]
                        audio = audio.mean(dim=1)
                        audio = audio.unsqueeze(1)
                        audio = torch.repeat_interleave(audio, self.frame_num, 1) # [16, 4, 32, 2048]
                        fusion_modal.append(audio)

            # visual only input
            if 'audio' not in self.modalities and 'pc' not in self.modalities:
                if 'espresso' in self.task:
                    fusion_modal = torch.cat(fusion_modal, dim=-1) # 16, 4, 32, 8192
                    inputs_t5_extra = self.fusion(fusion_modal)
                    inputs_t5_rgb += self.sigmoid(inputs_t5_extra) * inputs_t5_extra
                    inputs_t5 = torch.cat([vid_prefix_embed, inputs_t5_rgb], dim=2)
                    atts_t5 = torch.ones(inputs_t5.size()[:-1], dtype=torch.long).to(device)
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)
                else:
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)

            # [F1, F2, F3,..., A]
            elif 'audio' in self.modalities:
                if 'espresso' in self.task:
                    fusion_modal = torch.cat(fusion_modal, dim=-1) # 16, 4, 32, 8192
                    inputs_t5_extra = self.fusion(fusion_modal)
                    inputs_t5_rgb += self.sigmoid(inputs_t5_extra) * inputs_t5_extra
                    inputs_t5 = torch.cat([vid_prefix_embed, inputs_t5_rgb], dim=2)
                    atts_t5 = torch.ones(inputs_t5.size()[:-1], dtype=torch.long).to(device)
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)
                else:
                    audio_prefix_embed, audio_prefix_mask = self.get_prefix_embedding(self.audio_prefix, b, device)
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)
                    # seems no prefix works better for audio-video reasoning
                    inputs_t5 = torch.cat([inputs_t5, t5_inputs['audio'].reshape(b, -1, t5_inputs['audio'].shape[-1])], dim=1)
                    atts_t5 = torch.cat([atts_t5, t5_atts['audio'].reshape(b, -1)], dim=1)

            elif 'pc' in self.modalities:
                if 'espresso' in self.task:
                    fusion_modal = torch.cat(fusion_modal, dim=-1) # 16, 4, 32, 8192
                    inputs_t5_extra = self.fusion(fusion_modal)
                    inputs_t5_rgb += self.sigmoid(inputs_t5_extra) * inputs_t5_extra
                    inputs_t5 = torch.cat([vid_prefix_embed, inputs_t5_rgb], dim=2)
                    atts_t5 = torch.ones(inputs_t5.size()[:-1], dtype=torch.long).to(device)
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)
                else:
                    pc_prefix_embed, pc_prefix_mask = self.get_prefix_embedding(self.pc_prefix, b, device)
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)
                    inputs_t5 = torch.cat([inputs_t5, pc_prefix_embed.squeeze(), t5_inputs['pc'].reshape(b, -1, t5_inputs['pc'].shape[-1])], dim=1)
                    atts_t5 = torch.cat([atts_t5, pc_prefix_mask.squeeze(), t5_atts['pc'].reshape(b, -1)], dim=1)

        elif 'audio' in self.modalities: # audio
            inputs_t5 = t5_inputs['audio'].reshape(b, -1, t5_inputs['audio'].shape[-1])
            atts_t5 = t5_atts['audio'].reshape(b, -1)
            audio_prefix_embed, audio_prefix_mask = self.get_prefix_embedding(self.audio_prefix, b, device)
            inputs_t5 = torch.cat([audio_prefix_embed.squeeze(), inputs_t5], dim=1)
            atts_t5 = torch.cat([audio_prefix_mask.squeeze(), atts_t5], dim=1)

        elif 'pc' in self.modalities: # pc
            inputs_t5 = t5_inputs['pc'].reshape(b, -1, t5_inputs['pc'].shape[-1])
            atts_t5 = t5_atts['pc'].reshape(b, -1)
            pc_prefix_embed, pc_prefix_mask = self.get_prefix_embedding(self.audio_prefix, b, device)
            inputs_t5 = torch.cat([pc_prefix_embed.squeeze(), inputs_t5], dim=1)
            atts_t5 = torch.cat([pc_prefix_mask.squeeze(), atts_t5], dim=1)

        inputs_embeds = torch.cat([inputs_t5, input_text_embeds], dim=1)
        encoder_atts = torch.cat([atts_t5, input_text.attention_mask], dim=1)

        with torch.cuda.amp.autocast(dtype=torch.bfloat16):

            output_tokens = self.t5_tokenizer(
                answer, padding="longest", truncation=True,
                max_length=self.max_txt_len, return_tensors="pt").to(device)
            targets_qa = output_tokens.input_ids.masked_fill(
                output_tokens.input_ids == self.t5_tokenizer.pad_token_id, -100)
            output_tokens_mask = output_tokens.attention_mask

            outputs = self.t5_model(
                    inputs_embeds=inputs_embeds, attention_mask=encoder_atts,
                    decoder_attention_mask=output_tokens_mask, return_dict=True, labels=targets_qa)
            loss = outputs.loss

        return {'loss': loss}

    def encode_input(self, input, modality, training=True):

        ln = getattr(self, f"ln_{modality}")

        if modality in ['rgb', 'depth', 'flow', 'norm']:
            modality = 'visual'
        if modality in ['audio']:
            modality = 'audio'
        if modality in ['pc']:
            modality = 'pc'

        encoder = getattr(self, f"{modality}_encoder")

        if modality == 'visual':
            b, t, c, w, h = input.shape
            input = input.reshape(-1, c, w, h)
            if training:
                image_embeds = ln(encoder(input))
            else:
                with torch.cuda.amp.autocast(enabled=(self.device != torch.device("cpu"))):
                    image_embeds = ln(encoder(input))
            _, n, _ = image_embeds.shape
            image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(input.device) # bt n c
            return image_embeds, image_atts

        if modality == 'audio':
            embeds, atts = [], []
            for j in range(input.size(1)):
                this_frame = input[:,j,:,:]
                if training:
                    embeds.append(encoder(this_frame))
                else:
                    with torch.cuda.amp.autocast(enabled=(self.device != torch.device("cpu"))):
                        embeds.append(encoder(this_frame))
                atts.append(torch.ones(embeds[j].size()[:-1], dtype=torch.long).to(input.device))

            # print('here', len(embeds), embeds[0].shape) # 2, 3, 256, 768
            embeds = torch.stack(embeds
# ... truncated ...


"""
 Copyright (c) 2023, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""
import logging

import copy
import torch
import torch.nn as nn
from torch.cuda.amp import autocast as autocast
from transformers import T5TokenizerFast, BertTokenizer

from lavis.common.registry import registry
from lavis.models.blip2_models.Qformer import BertConfig, BertLMHeadModel
from lavis.models.blip2_models.blip2 import Blip2Base, disabled_train
from lavis.models.blip2_models.modeling_t5 import T5Config, T5ForConditionalGeneration
from positional_encodings.torch_encodings import PositionalEncoding1D

def print_trainable_parameters(model):
    trainable_params = 0
    all_param = 0
    for _, param in model.named_parameters():
        all_param += param.numel()
        if param.requires_grad:
            trainable_params += param.numel()
    print(
        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param:.2f}"
    )

@registry.register_model("crema")
class CREMA(Blip2Base):
    """
    BLIP2 T5 model.
    Supported model types:
        - pretrain_flant5xl: pretrained model with FlanT5-XL
        - pretrain_flant5xxl: pretrained model with FlanT5-XXL
        - caption_coco_flant5xl: fintuned image captioning model with FlanT5-XL
    Usage:
        >>> from lavis.models import load_model
        >>> model = load_model("blip2_t5", "pretrain_flant5xl")
    """
    PRETRAINED_MODEL_CONFIG_DICT = {
        "pretrain_flant5xl": "configs/models/blip2/blip2_pretrain_flant5xl.yaml",
        "pretrain_flant5xxl": "configs/models/blip2/blip2_pretrain_flant5xxl.yaml",
        "caption_coco_flant5xl": "configs/models/blip2/blip2_caption_flant5xl.yaml",
    }

    def __init__( self, img_size=224, drop_path_rate=0,


CREMA

from transformers import T5TokenizerFast, BertTokenizer

from lavis.common.registry import registry
from lavis.models.blip2_models.Qformer import BertConfig, BertLMHeadModel
from lavis.models.blip2_models.blip2 import Blip2Base, disabled_train
from lavis.models.blip2_models.modeling_t5 import T5Config, T5ForConditionalGeneration
from positional_encodings.torch_encodings import PositionalEncoding1D

def print_trainable_parameters(model):
    trainable_params = 0
    all_param = 0
    for _, param in model.named_parameters():
        all_param += param.numel()
        if param.requires_grad:
            trainable_params += param.numel()
    print(
        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param:.2f}"
    )

@registry.register_model("crema")
class CREMA(Blip2Base):
    """
    BLIP2 T5 model.
    Supported model types:
        - pretrain_flant5xl: pretrained model with FlanT5-XL
        - pretrain_flant5xxl: pretrained model with FlanT5-XXL
        - caption_coco_flant5xl: fintuned image captioning model with FlanT5-XL
    Usage:
        >>> from lavis.models import load_model
        >>> model = load_model("blip2_t5", "pretrain_flant5xl")
    """
    PRETRAINED_MODEL_CONFIG_DICT = {
        "pretrain_flant5xl": "configs/models/blip2/blip2_pretrain_flant5xl.yaml",
        "pretrain_flant5xxl": "configs/models/blip2/blip2_pretrain_flant5xxl.yaml",
        "caption_coco_flant5xl": "configs/models/blip2/blip2_caption_flant5xl.yaml",
    }

    def __init__( self, img_size=224, drop_path_rate=0,
        use_grad_checkpoint=False, vit_precision="fp16", freeze_vit=True,
        num_query_token=32, t5_model="google/flan-t5-xl", prompt="",
        max_txt_len=32, frame_num=8, answer_num=5, apply_lemmatizer=False,
        task='concate',
        modalities='rgb',
        downstream_task='mcqa', # caption / oeqa / mcqa
        lora_rank=64,
        lora_layer=None,
        lora_dropout=0.1,
        fuse_with_base_modality=False):

        """
        apply_lemmatizer: when set to True, postprocess predict_answers() result with lemmas.
        """
        super().__init__()

        self.task = task #.split('_')
        self.modalities = modalities.split('_')
        self.fuse_with_base_modality=fuse_with_base_modality

        print(self.modalities)
        num_features = 1408
        # ========= init vision encoder ============
        # init vision backbone for vision experts
        if 'rgb' in self.modalities or 'depth' in self.modalities or 'flow' in self.modalities or 'norm' in self.modalities:
            self.visual_encoder = self.init_vision_encoder_only(
                img_size, drop_path_rate, use_grad_checkpoint, vit_precision)
            for name, param in self.visual_encoder.named_parameters():
                param.requires_grad = False
            self.visual_encoder = self.visual_encoder.eval()
            self.visual_encoder.train = disabled_train
            logging.info("freeze vision encoder")
        if 'audio' in self.modalities:
            self.audio_encoder, self.ln_audio = self.init_audio_encoder('beats', cached_audio=False)
            for name, param in self.audio_encoder.named_parameters():
                param.requires_grad = False
            self.audio_encoder = self.audio_encoder.eval()
            self.audio_encoder.train = disabled_train
            logging.info("freeze audio encoder")

        if 'pc' in self.modalities:
            # pre-extracted features
            pass

        # print('num_features', self.visual_encoder.num_features) 1408
        # ========= init LLM ============
        # text backbone
        self.t5_tokenizer = T5TokenizerFast.from_pretrained(t5_model)
        t5_config = T5Config.from_pretrained(t5_model)
        t5_config.dense_act_fn = "gelu"
        self.t5_model = T5ForConditionalGeneration.from_pretrained(t5_model, config=t5_config)
        # freeze T5
        for name, param in self.t5_model.named_parameters():
            param.requires_grad = False
            param.data = param.data.bfloat16()

        # ========= init Qformer ============
       # if 'rgb' in self.modalities or 'depth' in self.modalities or 'flow' in self.modalities or 'norm' in self.modalities:
        self.Qformer, encoder_config = self.init_Multimodal_Qformer(
            num_query_token, num_features, #self.visual_encoder.num_features,
            modulars=self.modalities,
            r=lora_rank, lora_layer=lora_layer, lora_dropout=lora_dropout)

        self.Qformer.cls = None
        self.Qformer.bert.embeddings.word_embeddings = None
        self.Qformer.bert.embeddings.position_embeddings = None
        for layer in self.Qformer.bert.encoder.layer:
            layer.output = None
            layer.intermediate = None
        self.num_query_token = num_query_token

        if 'rgb' in self.modalities:
            self.query_tokens_rgb = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_rgb.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_rgb = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)
            self.ln_rgb = nn.LayerNorm(self.visual_encoder.num_features)

        if 'flow' in self.modalities:
            self.query_tokens_flow = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_flow.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_flow = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)
            self.ln_flow = nn.LayerNorm(self.visual_encoder.num_features)

        if 'norm' in self.modalities:
            self.query_tokens_norm = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_norm.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_norm = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)
            self.ln_norm = nn.LayerNorm(self.visual_encoder.num_features)

        if 'depth' in self.modalities:
            self.query_tokens_depth = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_depth.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_depth = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)
            self.ln_depth = nn.LayerNorm(self.visual_encoder.num_features)

        if 'audio' in self.modalities:
            self.query_tokens_audio = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_audio.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_audio = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)

            self.projection_audio = nn.Linear(self.audio_encoder.num_features, num_features)
            self.ln_audio = nn.LayerNorm(num_features)

        if 'pc' in self.modalities:
            self.query_tokens_pc = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_pc.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_pc = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)
            self.ln_pc = nn.LayerNorm(num_features)
            pos_model = PositionalEncoding1D(1408 // 3)
            x = torch.zeros(1, 256, 1408 // 3)
            self.pos_embedding = pos_model(x).squeeze().cuda()

        if 'espresso' in self.task:

            self.fusion = nn.Sequential(
                nn.Linear(2048*(len(self.modalities)-1), 2048),
                )
            self.sigmoid = nn.Sigmoid()

        self.downstream_task = downstream_task
        self.max_txt_len = 77
        answer_id = [71, 272, 205, 309, 262] # A B C D E
        self.answer_id = answer_id[:answer_num]

        self._apply_lemmatizer = apply_lemmatizer
        self._lemmatizer = None
        self.frame_num = frame_num
        self.ANS_MAP = {'A':0, 'B':1, 'C':2, 'D':3, 'E':4}

        if frame_num == 1:
            self.vid_prefix = ['Frame: ']
            self.depth_prefix = ['Depth Map: ']
            self.flow_prefix = ['Optical Flow: ']
            self.norm_prefix = ['Surface Normalization: ']
        else:
            self.vid_prefix = ['Frame {}: '.format(str(i+1)) for i in range(frame_num)]
            self.depth_prefix = ['Depth Map {}: '.format(str(i+1)) for i in range(frame_num)]
            self.flow_prefix = ['Optical Flow {}: '.format(str(i+1)) for i in range(frame_num)]
            self.norm_prefix = ['Surface Normalization {}: '.format(str(i+1)) for i in range(frame_num)]

        self.audio_prefix = ['Audio: ']
        self.pc_prefix = ['3D Model: ']

    def forward(self, samples):

        # rgb visual embedding
        qa_text, answer = samples['qa_input'], samples['qa_output']
        b = len(qa_text)

        input_embed_dict, input_att_dict = {}, {}

        for modal in self.modalities:
            input = samples[modal]
            if modal in ['rgb', 'depth', 'norm', 'flow']:
                # fix some loading issue
                if input.shape[1] == 3:
                    input = input.permute(0, 2, 1, 3, 4)

            # following 3D-LLM
            if modal == 'pc':
                with torch.cuda.amp.autocast(dtype=torch.float32):
                    pc_embeds = samples["pc_feat"]
                    pc = samples["pc"].long()
                    all_pcs = torch.zeros((pc_embeds.shape))
                    for j in range(pc.shape[0]):
                        pcs = []
                        for i in range(3):
                            pc_i = pc[j][:, i]
                            pcs.append(self.pos_embedding[pc_i])
                        pcs = torch.cat(pcs, -1)
                        all_pcs[j][:, :1407] = pcs
                    all_pcs = all_pcs.cuda()
                pc_embeds = pc_embeds + 0.01 * all_pcs
                atts = torch.ones(pc_embeds.size()[:-1], dtype=torch.long).to(pc_embeds.device)
                input_embed_dict[modal], input_att_dict[modal] = pc_embeds, atts
            else:
                input_embed_dict[modal], input_att_dict[modal] = self.encode_input(input, modal)

        device = input_embed_dict[list(input_embed_dict.keys())[0]].device

        input_text= self.t5_tokenizer(
            qa_text, padding="longest", truncation=True,
            max_length=self.max_txt_len, return_tensors="pt").to(device)
        input_text_embeds = self.t5_model.encoder.embed_tokens(input_text.input_ids)

        fusion_modal = []
        t5_inputs, t5_atts, t5_query = {}, {}, {}
        for modal in self.modalities:
            t5_inputs[modal], t5_atts[modal], t5_query[modal] = self.get_qformer_embedding(input_embed_dict[modal], input_att_dict[modal], device, modal, b)

        if 'rgb' in self.modalities:
            inputs_t5_rgb = t5_inputs['rgb']
            atts_t5_rgb = t5_atts['rgb']
            vid_prefix_embed, vid_prefix_mask = self.get_prefix_embedding(self.vid_prefix, b, device)
            inputs_t5 = torch.cat([vid_prefix_embed, inputs_t5_rgb], dim=2) # b, t, n_word + m, c
            atts_t5 = torch.cat([vid_prefix_mask, atts_t5_rgb], dim=2) # b, t, n_word + m

            for modal in self.modalities:
                if modal == 'rgb':
                    continue
                if modal in ['depth', 'norm', 'flow']:
                    if 'espresso' in self.task:
                        fusion_modal.append(t5_inputs[modal])
                    else:
                        inputs_t5 = torch.cat([inputs_t5, t5_inputs[modal]], dim=2)
                        atts_t5 = torch.cat([atts_t5, t5_atts[modal]], dim=2)

                if modal in ['pc']:
                    if 'espresso' in self.task:
                        pc = t5_inputs[modal]
                        pc = pc.unsqueeze(1)
                        pc = torch.repeat_interleave(pc, self.frame_num, 1) # [16, 4, 32, 2048]
                        fusion_modal.append(pc)

                if modal in ['audio']:
                    if 'espresso' in self.task:
                        audio = t5_inputs[modal]
                        audio = audio.mean(dim=1)
                        audio = audio.unsqueeze(1)
                        audio = torch.repeat_interleave(audio, self.frame_num, 1) # [16, 4, 32, 2048]
                        fusion_modal.append(audio)

            # visual only input
            if 'audio' not in self.modalities and 'pc' not in self.modalities:
                if 'espresso' in self.task:
                    fusion_modal = torch.cat(fusion_modal, dim=-1) # 16, 4, 32, 8192
                    inputs_t5_extra = self.fusion(fusion_modal)
                    inputs_t5_rgb += self.sigmoid(inputs_t5_extra) * inputs_t5_extra
                    inputs_t5 = torch.cat([vid_prefix_embed, inputs_t5_rgb], dim=2)
                    atts_t5 = torch.ones(inputs_t5.size()[:-1], dtype=torch.long).to(device)
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)
                else:
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)

            # [F1, F2, F3,..., A]
            elif 'audio' in self.modalities:
                if 'espresso' in self.task:
                    fusion_modal = torch.cat(fusion_modal, dim=-1) # 16, 4, 32, 8192
                    inputs_t5_extra = self.fusion(fusion_modal)
                    inputs_t5_rgb += self.sigmoid(inputs_t5_extra) * inputs_t5_extra
                    inputs_t5 = torch.cat([vid_prefix_embed, inputs_t5_rgb], dim=2)
                    atts_t5 = torch.ones(inputs_t5.size()[:-1], dtype=torch.long).to(device)
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)
                else:
                    audio_prefix_embed, audio_prefix_mask = self.get_prefix_embedding(self.audio_prefix, b, device)
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)
                    # seems no prefix works better for audio-video reasoning
                    inputs_t5 = torch.cat([inputs_t5, t5_inputs['audio'].reshape(b, -1, t5_inputs['audio'].shape[-1])], dim=1)
                    atts_t5 = torch.cat([atts_t5, t5_atts['audio'].reshape(b, -1)], dim=1)

            elif 'pc' in self.modalities:
                if 'espresso' in self.task:
                    fusion_modal = torch.cat(fusion_modal, dim=-1) # 16, 4, 32, 8192
                    inputs_t5_extra = self.fusion(fusion_modal)
                    inputs_t5_rgb += self.sigmoid(inputs_t5_extra) * inputs_t5_extra
                    inputs_t5 = torch.cat([vid_prefix_embed, inputs_t5_rgb], dim=2)
                    atts_t5 = torch.ones(inputs_t5.size()[:-1], dtype=torch.long).to(device)
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)
                else:
                    pc_prefix_embed, pc_prefix_mask = self.get_prefix_embedding(self.pc_prefix, b, device)
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)
                    inputs_t5 = torch.cat([inputs_t5, pc_prefix_embed.squeeze(), t5_inputs['pc'].reshape(b, -1, t5_inputs['pc'].shape[-1])], dim=1)
                    atts_t5 = torch.cat([atts_t5, pc_prefix_mask.squeeze(), t5_atts['pc'].reshape(b, -1)], dim=1)

        elif 'audio' in self.modalities: # audio
            inputs_t5 = t5_inputs['audio'].reshape(b, -1, t5_inputs['audio'].shape[-1])
            atts_t5 = t5_atts['audio'].reshape(b, -1)
            audio_prefix_embed, audio_prefix_mask = self.get_prefix_embedding(self.audio_prefix, b, device)
            inputs_t5 = torch.cat([audio_prefix_embed.squeeze(), inputs_t5], dim=1)
            atts_t5 = torch.cat([audio_prefix_mask.squeeze(), atts_t5], dim=1)

        elif 'pc' in self.modalities: # pc
            inputs_t5 = t5_inputs['pc'].reshape(b, -1, t5_inputs['pc'].shape[-1])
            atts_t5 = t5_atts['pc'].reshape(b, -1)
            pc_prefix_embed, pc_prefix_mask = self.get_prefix_embedding(self.audio_prefix, b, device)
            inputs_t5 = torch.cat([pc_prefix_embed.squeeze(), inputs_t5], dim=1)
            atts_t5 = torch.cat([pc_prefix_mask.squeeze(), atts_t5], dim=1)

        inputs_embeds = torch.cat([inputs_t5, input_text_embeds], dim=1)
        encoder_atts = torch.cat([atts_t5, input_text.attention_mask], dim=1)

        with torch.cuda.amp.autocast(dtype=torch.bfloat16):

            output_tokens = self.t5_tokenizer(
                answer, padding="longest", truncation=True,
                max_length=self.max_txt_len, return_tensors="pt").to(device)
            targets_qa = output_tokens.input_ids.masked_fill(
                output_tokens.input_ids == self.t5_tokenizer.pad_token_id, -100)
            output_tokens_mask = output_tokens.attention_mask

            outputs = self.t5_model(
                    inputs_embeds=inputs_embeds, attention_mask=encoder_atts,
                    decoder_attention_mask=output_tokens_mask, return_dict=True, labels=targets_qa)
            loss = outputs.loss

        return {'loss': loss}

    def encode_input(self, input, modality, training=True):

        ln = getattr(self, f"ln_{modality}")

        if modality in ['rgb', 'depth', 'flow', 'norm']:
            modality = 'visual'
        if modality in ['audio']:
            modality = 'audio'
        if modality in ['pc']:
            modality = 'pc'

        encoder = getattr(self, f"{modality}_encoder")

        if modality == 'visual':
            b, t, c, w, h = input.shape
            input = input.reshape(-1, c, w, h)
            if training:
                image_embeds = ln(encoder(input))
            else:
                with torch.cuda.amp.autocast(enabled=(self.device != torch.device("cpu"))):
                    image_embeds = ln(encoder(input))
            _, n, _ = image_embeds.shape
            image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(input.device) # bt n c
            return image_embeds, image_atts

        if modality == 'audio':
            embeds, atts = [], []
            for j in range(input.size(1)):
                this_frame = input[:,j,:,:]
                if training:
                    embeds.append(encoder(this_frame))
                else:
                    with torch.cuda.amp.autocast(enabled=(self.device != torch.device("cpu"))):
                        embeds.append(encoder(this_frame))
                atts.append(torch.ones(embeds[j].size()[:-1], dtype=torch.long).to(input.device))

            # print('here', len(embeds), embeds[0].shape) # 2, 3, 256, 768
            embeds = torch.stack(embeds, dim=1)
            # print('audio_embeds 1', embeds.shape) # 3, 2, 256, 768
            atts = torch.stack(atts, dim=1)
            embeds = self.projection_audio(embeds) # 3, 2, 256, 1408
            embeds = ln(embeds.reshape(-1, embeds.shape[-2], embeds.shape[-1]))
            atts = atts.reshape(-1, atts.shape[-1])
# ... truncated ...


    )

@registry.register_model("crema")
class CREMA(Blip2Base):
    """
    BLIP2 T5 model.
    Supported model types:
        - pretrain_flant5xl: pretrained model with FlanT5-XL
        - pretrain_flant5xxl: pretrained model with FlanT5-XXL
        - caption_coco_flant5xl: fintuned image captioning model with FlanT5-XL
    Usage:
        >>> from lavis.models import load_model
        >>> model = load_model("blip2_t5", "pretrain_flant5xl")
    """
    PRETRAINED_MODEL_CONFIG_DICT = {
        "pretrain_flant5xl": "configs/models/blip2/blip2_pretrain_flant5xl.yaml",
        "pretrain_flant5xxl": "configs/models/blip2/blip2_pretrain_flant5xxl.yaml",
        "caption_coco_flant5xl": "configs/models/blip2/blip2_caption_flant5xl.yaml",
    }

    def __init__( self, img_size=224, drop_path_rate=0,
        use_grad_checkpoint=False, vit_precision="fp16", freeze_vit=True,
        num_query_token=32, t5_model="google/flan-t5-xl", prompt="",
        max_txt_len=32, frame_num=8, answer_num=5, apply_lemmatizer=False,
        task='concate',
        modalities='rgb',
        downstream_task='mcqa', # caption / oeqa / mcqa
        lora_rank=64,
        lora_layer=None,
        lora_dropout=0.1,
        fuse_with_base_modality=False):

        """
        apply_lemmatizer: when set to True, postprocess predict_answers() result with lemmas.
        """
        super().__init__()

        self.task = task #.split('_')
        self.modalities = modalities.split('_')
        self.fuse_with_base_modality=fuse_with_base_modality

        print(self.modalities)
        num_features = 1408
        # ========= init vision encoder ============
        # init vision backbone for vision experts
        if 'rgb' in self.modalities or 'depth' in self.modalities or 'flow' in self.modalities or 'norm' in self.modalities:
            self.visual_encoder = self.init_vision_encoder_only(
                img_size, drop_path_rate, use_grad_checkpoint, vit_precision)
            for name, param in self.visual_encoder.named_parameters():
                param.requires_grad = False
            self.visual_encoder = self.visual_encoder.eval()
            self.visual_encoder.train = disabled_train
            logging.info("freeze vision encoder")
        if 'audio' in self.modalities:
            self.audio_encoder, self.ln_audio = self.init_audio_encoder('beats', cached_audio=False)
            for name, param in self.audio_encoder.named_parameters():
                param.requires_grad = False
            self.audio_encoder = self.audio_encoder.eval()
            self.audio_encoder.train = disabled_train
            logging.info("freeze audio encoder")

        if 'pc' in self.modalities:
            # pre-extracted features
            pass

        # print('num_features', self.visual_encoder.num_features) 1408
        # ========= init LLM ============
        # text backbone
        self.t5_tokenizer = T5TokenizerFast.from_pretrained(t5_model)
        t5_config = T5Config.from_pretrained(t5_model)
        t5_config.dense_act_fn = "gelu"
        self.t5_model = T5ForConditionalGeneration.from_pretrained(t5_model, config=t5_config)
        # freeze T5
        for name, param in self.t5_model.named_parameters():
            param.requires_grad = False
            param.data = param.data.bfloat16()

        # ========= init Qformer ============
       # if 'rgb' in self.modalities or 'depth' in self.modalities or 'flow' in self.modalities or 'norm' in self.modalities:
        self.Qformer, encoder_config = self.init_Multimodal_Qformer(
            num_query_token, num_features, #self.visual_encoder.num_features,
            modulars=self.modalities,
            r=lora_rank, lora_layer=lora_layer, lora_dropout=lora_dropout)

        self.Qformer.cls = None
        self.Qformer.bert.embeddings.word_embeddings = None
        self.Qformer.bert.embeddings.position_embeddings = None
        for layer in self.Qformer.bert.encoder.layer:
            layer.output = None
            layer.intermediate = None
        self.num_query_token = num_query_token

        if 'rgb' in self.modalities:
            self.query_tokens_rgb = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_rgb.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_rgb = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)
            self.ln_rgb = nn.LayerNorm(self.visual_encoder.num_features)

        if 'flow' in self.modalities:
            self.query_tokens_flow = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_flow.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_flow = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)
            self.ln_flow = nn.LayerNorm(self.visual_encoder.num_features)

        if 'norm' in self.modalities:
            self.query_tokens_norm = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_norm.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_norm = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)
            self.ln_norm = nn.LayerNorm(self.visual_encoder.num_features)

        if 'depth' in self.modalities:
            self.query_tokens_depth = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_depth.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_depth = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)
            self.ln_depth = nn.LayerNorm(self.visual_encoder.num_features)

        if 'audio' in self.modalities:
            self.query_tokens_audio = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_audio.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_audio = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)

            self.projection_audio = nn.Linear(self.audio_encoder.num_features, num_features)
            self.ln_audio = nn.LayerNorm(num_features)

        if 'pc' in self.modalities:
            self.query_tokens_pc = nn.Parameter(
                torch.zeros(1, num_query_token, encoder_config.hidden_size))
            self.query_tokens_pc.data.normal_(mean=0.0, std=encoder_config.initializer_range)
            self.t5_proj_pc = nn.Linear(
                self.Qformer.config.hidden_size, self.t5_model.config.hidden_size)
            self.ln_pc = nn.LayerNorm(num_features)
            pos_model = PositionalEncoding1D(1408 // 3)
            x = torch.zeros(1, 256, 1408 // 3)
            self.pos_embedding = pos_model(x).squeeze().cuda()

        if 'espresso' in self.task:

            self.fusion = nn.Sequential(
                nn.Linear(2048*(len(self.modalities)-1), 2048),
                )
            self.sigmoid = nn.Sigmoid()

        self.downstream_task = downstream_task
        self.max_txt_len = 77
        answer_id = [71, 272, 205, 309, 262] # A B C D E
        self.answer_id = answer_id[:answer_num]

        self._apply_lemmatizer = apply_lemmatizer
        self._lemmatizer = None
        self.frame_num = frame_num
        self.ANS_MAP = {'A':0, 'B':1, 'C':2, 'D':3, 'E':4}

        if frame_num == 1:
            self.vid_prefix = ['Frame: ']
            self.depth_prefix = ['Depth Map: ']
            self.flow_prefix = ['Optical Flow: ']
            self.norm_prefix = ['Surface Normalization: ']
        else:
            self.vid_prefix = ['Frame {}: '.format(str(i+1)) for i in range(frame_num)]
            self.depth_prefix = ['Depth Map {}: '.format(str(i+1)) for i in range(frame_num)]
            self.flow_prefix = ['Optical Flow {}: '.format(str(i+1)) for i in range(frame_num)]
            self.norm_prefix = ['Surface Normalization {}: '.format(str(i+1)) for i in range(frame_num)]

        self.audio_prefix = ['Audio: ']
        self.pc_prefix = ['3D Model: ']

    def forward(self, samples):

        # rgb visual embedding
        qa_text, answer = samples['qa_input'], samples['qa_output']
        b = len(qa_text)

        input_embed_dict, input_att_dict = {}, {}

        for modal in self.modalities:
            input = samples[modal]
            if modal in ['rgb', 'depth', 'norm', 'flow']:
                # fix some loading issue
                if input.shape[1] == 3:
                    input = input.permute(0, 2, 1, 3, 4)

            # following 3D-LLM
            if modal == 'pc':
                with torch.cuda.amp.autocast(dtype=torch.float32):
                    pc_embeds = samples["pc_feat"]


forward

        self._apply_lemmatizer = apply_lemmatizer
        self._lemmatizer = None
        self.frame_num = frame_num
        self.ANS_MAP = {'A':0, 'B':1, 'C':2, 'D':3, 'E':4}

        if frame_num == 1:
            self.vid_prefix = ['Frame: ']
            self.depth_prefix = ['Depth Map: ']
            self.flow_prefix = ['Optical Flow: ']
            self.norm_prefix = ['Surface Normalization: ']
        else:
            self.vid_prefix = ['Frame {}: '.format(str(i+1)) for i in range(frame_num)]
            self.depth_prefix = ['Depth Map {}: '.format(str(i+1)) for i in range(frame_num)]
            self.flow_prefix = ['Optical Flow {}: '.format(str(i+1)) for i in range(frame_num)]
            self.norm_prefix = ['Surface Normalization {}: '.format(str(i+1)) for i in range(frame_num)]

        self.audio_prefix = ['Audio: ']
        self.pc_prefix = ['3D Model: ']

    def forward(self, samples):

        # rgb visual embedding
        qa_text, answer = samples['qa_input'], samples['qa_output']
        b = len(qa_text)

        input_embed_dict, input_att_dict = {}, {}

        for modal in self.modalities:
            input = samples[modal]
            if modal in ['rgb', 'depth', 'norm', 'flow']:
                # fix some loading issue
                if input.shape[1] == 3:
                    input = input.permute(0, 2, 1, 3, 4)

            # following 3D-LLM
            if modal == 'pc':
                with torch.cuda.amp.autocast(dtype=torch.float32):
                    pc_embeds = samples["pc_feat"]
                    pc = samples["pc"].long()
                    all_pcs = torch.zeros((pc_embeds.shape))
                    for j in range(pc.shape[0]):
                        pcs = []
                        for i in range(3):
                            pc_i = pc[j][:, i]
                            pcs.append(self.pos_embedding[pc_i])
                        pcs = torch.cat(pcs, -1)
                        all_pcs[j][:, :1407] = pcs
                    all_pcs = all_pcs.cuda()
                pc_embeds = pc_embeds + 0.01 * all_pcs
                atts = torch.ones(pc_embeds.size()[:-1], dtype=torch.long).to(pc_embeds.device)
                input_embed_dict[modal], input_att_dict[modal] = pc_embeds, atts
            else:
                input_embed_dict[modal], input_att_dict[modal] = self.encode_input(input, modal)

        device = input_embed_dict[list(input_embed_dict.keys())[0]].device

        input_text= self.t5_tokenizer(
            qa_text, padding="longest", truncation=True,
            max_length=self.max_txt_len, return_tensors="pt").to(device)
        input_text_embeds = self.t5_model.encoder.embed_tokens(input_text.input_ids)

        fusion_modal = []
        t5_inputs, t5_atts, t5_query = {}, {}, {}
        for modal in self.modalities:
            t5_inputs[modal], t5_atts[modal], t5_query[modal] = self.get_qformer_embedding(input_embed_dict[modal], input_att_dict[modal], device, modal, b)

        if 'rgb' in self.modalities:
            inputs_t5_rgb = t5_inputs['rgb']
            atts_t5_rgb = t5_atts['rgb']
            vid_prefix_embed, vid_prefix_mask = self.get_prefix_embedding(self.vid_prefix, b, device)
            inputs_t5 = torch.cat([vid_prefix_embed, inputs_t5_rgb], dim=2) # b, t, n_word + m, c
            atts_t5 = torch.cat([vid_prefix_mask, atts_t5_rgb], dim=2) # b, t, n_word + m

            for modal in self.modalities:
                if modal == 'rgb':
                    continue
                if modal in ['depth', 'norm', 'flow']:
                    if 'espresso' in self.task:
                        fusion_modal.append(t5_inputs[modal])
                    else:
                        inputs_t5 = torch.cat([inputs_t5, t5_inputs[modal]], dim=2)
                        atts_t5 = torch.cat([atts_t5, t5_atts[modal]], dim=2)

                if modal in ['pc']:
                    if 'espresso' in self.task:
                        pc = t5_inputs[modal]
                        pc = pc.unsqueeze(1)
                        pc = torch.repeat_interleave(pc, self.frame_num, 1) # [16, 4, 32, 2048]
                        fusion_modal.append(pc)

                if modal in ['audio']:
                    if 'espresso' in self.task:
                        audio = t5_inputs[modal]
                        audio = audio.mean(dim=1)
                        audio = audio.unsqueeze(1)
                        audio = torch.repeat_interleave(audio, self.frame_num, 1) # [16, 4, 32, 2048]
                        fusion_modal.append(audio)

            # visual only input
            if 'audio' not in self.modalities and 'pc' not in self.modalities:
                if 'espresso' in self.task:
                    fusion_modal = torch.cat(fusion_modal, dim=-1) # 16, 4, 32, 8192
                    inputs_t5_extra = self.fusion(fusion_modal)
                    inputs_t5_rgb += self.sigmoid(inputs_t5_extra) * inputs_t5_extra
                    inputs_t5 = torch.cat([vid_prefix_embed, inputs_t5_rgb], dim=2)
                    atts_t5 = torch.ones(inputs_t5.size()[:-1], dtype=torch.long).to(device)
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)
                else:
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)

            # [F1, F2, F3,..., A]
            elif 'audio' in self.modalities:
                if 'espresso' in self.task:
                    fusion_modal = torch.cat(fusion_modal, dim=-1) # 16, 4, 32, 8192
                    inputs_t5_extra = self.fusion(fusion_modal)
                    inputs_t5_rgb += self.sigmoid(inputs_t5_extra) * inputs_t5_extra
                    inputs_t5 = torch.cat([vid_prefix_embed, inputs_t5_rgb], dim=2)
                    atts_t5 = torch.ones(inputs_t5.size()[:-1], dtype=torch.long).to(device)
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)
                else:
                    audio_prefix_embed, audio_prefix_mask = self.get_prefix_embedding(self.audio_prefix, b, device)
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)
                    # seems no prefix works better for audio-video reasoning
                    inputs_t5 = torch.cat([inputs_t5, t5_inputs['audio'].reshape(b, -1, t5_inputs['audio'].shape[-1])], dim=1)
                    atts_t5 = torch.cat([atts_t5, t5_atts['audio'].reshape(b, -1)], dim=1)

            elif 'pc' in self.modalities:
                if 'espresso' in self.task:
                    fusion_modal = torch.cat(fusion_modal, dim=-1) # 16, 4, 32, 8192
                    inputs_t5_extra = self.fusion(fusion_modal)
                    inputs_t5_rgb += self.sigmoid(inputs_t5_extra) * inputs_t5_extra
                    inputs_t5 = torch.cat([vid_prefix_embed, inputs_t5_rgb], dim=2)
                    atts_t5 = torch.ones(inputs_t5.size()[:-1], dtype=torch.long).to(device)
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)
                else:
                    pc_prefix_embed, pc_prefix_mask = self.get_prefix_embedding(self.pc_prefix, b, device)
                    inputs_t5 = inputs_t5.reshape(b, -1, inputs_t5.shape[-1])
                    atts_t5 = atts_t5.reshape(b, -1)
                    inputs_t5 = torch.cat([inputs_t5, pc_prefix_embed.squeeze(), t5_inputs['pc'].reshape(b, -1, t5_inputs['pc'].shape[-1])], dim=1)
                    atts_t5 = torch.cat([atts_t5, pc_prefix_mask.squeeze(), t5_atts['pc'].reshape(b, -1)], dim=1)

        elif 'audio' in self.modalities: # audio
            inputs_t5 = t5_inputs['audio'].reshape(b, -1, t5_inputs['audio'].shape[-1])
            atts_t5 = t5_atts['audio'].reshape(b, -1)
            audio_prefix_embed, audio_prefix_mask = self.get_prefix_embedding(self.audio_prefix, b, device)
            inputs_t5 = torch.cat([audio_prefix_embed.squeeze(), inputs_t5], dim=1)
            atts_t5 = torch.cat([audio_prefix_mask.squeeze(), atts_t5], dim=1)

        elif 'pc' in self.modalities: # pc
            inputs_t5 = t5_inputs['pc'].reshape(b, -1, t5_inputs['pc'].shape[-1])
            atts_t5 = t5_atts['pc'].reshape(b, -1)
            pc_prefix_embed, pc_prefix_mask = self.get_prefix_embedding(self.audio_prefix, b, device)
            inputs_t5 = torch.cat([pc_prefix_embed.squeeze(), inputs_t5], dim=1)
            atts_t5 = torch.cat([pc_prefix_mask.squeeze(), atts_t5], dim=1)

        inputs_embeds = torch.cat([inputs_t5, input_text_embeds], dim=1)
        encoder_atts = torch.cat([atts_t5, input_text.attention_mask], dim=1)

        with torch.cuda.amp.autocast(dtype=torch.bfloat16):

            output_tokens = self.t5_tokenizer(
                answer, padding="longest", truncation=True,
                max_length=self.max_txt_len, return_tensors="pt").to(device)
            targets_qa = output_tokens.input_ids.masked_fill(
                output_tokens.input_ids == self.t5_tokenizer.pad_token_id, -100)
            output_tokens_mask = output_tokens.attention_mask

            outputs = self.t5_model(
                    inputs_embeds=inputs_embeds, attention_mask=encoder_atts,
                    decoder_attention_mask=output_tokens_mask, return_dict=True, labels=targets_qa)
            loss = outputs.loss

        return {'loss': loss}

    def encode_input(self, input, modality, training=True):

        ln = getattr(self, f"ln_{modality}")

        if modality in ['rgb', 'depth', 'flow', 'norm']:
            modality = 'visual'
        if modality in ['audio']:
            modality = 'audio'
        if modality in ['pc']:
            modality = 'pc'

        encoder = getattr(self, f"{modality}_encoder")

        if modality == 'visual':
            b, t, c, w, h = input.shape
            input = input.reshape(-1, c, w, h)
            if training:
                image_embeds = ln(encoder(input))
            else: