
    Pg(                       d Z ddlZddlmZ ddlmZmZmZmZ ddl	Z	ddl
Z	ddl	mZ ddlmZ ddlmZ dd	lmZ dd
lmZmZmZmZ ddlmZ ddlmZmZmZ ddlmZmZm Z m!Z!m"Z"m#Z# ddl$m%Z%m&Z& ddl'm(Z(m)Z)m*Z*  e!jV                  e,      Z-dZ.e G d de             Z/ G d dej`                        Z1 G d dej`                        Z2 G d dej`                        Z3 G d dej`                        Z4 G d de      Z5dZ6d Z7d!Z8 G d" d#ej`                        Z9 G d$ d%e5      Z: G d& d'ej`                        Z; G d( d)ej`                        Z< G d* d+ej`                        Z= G d, d-ej`                        Z> G d. d/ej`                        Z? G d0 d1ej`                        Z@ G d2 d3ej`                        ZA G d4 d5ej`                        ZB G d6 d7e5      ZC ed8e6       G d9 d:e5e             ZDg d;ZEy)<zPyTorch InstructBLIP model.    N)	dataclass)AnyOptionalTupleUnion)nn)CrossEntropyLoss   )ACT2FN)GenerationMixin)BaseModelOutput)BaseModelOutputWithPastAndCrossAttentionsBaseModelOutputWithPooling,BaseModelOutputWithPoolingAndCrossAttentions)PreTrainedModel)apply_chunking_to_forward find_pruneable_heads_and_indicesprune_linear_layer)ModelOutputadd_start_docstrings%add_start_docstrings_to_model_forwardloggingreplace_return_docstrings	torch_int   )AutoModelForCausalLMAutoModelForSeq2SeqLM   )InstructBlipConfigInstructBlipQFormerConfigInstructBlipVisionConfigz"Salesforce/instructblip-flan-t5-xlc                      e Zd ZU dZdZeeej                        e	d<   dZ
eeej                        e	d<   dZeej                     e	d<   dZeeej                        e	d<   dZeeej                        e	d<   dee   fd	Zy)
/InstructBlipForConditionalGenerationModelOutputa)  
    Class defining the outputs of [`InstructBlipForConditionalGeneration`].

    Args:
        loss (`torch.FloatTensor`, *optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
            Language modeling loss from the language model.
        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
            Prediction scores of the language modeling head of the language model.
        vision_outputs (`BaseModelOutputWithPooling`):
            Outputs of the vision encoder.
        qformer_outputs (`BaseModelOutputWithPoolingAndCrossAttentions`):
            Outputs of the Q-Former (Querying Transformer).
        language_model_outputs (`CausalLMOutputWithPast` or `Seq2SeqLMOutput`):
            Outputs of the language model.
    Nlosslogitsvision_outputsqformer_outputslanguage_model_outputsreturnc                 H     t         fd j                         D              S )Nc              3   d   K   | ]'  }|d vr|   nt        |      j                          ) yw))r&   r'   r(   N)getattrto_tuple).0kselfs     }/var/www/html/suriana-translation/venv/lib/python3.12/site-packages/transformers/models/instructblip/modeling_instructblip.py	<genexpr>zKInstructBlipForConditionalGenerationModelOutput.to_tuple.<locals>.<genexpr>O   sB      
  WW Gq!**,-
s   -0)tuplekeysr0   s   `r1   r-   z8InstructBlipForConditionalGenerationModelOutput.to_tupleN   s%     
 YY[	
 
 	
    )__name__
__module____qualname____doc__r$   r   r   torchFloatTensor__annotations__r%   r&   r'   r(   r   r-    r6   r1   r#   r#   5   s      04D(5**+
,315FHU5,,-.526NHU../6:>OXeE$5$567>AEHU5+<+<%=>E
%* 
r6   r#   c                        e Zd Zdef fdZdej                  dededej                  fdZddej                  d	e
dej                  fd
Z xZS )InstructBlipVisionEmbeddingsconfigc                 p   t         |           || _        |j                  | _        |j
                  | _        |j                  | _        t        j                  t        j                  dd| j                              | _        t        j                  d| j                  | j                  | j                        | _        | j
                  | j                  z  dz  | _        | j                  dz   | _        t        j                  t        j                  d| j                  | j                              | _        y )Nr   r
   )in_channelsout_channelskernel_sizestrider   )super__init__rA   hidden_size	embed_dim
image_size
patch_sizer   	Parameterr;   randnclass_embeddingConv2dpatch_embeddingnum_patchesnum_positionsposition_embeddingr0   rA   	__class__s     r1   rH   z%InstructBlipVisionEmbeddings.__init__Y   s    ++ ++ ++!||EKK1dnn,MN!yyDOO\`\k\k 
 !OOt>1D!--1"$,,u{{1d>P>PRVR`R`/a"br6   
embeddingsheightwidthr)   c                    |j                   d   dz
  }| j                  j                   d   dz
  }t        j                  j	                         s||k(  r||k(  r| j                  S | j                  ddddf   }| j                  ddddf   }|j                   d   }|| j
                  z  }	|| j
                  z  }
t        |dz        }|j                  d|||      }|j                  dddd      }t        j                  j                  ||	|
fdd	
      }|j                  dddd      j                  dd|      }t        j                  ||fd      S )a   
        This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution
        images. This method is also adapted to support torch.jit tracing.

        Adapted from:
        - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
        - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
        r   Ng      ?r   r
   r   bicubicF)sizemodealign_cornersdim)shaperT   r;   jit
is_tracingrL   r   reshapepermuter   
functionalinterpolateviewcat)r0   rW   rX   rY   rR   rS   class_pos_embedpatch_pos_embedra   
new_height	new_widthsqrt_num_positionss               r1   interpolate_pos_encodingz5InstructBlipVisionEmbeddings.interpolate_pos_encodingk   s`    !&&q)A-//55a81< yy##%+*F6UZ?***11!RaR%811!QR%8r"t.
T__,	&}c'9:)11!5GI[]`a)11!Q1=--33i(	 4 
 *11!Q1=BB1b#Nyy/?;CCr6   pixel_valuesrp   c                    |j                   \  }}}}| j                  j                  j                  }| j                  |j	                  |            }|j                  d      j                  dd      }| j                  j                  |dd      j	                  |      }	t        j                  |	|gd      }
|r| j                  |
||      }n| j                  }|
|d d d |
j                  d      d d f   j	                  |      z   }
|
S )Ndtyper   r   r[   r`   )rb   rQ   weightrt   toflatten	transposerO   expandr;   rj   rp   rT   r]   )r0   rq   rp   
batch_size_rX   rY   target_dtypepatch_embedsclass_embedsrW   rT   s               r1   forwardz$InstructBlipVisionEmbeddings.forward   s    '3'9'9$
Avu++2288++LOO,O,OP#++A.88A>++22:q"EHHVYYl;C
#!%!>!>z6SX!Y!%!8!8"4Q8L*//!:L8La5O"P"S"ST`"aa
r6   F)r7   r8   r9   r!   rH   r;   Tensorintrp   r<   boolr   __classcell__rV   s   @r1   r@   r@   X   sm    c7 c$&D5<< &D &DUX &D]b]i]i &DPE$5$5 QU bgbnbn r6   r@   c                        e Zd ZdZ fdZdej                  dedefdZ	 	 ddej                  de	ej                     d	e	e
   d
