import math
from typing import List, Optional, Tuple

import numpy as np
import torch
from PIL import Image
from transformers import (
    Qwen2_5_VLForConditionalGeneration,
    Qwen2_5_VLModel,
    GenerationMixin,
    PreTrainedModel,
    PreTrainedTokenizer,
    Qwen2ForCausalLM,
)

from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VisionTransformerPretrainedModel

try:
    from transformers.models.qwen2_vl.image_processing_qwen2_vl import (  # noqa
        Qwen2VLImageProcessor,
    )
except ImportError:
    print('Please upgrade transformers to version 4.46.3 or higher')

from .configuration_tianjiangzhuge_VL_chat import TIANJIANGZHUGEVLChatConfig


class TianJiangZhuGe(Qwen2_5_VLForConditionalGeneration):
    config_class = TIANJIANGZHUGEVLChatConfig

    def __init__(self, config: TIANJIANGZHUGEVLChatConfig) -> None:
        super().__init__(config)
        config._attn_implementation = "flash_attention_2"
        config._attn_implementation_autoset = True
        self.model = Qwen2_5_VLModel(config)