eej                  e	ej                     e	eej                        f   fdZ xZS )InstructBlipAttentionz=Multi-headed attention from 'Attention Is All You Need' paperc                    t         |           || _        |j                  | _        |j
                  | _        | j                  | j                  z  | _        | j                  | j                  z  | j                  k7  r&t        d| j                   d| j                   d      | j                  dz  | _	        t        j                  |j                        | _        t        j                  | j                  d| j                  z  d      | _        |j                   ret        j"                  t%        j&                  | j                              }t        j"                  t%        j&                  | j                              }nd }d }|Qt%        j(                  |t%        j*                  |d      |f      }t        j"                  |      | j                  _        t        j                  | j                  | j                        | _        y )	Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: z).g      r
   F)bias)requires_grad)rG   rH   rA   rI   rJ   num_attention_heads	num_headshead_dim
ValueErrorscaler   Dropoutattention_dropoutdropoutLinearqkvqkv_biasrM   r;   zerosrj   
zeros_liker   
projection)r0   rA   q_biasv_biasr   rV   s        r1   rH   zInstructBlipAttention.__init__   sv   ++33$..8==4>>)T^^;MdnnM] ^NN#2'  ]]D(
zz&":":; 99T^^Q-?eL??\\%++dnn"=>F\\%++dnn"=>FFFyy&%*:*:6QV*WY_!`aHLL2DHHM))DNNDNNCr6   tensorseq_lenbszc                     |j                  ||| j                  | j                        j                  dd      j	                         S )Nr   r   )ri   r   r   rx   
contiguous)r0   r   r   r   s       r1   _shapezInstructBlipAttention._shape   s7    {{3GQQRSUVWbbddr6   hidden_states	head_maskoutput_attentionsr)   c                    |j                         \  }}}| j                  |      }|j                  ||d| j                  || j                  z        j	                  ddddd      }|d   |d   |d   }
}	}t        j                  ||	j                  dd            }|| j                  z  }t        j                  j                  |d      }| j                  |      }|||z  }t        j                  ||
      j	                  dddd      }|j                         d	d | j                  fz   }|j                  |      }| j                  |      }|r||f}|S |d	f}|S )
z#Input shape: Batch x Time x Channelr
   r   r   r      r[   r`   N)r]   r   re   r   rf   r;   matmulrx   r   r   rg   softmaxr   rJ   r   )r0   r   r   r   r   tgt_lenrJ   	mixed_qkvquery_states
key_statesvalue_statesattention_scoresattention_probscontext_layernew_context_layer_shapeoutputoutputss                    r1   r   zInstructBlipAttention.forward   su    #0"4"4"6WiHH]+	%%c7At~~yTXTbTbGbckkq!Q
	 2;1y|YWX\,j !<<j6J6J2r6RS+djj8 --//0@b/I ,,7  -	9O_lCKKAqRSUVW"/"4"4"6s";t~~>O"O%--.EF//@6?+ HNtnr6   NF)r7   r8   r9   r:   rH   r;   r   r   r   r   r   r   r   r   r   s   @r1   r   r      s    GD<eU\\ eC ec e -1,1	*||* ELL)* $D>	*
 
u||Xell3XeELL>Q5RR	S*r6   r   c                   V     e Zd Z fdZdej
                  dej
                  fdZ xZS )InstructBlipMLPc                    t         |           || _        t        |j                     | _        t        j                  |j                  |j                        | _
        t        j                  |j                  |j                        | _        y N)rG   rH   rA   r   
hidden_actactivation_fnr   r   rI   intermediate_sizefc1fc2rU   s     r1   rH   zInstructBlipMLP.__init__   sd    #F$5$5699V//1I1IJ99V55v7I7IJr6   r   r)   c                 l    | j                  |      }| j                  |      }| j                  |      }|S r   )r   r   r   r0   r   s     r1   r   zInstructBlipMLP.forward   s4    /**=9/r6   r7   r8   r9   rH   r;   r   r   r   r   s   @r1   r   r      s$    KU\\ ell r6   r   c            
            e Zd Zdef fdZ	 ddej                  dej                  dee   de	ej                     fdZ xZS )	InstructBlipEncoderLayerrA   c                 D   t         |           |j                  | _        t	        |      | _        t        j                  | j                  |j                        | _	        t        |      | _        t        j                  | j                  |j                        | _        y Neps)rG   rH   rI   rJ   r   	self_attnr   	LayerNormlayer_norm_epslayer_norm1r   mlplayer_norm2rU   s     r1   rH   z!InstructBlipEncoderLayer.__init__  sm    ++.v6<<F<Q<QR"6*<<F<Q<QRr6   r   attention_maskr   r)   c                     |}| j                  |      }| j                  |||      \  }}||z   }|}| j                  |      }| j                  |      }||z   }|f}|r||fz  }|S )aI  
        Args:
            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
            attention_mask (`torch.FloatTensor`): attention mask of size
                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
                `(config.encoder_attention_heads,)`.
            output_attentions (`bool`, *optional*):
                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                returned tensors for more detail.
        )r   r   r   )r   r   r   r   )r0   r   r   r   residualattn_weightsr   s          r1   r   z InstructBlipEncoderLayer.forward  s      !((7&*nn'$/ '5 '
#|
 &0 ((7/%0 "&Gr6   r   )r7   r8   r9   r   rH   r;   r   r   r   r   r<   r   r   r   s   @r1   r   r     sY    S1 S -2	$||$ $ $D>	$
 
u  	!$r6   r   c                   *    e Zd ZdZeZdZdZg dZd Z	y)InstructBlipPreTrainedModelz
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    blipT)InstructBlipQFormerEmbeddingsr   %InstructBlipQFormerMultiHeadAttentionInstructBlipQFormerSelfOutputc                 d   | j                   j                  }t        |t        j                        s4t        |t        j
                        st        |t        j                        rc|j                  j                  j                  d|       t        |d      r0|j                  $|j                  j                  j                          t        |t              rt        | j                   d      r:t        | j                   t              s | j                   j                  j                  }t        j                   j#                  |j$                  d|       t        j                   j#                  |j&                  d|       yt        |t        j(                        rJ|j                  j                  j                          |j                  j                  j+                  d       yt        |t        j                        r2|j                  %|j                  j                  j                          yyy)zInitialize the weightsg        )meanstdr   Nvision_config      ?)rA   initializer_range
isinstancer   rP   	Embeddingr   ru   datanormal_hasattrr   zero_r@   r!   r   inittrunc_normal_rT   rO   r   fill_)r0   modulefactors      r1   _init_weightsz)InstructBlipPreTrainedModel._init_weightsG  sp   ..fbii(Jvr||,LPZ[aceclclPmMM&&CV&<vv&6;;+B  &&(f:;t{{O4ZUm=n22DDGG!!&";";#6!RGG!!&"8"8s!O-KK""$MM$$S)		*v{{/FKK""$ 0G*r6   N)
r7   r8   r9   r:   r   config_classbase_model_prefixsupports_gradient_checkpointing_no_split_modulesr   r>   r6   r1   r   r   5  s'    
 &L&*#%r6   r   aE  
    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
    etc.)

    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
    and behavior.

    Parameters:
        config ([`InstructBlipConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
a  
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`InstructBlipProcessor`]. See
            [`InstructBlipProcessor.__call__`] for details.
        output_attentions (`bool`, *optional*):
            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
            tensors for more detail.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
        interpolate_pos_encoding (`bool`, *optional*, defaults to `False`):
            Whether to interpolate the pre-trained position encodings.
a  
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`InstructBlipProcessor`]. See
            [`InstructBlipProcessor.__call__`] for details.

        qformer_input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Indices of input sequence tokens in the vocabulary of the Q-Former. Input tokens can optionally be provided
            to serve as text prompt, which the Q-Former model will encode.

            Indices can be obtained using [`InstructBlipProcessor`]. See [`InstructBlipProcessor.__call__`] for
            details.

            [What are input IDs?](../glossary#input-ids)

        qformer_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.

            [What are attention masks?](../glossary#attention-mask)

        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Indices of input sequence tokens in the vocabulary of the language model. Input tokens can optionally be
            provided to serve as text prompt, which the language model can continue.

            Indices can be obtained using [`InstructBlipProcessor`]. See [`InstructBlipProcessor.__call__`] for
            details.

            [What are input IDs?](../glossary#input-ids)

        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.

            [What are attention masks?](../glossary#attention-mask)

        decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
            Indices of decoder input sequence tokens in the vocabulary of the language model. Only relevant in case an
            encoder-decoder language model (like T5) is used.

            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details. [What are decoder input IDs?](../glossary#decoder-input-ids)

        decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
            Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
            be used by default.

            Only relevant in case an encoder-decoder language model (like T5) is used.

        output_attentions (`bool`, *optional*):
            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
            tensors for more detail.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
        interpolate_pos_encoding (`bool`, *optional*, defaults to `False`):
            Whether to interpolate the pre-trained position encodings.
        use_cache (`bool`, *optional*):
            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
            `past_key_values`).
c                        e Zd ZdZdef fdZ	 	 	 	 d
deej                     dee	   dee	   dee	   de
eef   f
d	Z xZS )InstructBlipEncodera  
    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
    [`InstructBlipEncoderLayer`].

    Args:
        config (`InstructBlipConfig`):
            The corresponding vision configuration for the `InstructBlipEncoder`.
    rA   c                     t         |           || _        t        j                  t        |j                        D cg c]  }t        |       c}      | _        d| _	        y c c}w r   )
rG   rH   rA   r   
ModuleListrangenum_hidden_layersr   layersgradient_checkpointing)r0   rA   r{   rV   s      r1   rH   zInstructBlipEncoder.__init__  sP    mmuU[UmUmOn$o!%=f%E$op&+# %ps   A#r   r   output_hidden_statesreturn_dictr)   c                    ||n| j                   j                  }||n| j                   j                  }||n| j                   j                  }|rdnd}|rdnd}|}t	        | j
                        D ]`  \  }	}
|r||fz   }| j                  r+| j                  r| j                  |
j                  |||      }n |
|||      }|d   }|sX||d   fz   }b |r||fz   }|st        d |||fD              S t        |||      S )a  
        Args:
            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
                Embedded representation of the inputs. Should be float, not int tokens.
            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

                - 1 for tokens that are **not masked**,
                - 0 for tokens that are **masked**.

                [What are attention masks?](../glossary#attention-mask)
            output_attentions (`bool`, *optional*):
                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                returned tensors for more detail.
            output_hidden_states (`bool`, *optional*):
                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
                for more detail.
            return_dict (`bool`, *optional*):
                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
        Nr>   r   r   r   c              3   &   K   | ]	  }||  y wr   r>   r.   vs     r1   r2   z.InstructBlipEncoder.forward.<locals>.<genexpr>  s     eqWXWdes   )last_hidden_stater   
attentions)rA   r   r   use_return_dict	enumerater   r   training_gradient_checkpointing_func__call__r3   r   )r0   inputs_embedsr   r   r   r   encoder_statesall_attentionsr   idxencoder_layerlayer_outputss               r1   r   zInstructBlipEncoder.forward  sA   8 2C1N-TXT_T_TqTq$8$D $++JjJj 	 &1%<k$++B]B]30d%"+DKK"8 	FC#!/=2B!B**t}} $ A A!**!"%	! !.!"&7! *!,M !/=3C2E!E)	F,  +}.>>Ne]NN$Seee+>Vd
 	
r6   )NNNN)r7   r8   r9   r:   r   rH   r   r;   r   r   r   r   r   r   r   r   s   @r1   r   r     s    ,1 , 26,0/3&*C
 !.C
 $D>	C

 'tnC
 d^C
 
uo%	&C
r6   r   c                        e Zd ZdZeZdef fdZ ee       e	e
e      	 	 	 	 	 ddeej                     dee   dee   dee   ded	eee
f   fd
              Zd Z xZS )InstructBlipVisionModelrq   rA   c                     t         |   |       || _        |j                  }t	        |      | _        t        |      | _        t        j                  ||j                        | _        | j                          y r   )rG   rH   rA   rI   r@   rW   r   encoderr   r   r   post_layernorm	post_init)r0   rA   rJ   rV   s      r1   rH   z InstructBlipVisionModel.__init__  s]     &&	6v>*62 ll9&:O:OPr6   output_typer   r   r   r   rp   r)   c                    ||n| j                   j                  }||n| j                   j                  }||n| j                   j                  }|t	        d      | j                  ||      }| j                  ||||      }|d   }| j                  |      }|dddddf   }	| j                  |	      }	|s
||	f|dd z   S t        ||	|j                  |j                        S )z
        Returns:

        Nz You have to specify pixel_values)rp   )r  r   r   r   r   r   )r   pooler_outputr   r   )rA   r   r   r   r   rW   r  r  r   r   r   )
r0   rq   r   r   r   rp   r   encoder_outputsr   pooled_outputs
             r1   r   zInstructBlipVisionModel.forward(  s    2C1N-TXT_T_TqTq$8$D $++JjJj 	 &1%<k$++B]B]?@@Ogh,,'/!5#	 ' 
 ,A. //0AB)!Q'2++M:%}58KKK)/')77&11	
 	
r6   c                     | j                   S r   )rW   r5   s    r1   get_input_embeddingsz,InstructBlipVisionModel.get_input_embeddingsX  s    r6   NNNNF)r7   r8   r9   main_input_namer!   r   rH   r   $INSTRUCTBLIP_VISION_INPUTS_DOCSTRINGr   r   r   r;   r<   r   r   r   r   r  r   r   s   @r1   r
  r
    s    $O+L	7 	 ++OP+ETlm 59,0/3&*).,
u001,
 $D>,
 'tn	,

 d^,
 #',
 
u00	1,
 n Q,
\r6   r
  c                   R     e Zd Zd fd	Zd Zd Zd Zd Zd Z	 	 	 	 	 	 d	dZ	 xZ
S )
r   c                    t         |           || _        |j                  |j                  z  dk7  r0t        |d      s$t        d|j                  |j                  fz        |j                  | _        t        |j                  |j                  z        | _        | j                  | j                  z  | _	        t        j                  |j                  | j                        | _        |r_t        j                  |j                  | j                        | _        t        j                  |j                  | j                        | _        n^t        j                  |j                  | j                        | _        t        j                  |j                  | j                        | _        t        j                   |j"                        | _        t'        |dd      | _        | j(                  dk(  s| j(                  dk(  rF|j*                  | _        t        j,                  d|j*                  z  d	z
  | j                        | _        d
| _        y )Nr   embedding_sizezLThe hidden size (%d) is not a multiple of the number of attention heads (%d)position_embedding_typeabsoluterelative_keyrelative_key_queryr   r   F)rG   rH   rA   rI   r   r   r   r   attention_head_sizeall_head_sizer   r   queryencoder_hidden_sizekeyvaluer   attention_probs_dropout_probr   r,   r  max_position_embeddingsr   distance_embeddingsave_attentionr0   rA   is_cross_attentionrV   s      r1   rH   z.InstructBlipQFormerMultiHeadAttention.__init__]  s    : ::a?PVXhHi^%%v'A'ABC 
 $*#=#= #&v'9'9F<V<V'V#W !558P8PPYYv1143E3EF
yy!;!;T=O=OPDH6#=#=t?Q?QRDJyy!3!3T5G5GHDH6#5#5t7I7IJDJzz&"E"EF'.v7PR\']$''>9T=Y=Y]q=q+1+I+ID(&(ll1v7U7U3UXY3Y[_[s[s&tD##r6   c                     || _         y r   attn_gradients)r0   r/  s     r1   save_attn_gradientsz9InstructBlipQFormerMultiHeadAttention.save_attn_gradientsy  s
    ,r6   c                     | j                   S r   r.  r5   s    r1   get_attn_gradientsz8InstructBlipQFormerMultiHeadAttention.get_attn_gradients|  s    """r6   c                     || _         y r   attention_map)r0   r5  s     r1   save_attention_mapz8InstructBlipQFormerMultiHeadAttention.save_attention_map  s
    *r6   c                     | j                   S r   r4  r5   s    r1   get_attention_mapz7InstructBlipQFormerMultiHeadAttention.get_attention_map  s    !!!r6   c                     |j                         d d | j                  | j                  fz   } |j                  | }|j	                  dddd      S )Nr[   r   r   r   r
   )r]   r   r!  ri   rf   )r0   xnew_x_shapes      r1   transpose_for_scoresz:InstructBlipQFormerMultiHeadAttention.transpose_for_scores  sN    ffhsmt'?'?AYAY&ZZAFFK yyAq!$$r6   c                    |d u}|rC| j                  | j                  |            }	| j                  | j                  |            }
|}n|y| j                  | j                  |            }	| j                  | j                  |            }
t        j                  |d   |	gd      }	t        j                  |d   |
gd      }
n@| j                  | j                  |            }	| j                  | j                  |            }
| j                  |      }| j                  |      }|	|
f}t        j                  ||	j                  dd            }| j                  dk(  s| j                  dk(  rF|j                         d   }t        j                  |t        j                  |j                  	      j                  dd      }t        j                  |t        j                  |j                  	      j                  dd      }||z
  }| j                  || j                  z   dz
        }|j!                  |j"                  
      }| j                  dk(  rt        j$                  d||      }||z   }nE| j                  dk(  r6t        j$                  d||      }t        j$                  d|	|      }||z   |z   }|t'        j(                  | j*                        z  }|j"                  }|||z   } t-        j.                  d      |      j!                  |      }|r8| j0                  r,| j3                  |       |j5                  | j6                         | j9                  |      }|||z  }t        j                  ||
      }|j;                  dddd      j=                         }|j                         d d | j>                  fz   } |j                  | }|r||fn|f}||fz   }|S )Nr   r   r`   r   r[   r   r  r   rt   devicers   zbhld,lrd->bhlrzbhrd,lrd->bhlrr
   ) r<  r%  r&  r;   rj   r#  r   rx   r  r]   arangelongr?  ri   r)  r(  rv   rt   einsummathsqrtr!  r   Softmaxr*  r6  register_hookr0  r   rf   r   r"  )r0   r   r   r   encoder_hidden_statesencoder_attention_maskpast_key_valuer   r,  	key_layervalue_layermixed_query_layerquery_layerr   
seq_lengthposition_ids_lposition_ids_rdistancepositional_embeddingrelative_position_scoresrelative_position_scores_queryrelative_position_scores_keyattention_scores_dtyper   attention_probs_droppedr   r   r   s                               r1   r   z-InstructBlipQFormerMultiHeadAttention.forward  s    3$>11$((;P2QRI33DJJ?T4UVK3N'11$((=2IJI33DJJ}4MNK		>!#4i"@aHI))^A%6$D!LK11$((=2IJI33DJJ}4MNK JJ}5//0AB#[1 !<<Y5H5HR5PQ''>9T=Y=Y]q=q&++-a0J"\\*EJJ}OcOcdiijlnopN"\\*EJJ}OcOcdiijkmopN%6H#'#:#:8dFbFb;bef;f#g #7#:#:ARAR#:#S ++~=+0<<8H+Wk+l(#36N#N --1EE16>NP[]q1r./4||<LiYm/n,#36T#TWs#s +dii8P8P.QQ!1!7!7%/.@ -"**,-=>AABXY$"5"5##O4))$*B*BC #',,"?  &=	&I#%<kJ%--aAq9DDF"/"4"4"6s";t?Q?Q>S"S***,CD6G=/2mM]^--r6   r   NNNNNF)r7   r8   r9   rH   r0  r2  r6  r8  r<  r   r   r   s   @r1   r   r   \  s;    $8-#+"% "#Tr6   r   c                   n     e Zd Z fdZdej
                  dej
                  dej
                  fdZ xZS )r   c                 (   t         |           t        j                  |j                  |j                        | _        t        j                  |j                  |j                        | _        t        j                  |j                        | _
        y r   )rG   rH   r   r   rI   denser   r   r   hidden_dropout_probr   rU   s     r1   rH   z&InstructBlipQFormerSelfOutput.__init__  s`    YYv1163E3EF
f&8&8f>S>STzz&"<"<=r6   r   input_tensorr)   c                 r    | j                  |      }| j                  |      }| j                  ||z         }|S r   r[  r   r   r0   r   r]  s      r1   r   z%InstructBlipQFormerSelfOutput.forward  7    

=1]3}|'CDr6   r   r   s   @r1   r   r     1    >U\\  RWR^R^ r6   r   c                       e Zd Zd fd	Zd Z	 	 	 	 	 	 ddej                  deej                     deej                     deej                     deej                     dee	e	ej                           d	ee
   d
e	ej                     fdZ xZS )InstructBlipQFormerAttentionc                     t         |           t        ||      | _        t	        |      | _        t               | _        y r   )rG   rH   r   	attentionr   r   setpruned_headsr+  s      r1   rH   z%InstructBlipQFormerAttention.__init__  s3    >vGYZ3F;Er6   c                 >   t        |      dk(  ry t        || j                  j                  | j                  j                  | j
                        \  }}t        | j                  j                  |      | j                  _        t        | j                  j                  |      | j                  _        t        | j                  j                  |      | j                  _	        t        | j                  j                  |d      | j                  _        | j                  j                  t        |      z
  | j                  _        | j                  j                  | j                  j                  z  | j                  _        | j
                  j                  |      | _        y )Nr   r   r`   )lenr   rf  r   r!  rh  r   r#  r%  r&  r   r[  r"  union)r0   headsindexs      r1   prune_headsz(InstructBlipQFormerAttention.prune_heads  s   u:?74>>55t~~7Y7Y[_[l[l
u
  2$..2F2FN/0B0BEJ1$..2F2FN.t{{/@/@%QO .2^^-O-ORUV[R\-\*'+~~'I'IDNNLnLn'n$ --33E:r6   r   r   r   rG  rH  rI  r   r)   c           	      p    | j                  |||||||      }| j                  |d   |      }	|	f|dd  z   }
|
S )Nr   r   )rf  r   )r0   r   r   r   rG  rH  rI  r   self_outputsattention_outputr   s              r1   r   z$InstructBlipQFormerAttention.forward
  sW     ~~!"
  ;;|AF#%QR(88r6   r   rX  )r7   r8   r9   rH   rn  r;   r   r   r<   r   r   r   r   r   s   @r1   rd  rd    s    ";* 7;15=A>BDH,1|| !!2!23 E--.	
  ((9(9: !)):): ; !uU->->'?!@A $D> 
u||	r6   rd  c                   V     e Zd Z fdZdej
                  dej
                  fdZ xZS )InstructBlipQFormerIntermediatec                    t         |           t        j                  |j                  |j
                        | _        t        |j                  t              rt        |j                     | _        y |j                  | _        y r   )rG   rH   r   r   rI   r   r[  r   r   strr   intermediate_act_fnrU   s     r1   rH   z(InstructBlipQFormerIntermediate.__init__$  s]    YYv1163K3KL
f''-'-f.?.?'@D$'-'8'8D$r6   r   r)   c                 J    | j                  |      }| j                  |      }|S r   )r[  rv  r   s     r1   r   z'InstructBlipQFormerIntermediate.forward,  s&    

=100?r6   r   r   s   @r1   rs  rs  #  s#    9U\\ ell r6   rs  c                   n     e Zd Z fdZdej
                  dej
                  dej
                  fdZ xZS )InstructBlipQFormerOutputc                 (   t         |           t        j                  |j                  |j
                        | _        t        j                  |j
                  |j                        | _        t        j                  |j                        | _        y r   )rG   rH   r   r   r   rI   r[  r   r   r   r\  r   rU   s     r1   rH   z"InstructBlipQFormerOutput.__init__4  s`    YYv779K9KL
f&8&8f>S>STzz&"<"<=r6   r   r]  r)   c                 r    | j                  |      }| j                  |      }| j                  ||z         }|S r   r_  r`  s      r1   r   z!InstructBlipQFormerOutput.forward:  ra  r6   r   r   s   @r1   ry  ry  3  rb  r6   ry  c                   @     e Zd Z fdZ	 	 	 	 	 	 	 ddZd Zd Z xZS )InstructBlipQFormerLayerc                 f   t         |           |j                  | _        d| _        t	        |      | _        || _        ||j                  z  dk(  rt	        |d      | _        d| _	        nd| _	        t        |      | _        t        |      | _        t        |      | _        t        |      | _        y )Nr   r   T)r,  F)rG   rH   chunk_size_feed_forwardseq_len_dimrd  rf  	layer_idxcross_attention_frequencycrossattentionhas_cross_attentionrs  intermediatery  r   intermediate_queryoutput_queryr0   rA   r  rV   s      r1   rH   z!InstructBlipQFormerLayer.__init__B  s    '-'E'E$5f="v7771<">vZ^"_D'+D$',D$;FC/7"A&"I5f=r6   c	           
         ||d d nd }	| j                  |||||	      }
|
d   }|
dd }|
d   }|dkD  r|d d d |d d f   }| j                  r1|t        d      | j                  ||||||      }|d   }||dd z   }t	        | j
                  | j                  | j                  |      }|j                  d   |kD  r~t	        | j                  | j                  | j                  |d d |d d d f         }t        j                  ||gd      }n,t	        | j                  | j                  | j                  |      }|f|z   }||fz   }|S )	Nr   )r   rI  r   r   r[   z>encoder_hidden_states must be given for cross-attention layersr   r`   )rf  r  r   r  r   feed_forward_chunk_queryr  r  rb   feed_forward_chunkr;   rj   )r0   r   r   r   rG  rH  rI  r   query_lengthself_attn_past_key_valueself_attention_outputsrq  r   present_key_valuequery_attention_outputcross_attention_outputslayer_outputlayer_output_texts                     r1   r   z InstructBlipQFormerLayer.forwardV  s    :H9S>"1#5Y] !%/3 "0 "
 2!4(2.226!%5a,6I%J"''(0$%eff*.*=*=*")*&7 +> +' *A)C&!$;Ab$AA4--,,  &	L  %%a(<7$=++00$$$Qq%89	%!  %yy,8I)JPQR4'',,   	L  /G+.00r6   c                 L    | j                  |      }| j                  ||      }|S r   )r  r   r0   rq  intermediate_outputr  s       r1   r  z+InstructBlipQFormerLayer.feed_forward_chunk  s,    "//0@A{{#68HIr6   c                 L    | j                  |      }| j                  ||      }|S r   )r  r  r  s       r1   r  z1InstructBlipQFormerLayer.feed_forward_chunk_query  s.    "556FG(()<>NOr6   )NNNNNFr   )r7   r8   r9   rH   r   r  r  r   r   s   @r1   r}  r}  A  s0    >. "#EN
r6   r}  c                   :     e Zd Z fdZ	 	 	 	 	 	 	 	 	 	 ddZ xZS )InstructBlipQFormerEncoderc           	          t         |           || _        t        j                  t        |j                        D cg c]  }t        ||       c}      | _        d| _	        y c c}w r   )
rG   rH   rA   r   r   r   r   r}  layerr   r  s      r1   rH   z#InstructBlipQFormerEncoder.__init__  sW    ]]JOPVPhPhJijY%fi8j

 ',# ks   A$c                 j   |	rdnd }|rdnd }|rdnd }|rdnd }t        | j                  j                        D ]  }| j                  |   }|	r||fz   }|||   nd }|||   nd }t	        | j                  dd      rF| j
                  r:|rt        j                  d       d}| j                  |j                  |||||      }n |||||||||      }|d   }|r	||d   fz  }|s||d   fz   }|j                  s||d   fz   } |	r||fz   }|
st        d	 |||||fD              S t        |||||
      S )Nr>   r   FzZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...r   r[   r   r   c              3   $   K   | ]  }|| 
 y wr   r>   r   s     r1   r2   z5InstructBlipQFormerEncoder.forward.<locals>.<genexpr>  s      
 = 
s   )r   past_key_valuesr   r   cross_attentions)r   rA   r   r  r,   r   loggerwarningr  r  r  r3   r   )r0   r   r   r   rG  rH  r  	use_cacher   r   r   r  all_hidden_statesall_self_attentionsall_cross_attentionsnext_decoder_cacheilayer_modulelayer_head_maskrI  r  s                        r1   r   z"InstructBlipQFormerEncoder.forward  s    #7BD$5b4%6rD#,R$t{{445 (	VA::a=L#$58H$H!.7.CilO3B3N_Q/TXNt{{$<eDNNt !&I $ A A ))!"#)*! !-!"#)*"% 	! *!,M"}R'8&::" &9]1=M<O&O#33+?=QRCSBU+U(Q(	VT   1]4D D 
 "&%'(
 
 
 9+.+*1
 	
r6   )
NNNNNNFFTr   )r7   r8   r9   rH   r   r   r   s   @r1   r  r    s.    , "#"S
r6   r  c                   2     e Zd ZdZ fdZ	 	 	 	 ddZ xZS )r   z;Construct the embeddings from word and position embeddings.c                 P   t         |           t        j                  |j                  |j
                  |j                        | _        t        j                  |j                  |j
                        | _	        t        j                  |j
                  |j                        | _        t        j                  |j                        | _        | j!                  dt#        j$                  |j                        j'                  d      d       t)        |dd      | _        || _        y )	N)padding_idxr   position_ids)r   r[   F)
persistentr  r  )rG   rH   r   r   
vocab_sizerI   pad_token_idword_embeddingsr(  position_embeddingsr   r   	layernormr   r\  r   register_bufferr;   r@  ry   r,   r  rA   rU   s     r1   rH   z&InstructBlipQFormerEmbeddings.__init__  s    !||F,=,=v?Q?Q_e_r_rs#%<<0N0NPVPbPb#c f&8&8f>S>STzz&"<"<= 	ELL)G)GHOOPWXej 	 	
 (/v7PR\']$r6   c                    ||j                         d   }nd}|&| j                  d d |||z   f   j                         }|k| j                  |      }| j                  dk(  r/| j                  |j                  |j                              }||z   }|t        j                  ||fd      }n|}|j                  | j                  j                  j                        }| j                  |      }| j                  |      }|S )Nr   r   r  r`   )r]   r  cloner  r  r  rv   r?  r;   rj   r  ru   rt   r   )r0   	input_idsr  query_embedspast_key_values_lengthrN  rW   r  s           r1   r   z%InstructBlipQFormerEmbeddings.forward  s      ")!,JJ,,Q0FVlIl0l-lmssuL --i8J++z9&*&>&>|zO`O`?a&b#'*==
'"YYj'AqI
%J]]4>>#8#8#>#>?
^^J/
\\*-
r6   )NNNr   )r7   r8   r9   r:   rH   r   r   r   s   @r1   r   r     s    E$  r6   r   c                       e Zd ZdZdef fdZd Zd Zd Z	 dde	j                  dee   d	e	j                  d
ede	j                  f
dZ	 	 	 	 	 	 	 	 	 	 	 dde	j                   dee	j$                     dee	j                      dee	j                     dee	j$                     dee	j$                     dee	j$                     deeee	j$                           dee   dee   dee   dee   deee	j$                     ef   fdZ xZS )InstructBlipQFormerModelz
    Querying Transformer (Q-Former), used in InstructBLIP. Slightly modified from BLIP-2 as it also takes the
    instruction as input.
    rA   c                     t         |   |       || _        t        |      | _        t        |      | _        | j                          y r   )rG   rH   rA   r   rW   r  r  r  rU   s     r1   rH   z!InstructBlipQFormerModel.__init__A  s9     7?1&9r6   c                 .    | j                   j                  S r   rW   r  r5   s    r1   r  z-InstructBlipQFormerModel.get_input_embeddingsK  s    ...r6   c                 &    || j                   _        y r   r  r0   r&  s     r1   set_input_embeddingsz-InstructBlipQFormerModel.set_input_embeddingsN  s    */'r6   c                     |j                         D ]7  \  }}| j                  j                  |   j                  j	                  |       9 y)z
        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
        class PreTrainedModel
        N)itemsr  r  rf  rn  )r0   heads_to_pruner  rl  s       r1   _prune_headsz%InstructBlipQFormerModel._prune_headsQ  sE    
 +002 	CLE5LLu%//;;EB	Cr6   r   input_shaper?  	has_queryr)   c                    |j                         dk(  r|dddddddf   }n=|j                         dk(  r|ddddddf   }nt        d| d|j                   d      |j                  | j                        }d|z
  d	z  }|S )
a>  
        Makes broadcastable attention and causal masks so that future and masked tokens are ignored.

        Arguments:
            attention_mask (`torch.Tensor`):
                Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
            input_shape (`Tuple[int]`):
                The shape of the input to the model.
            device: (`torch.device`):
                The device of the input to the model.

        Returns:
            `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`.
        r
   Nr   z!Wrong shape for input_ids (shape z) or attention_mask (shape )rs   r   g     )ra   r   rb   rv   rt   )r0   r   r  r?  r  extended_attention_masks         r1   get_extended_attention_maskz4InstructBlipQFormerModel.get_extended_attention_maskY  s    . 1$&4Qa]&C#!Q& '5QdA5E&F#3K=@[\j\p\p[qqrs  #:"<"<4::"<"N#&)@#@H"L&&r6   r  r  r  r   rG  rH  r  r  r   r   r   c                    |
|
n| j                   j                  }
||n| j                   j                  }||n| j                   j                  }||t	        d      |,|d   d   j
                  d   | j                   j                  z
  nd}||j
                  d   nd}| j                  ||||      }|j                         dd }|\  }}|j                  }|t        j                  |||z   f|      }| j                  |||      }|t        |t              r|d   j                         \  }}}n|j                         \  }}}||f}t        |t              r|D cg c]  }| j                  |       }}n?|)t        j                  ||      }| j                  |      }n| j                  |      }nd}| j!                  || j                   j"                        }| j%                  |||||||	|
|||	      }|d   }|dddddf   }|s
||f|dd z   S t'        |||j(                  |j*                  |j,                  |j.                  
      S c c}w )a  
        encoder_hidden_states  (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
            the model is configured as a decoder.
        encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
            the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.
        past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of:
            shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): Contains precomputed key and
            value hidden states of the attention blocks. Can be used to speed up decoding. If `past_key_values` are
            used, the user can optionally input only the last `decoder_input_ids` (those that don't have their past key
            value states given to this model) of shape `(batch_size, 1)` instead of all `decoder_input_ids` of shape
            `(batch_size, sequence_length)`.
        use_cache (`bool`, *optional*):
            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
            `past_key_values`).
        Nz7You have to specify query_embeds when input_ids is Noner   r   r   )r  r  r  r  r[   )r?  )
r   r   rG  rH  r  r  r   r   r   r  )r   r  r  r   r   r  )rA   r   r   r   r   rb   r  rW   r]   r?  r;   onesr  r   listinvert_attention_maskget_head_maskr   r  r   r  r   r   r  )r0   r  r   r  r  r   rG  rH  r  r  r   r   r   r  r  embedding_outputr  rz   rN  r?  r  encoder_batch_sizeencoder_sequence_lengthr{   encoder_hidden_shapemaskencoder_extended_attention_maskr  sequence_outputr  s                                 r1   r   z InstructBlipQFormerModel.forward  s   D 2C1N-TXT_T_TqTq$8$D $++JjJj 	 &1%<k$++B]B]!5VWW JYIdOAq!''*T[[-E-EEjk 	 1=0H|))!,a??%%#9	 + 
 '++-cr2!,
J!((!"ZZ*jCY6Y)ZdjkN #'"B"B>S^`f"g !,/6AVWXAYA^A^A`>"$;QAVA[A[A]>"$;Q$68O#P 0$7`v2wX\43M3Md3S2w/2w'/).4HQW)X&262L2LMc2d/262L2LMc2d/.2+ &&y$++2O2OP	,,2"7#B+/!5#% ' 
 *!,'1a0#]3oab6III;-'+;;)77&11,==
 	
G 3xs   !Ir   )NNNNNNNNNNN)r7   r8   r9   r:   r    rH   r  r  r  r;   r   r   r   r?  r   r  
LongTensorr   r<   r   r   r   r   r   s   @r1   r  r  ;  s   
8 /0C  )')' 3Z)' 	)'
 )' 
)'\ 7;37/315=A>BEI$(,0/3&*x
##x
 !!2!23x
 u//0	x

 u||,x
 E--.x
  ((9(9:x
 !)):): ;x
 "%e.?.?(@"ABx
 D>x
 $D>x
 'tnx
 d^x
 
uU&&')UU	Vx
r6   r  a  
    InstructBLIP Model for generating text given an image and an optional text prompt. The model consists of a vision
    encoder, Querying Transformer (Q-Former) and a language model.

    One can optionally pass `input_ids` to the model, which serve as a text prompt, to make the language model continue
    the prompt. Otherwise, the language model starts generating text from the [BOS] (beginning-of-sequence) token.
    c            !           e Zd ZeZdZdZdZdZdgZ	def fdZ
d Zd Zd	 Zd
ej                  fdZd Zd Zd Zd Z ee       eee      	 	 	 	 	 	 	 	 	 	 	 ddej6                  dej6                  deej:                     deej6                     deej:                     deej:                     deej:                     dee   dee   deej:                     dee   dedee   d
ee ef   fd              Z! ejD                         	 	 	 	 	 d dej6                  deej:                     deej:                     deej:                     deej:                     ded
ej:                  fd       Z# xZ$S )!$InstructBlipForConditionalGenerationrq   TFquery_tokensrA   c                 4   t         |   |       t        |j                        | _        t        j                  t        j                  d|j                  |j                  j                              | _        t        |j                        | _        t        j                  |j                  j                  |j                   j                        | _        |j$                  r t'        j(                  |j                         }nt+        j(                  |j                         }|j,                  %| j,                  j/                  |j,                         |j0                  %| j0                  j/                  |j0                         || _        | j5                          y )Nr   )rG   rH   r
  r   vision_modelr   rM   r;   r   num_query_tokensqformer_configrI   r  r  qformerr   text_configlanguage_projectionuse_decoder_only_language_modelr   from_configr   r   extend_keep_in_fp32_moduleslanguage_modelr  )r0   rA   r  rV   s      r1   rH   z-InstructBlipForConditionalGeneration.__init__  s'    3F4H4HILLQ8O8OQWQfQfQrQr)st/0E0EF#%99V-B-B-N-NPVPbPbPnPn#o 111==f>P>PQN2>>v?Q?QRN++7"")).*J*JK//;&&--n.R.RS, 	r6   c                 6    | j                   j                         S r   )r  r  r5   s    r1   r  z9InstructBlipForConditionalGeneration.get_input_embeddings+  s    ""7799r6   c                 :    | j                   j                  |       y r   )r  r  r  s     r1   r  z9InstructBlipForConditionalGeneration.set_input_embeddings.  s    007r6   c                 :    | j                   j                  |       y r   )r  set_output_embeddings)r0   new_embeddingss     r1   r  z:InstructBlipForConditionalGeneration.set_output_embeddings1  s    11.Ar6   r)   c                 6    | j                   j                         S r   )r  get_output_embeddingsr5   s    r1   r  z:InstructBlipForConditionalGeneration.get_output_embeddings4  s    ""88::r6   c                 6    | j                   j                         S r   )r  get_encoderr5   s    r1   r  z0InstructBlipForConditionalGeneration.get_encoder7      ""..00r6   c                 6    | j                   j                         S r   )r  get_decoderr5   s    r1   r  z0InstructBlipForConditionalGeneration.get_decoder:  r  r6   c                     | j                   j                  s_| j                  j                  | j                  j                  _        | j                  j                  | j                  j                  _        y y r   )rA   r  r  sharedr  embed_tokensdecoderr5   s    r1   _tie_weightsz1InstructBlipForConditionalGeneration._tie_weights=  sT    {{::7;7J7J7Q7QD''47;7J7J7Q7QD''4 ;r6   c                    | j                   }t        |      dkD  r:d|vr6t        j                  j	                         dkD  rt
        j                  d       t        | j                  d      rd| j                  j                  _
        yy)z
        Some pre-processing hacks to make the model `accelerate` compatible. Check
        https://github.com/huggingface/transformers/pull/21707 for more details.
        r   r  a  The `language_model` is not in the `hf_device_map` dictionary and you are running your script in a multi-GPU environment. this may lead to unexpected behavior when using `accelerate`. Please pass a `device_map` that contains `language_model` to remove this warning. Please refer to https://github.com/huggingface/blog/blob/main/accelerate-large-models.md for more details on creating a `device_map` for large models._hf_hookTN)hf_device_maprj  r;   cudadevice_countr  r  r   r  r  io_same_device)r0   r  s     r1   _preprocess_acceleratez;InstructBlipForConditionalGeneration._preprocess_accelerateB  sy    
 **}!&6m&KPUPZPZPgPgPilmPmNNM 4&&
3:>D((7 4r6   r  qformer_input_idsqformer_attention_maskr  r   decoder_input_idsdecoder_attention_maskr   r   labelsr   rp   r  c                 Z   ||n| j                   j                  }| j                  |||	||      }|d   }t        j                  |j                         dd t        j                  |j                        }| j                  j                  |j                  d   dd      }t        j                  |j                         dd t        j                  |j                        }|t        j                  |      }t        j                  ||gd      }| j                  |||||||	|      }|d   ddd|j                  d      ddf   }| j                  |      }t        j                  |j                         dd t        j                  |j                        } | j                  j!                         |      }|t        j                  |      }t#        | j                   d	d      K|| j                   j$                  k(  j'                  d      j)                  |      }|j+                         ||<   nyt,        j/                  d
       t        j                  ||j1                  |j                        gd      }t        j                  ||j1                  |j                        gd      }| j                   j2                  r| j                  ||||	||      }|r|j4                  n|d   }d}|
|
j1                  |j                        }
|dd|
j                  d       dddf   }|dddddf   j7                         }|
dddf   j7                         j1                  |j                        }t9        d      } ||j;                  d| j                   j<                  j>                        |j;                  d            }n@| j                  ||||||	||
|	      }|r|j@                  n|d   }|r|j4                  n|d   }|s||||f}||f|z   S |S tC        |||||      S )a  
        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Labels for computing the language modeling loss. Indices should be in `[-100, 0, ..., config.vocab_size -
            1]`. All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ...,
            config.vocab_size]`

        Returns:

        Examples:

        ```python
        >>> from transformers import InstructBlipProcessor, InstructBlipForConditionalGeneration
        >>> import torch
        >>> from PIL import Image
        >>> import requests

        >>> model = InstructBlipForConditionalGeneration.from_pretrained("Salesforce/instructblip-vicuna-7b")
        >>> processor = InstructBlipProcessor.from_pretrained("Salesforce/instructblip-vicuna-7b")

        >>> device = "cuda" if torch.cuda.is_available() else "cpu"
        >>> model.to(device)  # doctest: +IGNORE_RESULT

        >>> url = "https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg"
        >>> image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
        >>> prompt = "What is unusual about this image?"
        >>> inputs = processor(images=image, text=prompt, return_tensors="pt").to(device)

        >>> outputs = model.generate(
        ...     **inputs,
        ...     do_sample=False,
        ...     num_beams=5,
        ...     max_length=256,
        ...     min_length=1,
        ...     top_p=0.9,
        ...     repetition_penalty=1.5,
        ...     length_penalty=1.0,
        ...     temperature=1,
        ... )
        >>> generated_text = processor.batch_decode(outputs, skip_special_tokens=True)[0].strip()
        >>> print(generated_text)
        The unusual aspect of this image is that a man is ironing clothes on the back of a yellow SUV, which is parked in the middle of a busy city street. This is an unconventional approach to ironing clothes, as it requires the man to balance himself and his ironing equipment on top of the vehicle while navigating through traffic. Additionally, the presence of taxis and other vehicles in the scene further emphasizes the unusual nature of this situation.
        ```N)rq   r   r   r   rp   r   r[   r>  r   r`   )r  r   r  rG  rH  r   r   r   image_token_indexA  Expanding inputs for image tokens in InstructBLIP should be done in processing. Please follow instruction here (https://gist.github.com/zucchini-nlp/e9f20b054fa322f84ac9311d9ab67042) to update your InstructBLIP model. Using processors without these attributes in the config is deprecated and will throw an error in v4.50.)r  r   r   r   r   r  .r   )	reduction)	r  r   r  r  r   r   r   r  r  )r$   r%   r&   r'   r(   )"rA   r   r  r;   r  r]   rA  r?  r  ry   rb   	ones_likerj   r  r  r  r  r,   r  	unsqueeze	expand_asrw   r  warning_oncerv   r  r%   r   r	   ri   r  r  r$   r#   ) r0   rq   r   r  r  r   r  r  r   r   r  r   rp   r  r&   image_embedsimage_attention_maskr  query_attention_maskquery_outputsquery_outputlanguage_model_inputslanguage_model_attention_maskr  special_image_maskr   r%   r$   shift_logitsshift_labelsloss_fctr   s                                    r1   r   z,InstructBlipForConditionalGeneration.forwardV  s.   | &1%<k$++B]B] **%/!5#%= + 
 &a(  %zz,*;*;*=cr*B%**]i]p]pq ((//0B0B10Er2N$zz,*;*;*=cr*B%**]i]p]pq!)%*__5F%G"!&,@BX+Y_`!a'1%".#7/!5# % 	
 %Q'+A\->->q-A+A1(DE !% 8 8 F(-

!&&("-UZZH]HdHd)
% C++@@B9M!"__Y7N 4;; 3T:F"+t{{/L/L"L!W!WXZ![!e!efs!t0E0M0M0OM,-z
 "II'<m>N>NOdOkOk>l&mstuM"YY.0A0AB_BfBf0ghnoN ;;66))+-"3%9'# * G (3W^^
FD!6==1FKKN?#4a 78%c3B3k2==?%c12g99;>>v}}M ,f= 1 1"dkk6M6M6X6X Y[g[l[lmo[pq))+-"3'="3%9'# * 
G $/7<<GAJD'2W^^
FnmWEF)-)9TGf$EvE>))#*
 	
r6   c                 x   t        | d      r| j                          |j                  d   }| j                  |d|      j                  }	t        j                  |	j                         dd t
        j                  |	j                        }
| j                  j                  |	j                  d   dd      }t        j                  |j                         dd t
        j                  |	j                        }|t        j                  |      }t        j                  ||gd	      }| j                  ||||	|
d
      }|j                  ddd|j                  d      ddf   }| j                  |      }t        j                  |j                         dd t
        j                  |j                        }|| j                   j"                  j$                  g}t'        | j                   dd      1| j                   j(                  g| j                   j*                  z  |z   }t        j,                  |gt
        j                  |	j                        }|j/                  |d      }|t        j                  |      } | j1                         |      }t'        | j                   dd      d|| j                   j(                  k(  j3                  d      j5                  |      }|j7                         j9                  |j                        ||<   nt:        j=                  d       t        j                  ||j9                  |j                        gd	      }t        j                  ||j9                  |j                        gd	      }| j>                  j                   j@                  sM|jC                  dd      |j                  d   z   dz
  |d<   |jC                  dd      |j                  d   z   |d<   ||d}| j>                  j                   j@                  s||d<    | j>                  jD                  di ||}|S )a  
        Overrides `generate` function to be able to use the model as a conditional generator.

        Args:
            pixel_values (`torch.FloatTensor` of shape (batch_size, num_channels, height, width)):
                Input images to be processed.
            qformer_input_ids (`torch.LongTensor` of shape (batch_size, sequence_length), *optional*):
                The sequence used as a prompt to be fed to the Q-Former module.
            qformer_attention_mask (`torch.LongTensor` of shape (batch_size, sequence_length), *optional*):
                Mask to avoid performing attention on padding token indices.
            input_ids (`torch.LongTensor` of shape (batch_size, sequence_length), *optional*):
                The sequence used as a prompt for the generation.
            attention_mask (`torch.LongTensor` of shape (batch_size, sequence_length), *optional*):
                Mask to avoid performing attention on padding token indices.
            interpolate_pos_encoding (`bool`, *optional*, defaults to `False`):
                Whether to interpolate the positional encoding of the image embeddings.

        Returns:
            captions (list): A list of strings of length batch_size * num_captions.
        r  r   T)r   rp   Nr[   r>  r   r`   )r  r   r  rG  rH  r   r  r  
max_length   
min_length)r  r   r  r>   )#r   r  rb   r  r   r;   r  r]   rA  r?  r  ry   r	  rj   r  r  rA   r  bos_token_idr,   r  r  r   repeatr  r
  r  rw   rv   r  r  r  is_encoder_decodergetgenerate)r0   rq   r   r  r  r   rp   generate_kwargsrz   r  r  r  r  r  r  r  language_attention_maskstart_tokensr  r  inputsr   s                         r1   r   z-InstructBlipForConditionalGeneration.generate  s   > 4)'')!''*
((%= ) 
 
	 	  %zz,*;*;*=cr*B%**]i]p]pq((//0B0B10Er2N$zz,*;*;*=cr*B%**]i]p]pq!)%*__5F%G"!&,@BX+Y_`!a'1%".#7 % 
 %66q:PL<M<Ma<P:PRS7ST $ 8 8 F"'**!&&("-UZZH]HdHd#
  KK33@@ALt{{$7>J $ = =>A]A]]`lll^5::lNaNabI!((Q7I!"__Y7N3113I> 4;; 3T:F"+t{{/L/L"L!W!WXZ![!e!efs!t0E0M0M0O0R0RS`SgSg0hM,-z
 "II'<m>N>NOdOkOk>l&mstuM"YY(.*;*;<S<Z<Z*[\bcN &&--@@#''b9<Q<W<WXY<ZZ]^^  - 1@0C0CLRS0TWlWrWrstWu0u-#0NS""))<<"+F;.$%%..KK?Kr6   )NNNNNNNNNFNr  )%r7   r8   r9   r   r   r  _supports_cache_class_supports_static_cache_supports_quantized_cacher  rH   r  r  r  r   Moduler  r  r  r  r  r   INSTRUCTBLIP_INPUTS_DOCSTRINGr   r#   r!   r;   r<   r   r  r   r   r   r   no_gradr   r   r   s   @r1   r  r    st    &L$O ! %+,1 4:8B;ryy ;11R
?( ++HICRj >B15598<=A,0/3-1&*).$(f
''f
 !,,f
 !))9)9 :	f

 E--.f
 !!1!12f
 $E$4$45f
 !))9)9 :f
 $D>f
 'tnf
 ))*f
 d^f
 #'f
 D>f
 
uEE	Ff
 Jf
P U]]_ 9==A0459).i''i $E$4$45i !))9)9 :	i
 E,,-i !!1!12i #'i 
		i ir6   r  )r  r   r  r
  )Fr:   rC  dataclassesr   typingr   r   r   r   r;   torch.utils.checkpointr   torch.nnr	   activationsr   
generationr   modeling_outputsr   r   r   r   modeling_utilsr   pytorch_utilsr   r   r   utilsr   r   r   r   r   r   autor   r   configuration_instructblipr   r    r!   
get_loggerr7   r  _CHECKPOINT_FOR_DOCr#   r(  r@   r   r   r   r   INSTRUCTBLIP_START_DOCSTRINGr  r)  r   r
  r   r   rd  rs  ry  r}  r  r   r  r  __all__r>   r6   r1   <module>r;     s   "  ! . .    % ! )  . l l  ? o o 
		H	%:  
k 
 
BG299 GVNBII Ndbii  -ryy -`$%/ $%N  ( $"B! LS
")) S
n@9 @FBBII BLBII .299 .dbii  		 dryy dP\
 \
~0BII 0fA
: A
H  !	c+F c	cLr6   