o Zh@sLddlZddlmZddlmZmZmZmZmZddl Z ddl m Z ddl m Z ddl mZddlmZmZmZmZdd lmZdd lmZdd lmZdd lmZmZmZdd lmZm Z ddl!m"Z"ddl#m$Z$m%Z%m&Z&m'Z'm(Z(ddl)m*Z*ddl+m,Z,m-Z-m.Z.e&rddl/m m0Z1e'2e3Z4e$Gddde Z5eGdddeZ6eGdddeZ7eGdddeZ8Gddde j9Z:de j;ded*e>fd+d,Z?Gd-d.d.e j9Z@Gd/d0d0e j9ZAGd1d2d2eZBGd3d4d4e j9ZCe$Gd5d6d6e5ZDGd7d8d8e j9ZEGd9d:d:e j9ZFGd;d<dd>e j9ZHGd?d@d@e j9ZIGdAdBdBe j9ZJGdCdDdDe j9ZKGdEdFdFe j9ZLGdGdHdHe j9ZMe$dIdJGdKdLdLe5ZNGdMdNdNe j9ZOGdOdPdPe j9ZPe$dQdJGdRdSdSe5ZQGdTdUdUe5eZRgdVZSdS)XN) dataclass)CallableListOptionalTupleUnion)nn)ACT2FN)Cache)%ClassifierFreeGuidanceLogitsProcessorGenerationMixinGenerationModeLogitsProcessorList)GenerateDecoderOnlyOutput)FlashAttentionKwargs)GradientCheckpointingLayer)BaseModelOutputBaseModelOutputWithPooling ModelOutput)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)auto_docstringcan_return_tupleis_torch_availablelogging torch_int) AutoModel) JanusConfigJanusVisionConfigJanusVQVAEConfigc@sFeZdZeZdZdZdgZddgZdZ dZ dZ dZ dZ dZddZd S) JanusPreTrainedModelmodelTZLlamaDecoderLayerpast_key_values causal_maskFcCst|jdr |jjjn|jj}t|tjtjfr0|jj j d|d|j dur.|j j dSdSt|tj tjfrH|j j |jj ddSt|tjrg|jj j d|d|jduri|jj |j dSdSdS)N vision_config)meanstdg?)hasattrconfigr(Zinitializer_range isinstancerLinearConv2dweightdataZnormal_biasZzero_ GroupNorm LayerNormZfill_ EmbeddingZ padding_idx)selfmoduler+r9W/var/www/auris/lib/python3.10/site-packages/transformers/models/janus/modeling_janus.py _init_weightsFs$      z"JanusPreTrainedModel._init_weightsN)__name__ __module__ __qualname__r! config_classZbase_model_prefixZsupports_gradient_checkpointing_no_split_modulesZ_skip_keys_device_placementZ_supports_flash_attn_2Z_supports_sdpaZ_supports_quantized_cacheZ_supports_cache_class_supports_static_cacheZ!_supports_param_buffer_assignmentr;r9r9r9r:r$8s r$c@s2eZdZUdZdZeejed<dZ ejed<dS)JanusVQVAEOutputaM Base class for Janus VQ-VAE mode model outputs. Args: decoded_pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): Reconstructed pixel values after encoding and decoding the input. embedding_loss (`torch.FloatTensor`): Embedding loss. Ndecoded_pixel_valuesembedding_loss) r<r=r>__doc__rCrtorch FloatTensor__annotations__rDr9r9r9r:rBYs  rBc@seZdZUdZdZeejed<dZ ee e ejed<dZ ee ejed<dZ ee ejed<dZ ee ejed<dS)JanusBaseModelOutputWithPasta Base class for Janus model's outputs that may also contain a past key/values (to speed up sequential decoding). Args: last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1, hidden_size)` is output. past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and optionally if `config.is_encoder_decoder=True` 2 additional tensors of shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if `config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. image_hidden_states (`tuple(torch.FloatTensor)`, *optional*): Tuple of `torch.FloatTensor` (one for the output of the image embeddings, `(batch_size, num_images, sequence_length, hidden_size)`. image_hidden_states of the model produced by the vision encoder, and optionally by the perceiver Nlast_hidden_stater& hidden_states attentionsimage_hidden_states)r<r=r>rErJrrFrGrHr&rrKrLrMr9r9r9r:rIhs $rIc@seZdZUdZdZeejed<dZ eejed<dZ ee ejed<dZ ee ejed<dZee ejed<dZee ejed<dS) JanusCausalLMOutputWithPasta Base class for Janus causal language model (or autoregressive) outputs. Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Language modeling loss (for next-token prediction). logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. image_hidden_states (`tuple(torch.FloatTensor)`, *optional*): Tuple of `torch.FloatTensor` (one for the output of the image embeddings, `(batch_size, num_images, sequence_length, hidden_size)`. image_hidden_states of the model produced by the vision encoder, and optionally by the perceiver Nlosslogitsr&rKrLrM)r<r=r>rErOrrFrGrHrPr&rrKrrLrMr9r9r9r:rNs  rNcs\eZdZdeffdd Zdejdededejfdd Zdd ejd e dejfd dZ Z S)JanusVisionEmbeddingsr-cst||_|j|_|j|_|j|_tj|j |j|j|jdd|_ |j|jd|_ |j |_ t |j |j|_|jdt|j ddddS)NZvalid) in_channels out_channels kernel_sizestridepaddingr position_ids)r F) persistent)super__init__r- hidden_size embed_dimZ image_size patch_sizerr0 num_channelspatch_embedding num_patches num_positionsr6position_embeddingZregister_bufferrFZarangeexpandr7r- __class__r9r:r[s "zJanusVisionEmbeddings.__init__ embeddingsheightwidthreturnc Cs|jd}|jjjd}tjs||kr||kr||jS|jjd}|jd}||j}||j} t |d} | d| | |}| dddd}t j j||| fddd }| dddddd|}|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 and no class embeddings. 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 rrXg?r rZbicubicF)sizemodeZ align_corners)shapercr1rFZjit is_tracingrW unsqueezer^rreshapepermuter functional interpolateview) r7rhrirjrarbZpatch_pos_embeddimZ new_heightZ new_widthZsqrt_num_positionsr9r9r:interpolate_pos_encodings&      z.JanusVisionEmbeddings.interpolate_pos_encodingF pixel_valuesrwc Csh|j\}}}}|jjj}||j|d}|ddd}|r(||||} n||j } || }|S)N)dtyperr ) rnr`r1rytoflatten transposerwrcrW) r7rxrw_rirjZ target_dtypeZ patch_embedsrhZ pos_embedsr9r9r:forwards  zJanusVisionEmbeddings.forwardF) r<r=r>r"r[rFTensorintrwboolr~ __classcell__r9r9rfr:rQs$&rQrKn_reprkcCs^|j\}}}}|dkr |S|dddddddddf|||||}||||||S)z This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) r N)rnrdrq)rKrbatchZnum_key_value_headsslenhead_dimr9r9r: repeat_kv s 0rr)r8querykeyvalueattention_maskscalingdropoutc Kst||j}t||j} t||dd|} |dur3|ddddddd|jdf} | | } tjj| dtj d |j } tjj | ||j d} t| | } | dd} | | fS)Nrr rX)rvry)ptrainingr )rnum_key_value_groupsrFmatmulr|rnrrssoftmaxZfloat32rzryrr contiguous) r8rrrrrrkwargs key_states value_states attn_weightsr' attn_outputr9r9r:eager_attention_forwards & rc sXeZdZdZdeffdd Z  d dejdeejdeejd e e fd d Z Z S) JanusVisionAttentionz(Attention Class for Janus Vision Encoderr-csLt||_|j|_|j|_|j|j|_|j|j|jkr-td|jd|jd|jd|_ |j |_ |j }|j }d|_ d|_tj|j|j|j|jd|_tj|j|j|j|jd|_tj|j|j|j|jd|_t|j|j|_|dkrt|nt|_ |rt|jnt|_|rt|j|_dSt|_dS) Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: z).Fr r3r)rZr[r-r\r]Znum_attention_heads num_headsr ValueErrorscaleattention_dropoutprojection_dropoutZ use_qk_norm is_causalrrr/Zattention_biasq_projk_projv_projprojection_layerDropoutZIdentityr5q_normk_norm)r7r-Z proj_dropoutZqk_normrfr9r:r[3s0   $zJanusVisionAttention.__init__NrKroutput_attentionsrcKsl|\}}}||}||} ||} |d|j|j}||}| d|j|j} || } ||||j|j dd}| |||j|j dd} | |||j|j dd} t } |j j dkr||j j dkrv|ddrvtdnt|j j } | ||| | |f|jsd n|j|j|jd |\} } | |||j} || }||}|r|| f}|S|df}|S) NrXr reagerZsdparFz`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.r))rrr)rlrrrrqrrrrr|rurr-Z_attn_implementationgetlogger warning_oncerrrrrr]rr)r7rKrrr batch_sizeseq_lenr} query_statesrrZattention_interfacerroutputoutputsr9r9r:r~PsL           zJanusVisionAttention.forward)NN) r<r=r>rEr"r[rFrrrrr~rr9r9rfr:r0s rcs8eZdZdeffdd ZdejdejfddZZS)JanusVisionMLPr-csrt||_t|j|j|_t|j|_ t |j|j|_ t |j|j|_ t |j|_t |j|_dSN)rZr[r-rr\Z mlp_ratioZintermediate_sizer hidden_act activation_fnrr/fc1fc2rZhidden_dropout_ratedropout1dropout2rerfr9r:r[s  zJanusVisionMLP.__init__rKrkcCs6||}||}||}||}||}|Sr)rrrrrr7rKr9r9r:r~s     zJanusVisionMLP.forward) r<r=r>r"r[rFrr~rr9r9rfr:rs rc sNeZdZdeffdd Z d dejdejdeede ej fd d Z Z S) JanusVisionEncoderLayerr-csXt|j|_tj|j|jd|_t||_ tj|j|jd|_ t ||_ ||_ dSN)eps)rZr[r\r]rr5layer_norm_eps layer_norm1r self_attn layer_norm2rmlpr-rerfr9r:r[s    z JanusVisionEncoderLayer.__init__FrKrrrkcCsb|}||}|j|||d\}}||}|}||}||}||}|f}|r/||f7}|S)a= Args: hidden_states (`torch.FloatTensor`): Input to the layer of shape `(batch, seq_len, embed_dim)`. attention_mask (`torch.FloatTensor`): Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where padding elements are indicated by very large negative values. output_attentions (`bool`, *optional*, defaults to `False`): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. )rKrr)rrrr)r7rKrrresidualrrr9r9r:r~s      zJanusVisionEncoderLayer.forwardr) r<r=r>r"r[rFrrrrrGr~rr9r9rfr:rs rc sZeZdZdZdeffdd Ze   d deej dee dee d e fd d Z Z S) JanusVisionEncoderz Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a [`JanusVisionEncoderLayer`]. Args: config: JanusVisionConfig r-cs:t|_tfddtjD|_d|_dS)Ncsg|]}tqSr9)r.0r}r-r9r: sz/JanusVisionEncoder.__init__..F) rZr[r-r ModuleListrangenum_hidden_layerslayersgradient_checkpointingrerfrr:r[s   zJanusVisionEncoder.__init__Nrroutput_hidden_statesrkc Cs|dur|n|jj}|dur|n|jj}|rdnd}|rdnd}|}|jD]}|r.||f}||||d} | d}|rB|| df}q%|rJ||f}t|||dS)ad Args: inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert `input_ids` indices into associated vectors than the model's internal embedding lookup matrix. 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. Nr9)rrr )rJrKrL)r-rrrr) r7 inputs_embedsrrrZencoder_statesZall_attentionsrKZ encoder_layerZ layer_outputsr9r9r:r~s2     zJanusVisionEncoder.forwardNNN)r<r=r>rEr"r[rrrFrrrr~rr9r9rfr:rs rcs~eZdZdZeZdeffdd Ze     ddee j dee dee d ee d e d e e eff d d ZddZZS)JanusVisionModelrxr-csJt|||_|j}t||_t||_tj ||j d|_ | dSr) rZr[r-r\rQrhrencoderrr5rpost_layernorm post_init)r7r-r]rfr9r:r[s    zJanusVisionModel.__init__NFrr return_dictrwrkc Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}|dur&td|j||d}|j||||d}|d}||}|dddddf} || } |s[|| f|ddSt|| |j |j dS)Nz You have to specify pixel_values)rw)rrrrrr )rJZ pooler_outputrKrL) r-rrZuse_return_dictrrhrrrrKrL) r7rxrrrrwrKZencoder_outputsrJZ pooled_outputr9r9r:r~*s2   zJanusVisionModel.forwardcC|jSr)rhr7r9r9r:get_input_embeddingsUz%JanusVisionModel.get_input_embeddings)NNNNF)r<r=r>main_input_namer"r?r[rrrFrGrrrrr~rrr9r9rfr:rs0   *rc*eZdZdeffdd ZddZZS)JanusVisionAlignerMLPr-cNttjj|_tfddtdj D|_ t j |_ dS)Ncg|] }tjjqSr9rr/projection_dimrrr9r:r_z2JanusVisionAlignerMLP.__init__..r )rZr[rr/r\rrrrdepth hidden_layersr rrrerfrr:r[Z zJanusVisionAlignerMLP.__init__cC,||}|jD] }||}||}q|Srrrrr7rKlayerr9r9r:r~c    zJanusVisionAlignerMLP.forward)r<r=r>r"r[r~rr9r9rfr:rY rcsLeZdZdZdeffdd ZdejfddZdej d ej fd d Z Z S) JanusVQVAEVectorQuantizera A module for vector quantization using learned embedding vectors. This module implements the quantization process similar to te one described in the VQ-VAE (Vector Quantized Variational AutoEncoder) paper. It quantizes continuous input vectors into discrete codebook vectors, which are learned during training. Current implementation improves over previous ones by avoiding costly matrix multiplications and allowing for post-hoc remapping of indices. r-csLt|j|_|j|_t|dd|_t|j|j|_ |j gd|_ dS)Nbetag?r) rZr[num_embeddingsr] embedding_dimgetattrrrr6 embeddingraquant_state_dimsrerfr9r:r[vs z"JanusVQVAEVectorQuantizer.__init__ hidden_statec Cs|dddd}|d|j}tj|ddddtj|jjddddtd ||jj dd}tj |dd}|||j }t | |d|jt || d}||| }|dddd}|||fS) Nrrr r rXT)rvZkeepdimrvz bd,dn->bn)rrrrurrFsumrr1Zeinsumr|Zargminrnr*detachr)r7rZhidden_state_flattenedZ distancesZmin_encoding_indiceshidden_state_quantrOr9r9r:r~s  z!JanusVQVAEVectorQuantizer.forward image_tokensrkcCsb|jd}|jjjd}||}tj|ddd}||g|j|R}|dddd}|S)NrrXr)rrvr r ) rnrr1F normalizerurrrr)r7rrZemb_dimrr9r9r:get_codebook_entrys  z,JanusVQVAEVectorQuantizer.get_codebook_entry) r<r=r>rEr#r[rFrr~ LongTensorrGrrr9r9rfr:rks   rcs*eZdZ  dfdd ZddZZS)JanusVQVAEResnetBlockNFcst||_|dur|n||_||_tjjd|ddd|_tjj ||dddd|_ tjjd|ddd|_ tj |j |_ tjj ||dddd|_|j|jkro|jratjj ||dddd|_dStjj ||dddd|_dSdS) N ư>TZ num_groupsr_rZaffiner r rTrUrVr)rZr[rRrSuse_conv_shortcutrFrr4norm1r0conv1norm2rrconv2 conv_shortcut nin_shortcut)r7r-rRrSr rfr9r:r[s  zJanusVQVAEResnetBlock.__init__cCs|}||}|t|9}||}||}|t|9}||}||}|j|jkr@|j r;| |}||S| |}||Sr) rrFsigmoidrrrr rRrSrr r )r7rKrr9r9r:r~s        zJanusVQVAEResnetBlock.forward)NFr<r=r>r[r~rr9r9rfr:rs rc$eZdZfddZddZZS)JanusVQVAEAttnBlockcst||_tjjd|ddd|_tjj||dddd|_tjj||dddd|_ tjj||dddd|_ tjj||dddd|_ dS)NrrTrr rr) rZr[rRrFrr4normr0qkvproj_outr7rRrfr9r:r[s zJanusVQVAEAttnBlock.__init__c Cs|}||}||}||}||}|j\}}}} ||||| ddd}||||| }t||} | t |d} t j | dd} ||||| }| ddd} t|| |||| } | | } || S)Nrrr rr) rrrrrnrqrrrFZbmmrrrr) r7rKrrrrrchannelsrirjrrr9r9r:r~s       zJanusVQVAEAttnBlock.forwardr r9r9rfr:rs  rcr)JanusVQVAEConvDownsamplecs$ttj||dddd|_dS)Nr rrr)rZr[rr0convrrfr9r:r[s z!JanusVQVAEConvDownsample.__init__cCs tj|dddd}||}|S)N)rr rr Zconstantr)padrmr)rrrrr9r9r:r~s z JanusVQVAEConvDownsample.forwardr r9r9rfr:r rcr)JanusVQVAEConvUpsamplecs&ttjj||dddd|_dS)Nr r r)rZr[rFrr0rrrfr9r:r[s zJanusVQVAEConvUpsample.__init__cCstj|ddd}||}|S)Ng@Znearest)Z scale_factorrm)rrtrrr9r9r:r~ s zJanusVQVAEConvUpsample.forwardr r9r9rfr:rrrcs<eZdZdedeffdd ZdejdejfddZZ S) JanusVQVAEMidBlockr-rcs8tt|||d|_t||_t|||d|_dS)Nr-rRrS)rZr[rblock_1rattn_1block_2)r7r-rrfr9r:r[s   zJanusVQVAEMidBlock.__init__rKrkcC"||}||}||}|Sr)rrr rr9r9r:r~   zJanusVQVAEMidBlock.forward) r<r=r>r#rr[rFrr~rr9r9rfr:rsrcs,eZdZfddZdejfddZZS)JanusVQVAEEncoderc sntt|j|_|j|_|j}|j}|j}|j }|j}t j j ||dddd|_ dt|}||_t |_t|jD]T}t } t } |||} |||} t|jD]} | t|| | d| } ||jdkrt| t| qXt }| |_| |_||jdkrt| |_|j|q=t|| |_t j jd| ddd |_t j j | |rd |n|dddd|_dS) Nr r r)r rrrTrr) rZr[lenchannel_multipliernum_resolutionsnum_res_blocks base_channelsrR double_latentlatent_channelsrFrr0conv_intuplein_channel_multiplierrdownrappendrrModuleblockattnr downsamplermidr4norm_outconv_out)r7r-r(rRr)r*r%r-i_levelr1r2block_in block_outi_blockr.rfr9r:r['sX         zJanusVQVAEEncoder.__init__rxcCs||g}t|jD]C}t|jD]'}|j|j||d}t|j|jdkr4|j|j||}||q||jdkrN||j| |dq |d}| |}| |}|t |9}||}|S)NrXrr )r+rr&r'r.r1r$r2r/r3r4r5rFr r6)r7rxrKr7r:rrJr9r9r:r~Zs$     zJanusVQVAEEncoder.forward)r<r=r>r[rFrr~rr9r9rfr:r#&s 3r#cs2eZdZfddZdejdejfddZZS)JanusVQVAEDecoderc sPtt|j|_|j|_|j}|j}|j}||j|jd}t j j ||dddd|_ t |||_t |_tt|jD]N}t }t }||j|} t|jdD]} |t||| d| }||jdkrt|t|qXt } || _|| _|dkrt|| _|j| q@t j jd|ddd |_t j j ||dddd|_dS) Nr r rrrrrTr)rZr[r$r%r&r'r(r*rSrFrr0r+rr4rupreversedrr/rrr0r1r2rupsampler4r5r6) r7r-r(r*rSr8r7r1r2r9r:r<rfr9r:r[tsD     zJanusVQVAEDecoder.__init__rrkcCs||}||}t|jD]9}t|jdD] }|j|j||}t|j|jdkr8|j|j||}q||jdkrH|j| |}q| |}|t |9}| |}|S)Nr r)r+r4rr&r'r<r1r$r2r>r5rFr r6)r7rr7r:r9r9r:r~s    zJanusVQVAEDecoder.forward)r<r=r>r[rFrGr~rr9r9rfr:r;ss .r;aG The VQ-VAE model used in Janus for encoding/decoding images into discrete tokens. This model follows the "Make-a-scene: Scene-based text-to-image generation with human priors" paper from [ Oran Gafni, Adam Polyak, Oron Ashual, Shelly Sheynin, Devi Parikh, and Yaniv Taigman](https://arxiv.org/abs/2203.13131). )Z custom_introcseZdZeZgdZdZdeffdd Zdej fddZ dej d ej fd d Z e edej d eej ej ffd d ZZS) JanusVQVAE)rrrrxr-cspt|t||_t||_tj|j |j d|_ tj|j |j d|_ | t||_d|_|dS)Nr F)rZr[r#rrquantizerFrr0r*r] quant_convpost_quant_convevalr;decoderrrrerfr9r:r[s     zJanusVQVAE.__init__cCs.||}||}||\}}}|||fSr)rrAr@)r7rxrKquantZemb_lossindicesr9r9r:encodes   zJanusVQVAE.encoderrkcCsr|jd|jjd|jjdkr'td|jjd|jjdd|jd|j|}||}||}|S)aG Decodes quantized token IDs into pixel values. Args: image_tokens (torch.LongTensor): Batch of token IDs. Returns: pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): Pixel values decoded from the token IDs. r rz4Expected `image_tokens` to have shape `(batch_size, z)`, but got shape `z`.)rnr@rrrrBrD)r7rZcodebook_entryrKrxr9r9r:decodes"    zJanusVQVAE.decodecCs:|jd}||\}}}|||d}t||}|S)NrrX)rnrGrHrurB)r7rxrrErDrFrCrr9r9r:r~s  zJanusVQVAE.forward)r<r=r>r#r?r@rr[rFrrGrGrHrrrr~rr9r9rfr:r?sr?cr)JanusVQVAEAlignerMLPr-cr)Ncrr9rrrr9r:rrz1JanusVQVAEAlignerMLP.__init__..r )rZr[rr/r]rrrrrrr rrrerfrr:r[rzJanusVQVAEAlignerMLP.__init__cCrrrrr9r9r:r~rzJanusVQVAEAlignerMLP.forward)r<r=r>r#r[r~rr9r9rfr:rIrrIcs<eZdZdZdeffdd ZdejdejfddZ Z S) JanusVQVAEHeadzOHead used for sampling tokens in image generation, replacing the usual lm head.r-cs>tt|j|j|_t|j|_ t|j|j |_ dSr) rZr[rr/Zimage_token_embed_dimrrr rrr vision_headrerfr9r:r[s  zJanusVQVAEHead.__init__rKrkcCr!r)rrrKrr9r9r:r~r"zJanusVQVAEHead.forward) r<r=r>rEr#r[rFrZtensorr~rr9r9rfr:rJsrJzl The Janus model which consists of a siglip vision backbone, a Llama language model and a VQ model. cseZdZdeffdd ZddZddZdd Zee dd e j d e j de e jde e j de ede e j de e j de ede ede edeee jffddZZS) JanusModelr-cst|||_t|j|_t|jj|_t |j |_ t |j jj|j jj|_t|j j|_t|j j|_tj|jd|_d|_|dS)NrF)rZr[r-rZ _from_configr( vision_modelralignerr?Z vq_configvqmodelrr6rr]generation_embeddingsrIgeneration_alignerrJgeneration_headrZ from_config text_configlanguage_modelrrrerfr9r:r[$s  zJanusModel.__init__cCs |jSr)rTrrr9r9r:r9 zJanusModel.get_input_embeddingscCs|j|dSr)rTset_input_embeddingsr7rr9r9r:rV<szJanusModel.set_input_embeddingscCs||}||j}|Sr)rMrNrJ)r7rx image_embedsr9r9r:get_image_features?s  zJanusModel.get_image_featuresNr input_idsrxrrWr&cache_positionr use_cacherrlogits_to_keepc  Ks<| dur| n|jj} | dur| n|jj} |du|duAr td|jr/|jr/|r/tdd}|dur;|dur;td|durE||}|durw| |} ||jj k}|j d}| d|}| ddd|}||j|j}|||}|jd|||||| | || d | }t|j|j|j|j|dur| ndd}|S) NzaYou cannot specify both input_ids and inputs_embeds at the same time, and must specify either onezZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...FzdYou cannot specify both pixel_values and inputs_embeds at the same time, and must specify either onerX) rrrWr&r\rrr[r])rJr&rKrLrMr9)r-rrrrrrrrrYZimage_token_idrnrqrprdrzdeviceryZmasked_scatterrTrIrJr&rKrL)r7rZrxrrWr&r[rr\rrr]rrXZimage_attention_maskr]Zimage_featuresZ lm_outputrr9r9r:r~Ds`         zJanusModel.forward) NNNNNNNNNNr)r<r=r>r!r[rrVrYrrrFrrGrrr rrrr~rr9r9rfr:rLsR    rLcsjeZdZddgZdZdeffdd ZddZd d Zd e j d e j fd dZ ddZ ddZ ddZddZee            d/de jde jdee j dee jdeedee jdee jd ee jd!eed"eed#eed$eee j ffd%d&Z      d0fd'd( Zd)e j fd*d+Ze j   d1d e j dee jd,eeffd-d. ZZS)2JanusForConditionalGenerationz(model.language_model.embed_tokens.weightzlm_head.weightTr-csBt|||_t||_tj|jj|jj dd|_ | dS)NFr) rZr[r-rLr%rr/rSr\ vocab_sizelm_headrrerfr9r:r[s   z&JanusForConditionalGeneration.__init__cCs |jjSr)r%rTrrr9r9r:rs z2JanusForConditionalGeneration.get_input_embeddingscCs|jj|dSr)r%rTrVrWr9r9r:rVsz2JanusForConditionalGeneration.set_input_embeddingsinputsrkcCs|j|}|j|}|Sr)r%rPrQ)r7rbrr9r9r:'prepare_embeddings_for_image_generations  zEJanusForConditionalGeneration.prepare_embeddings_for_image_generationcCrrrarr9r9r:get_output_embeddingsrz3JanusForConditionalGeneration.get_output_embeddingscC ||_dSrrd)r7Znew_embeddingsr9r9r:set_output_embeddingsrUz3JanusForConditionalGeneration.set_output_embeddingscCrfrr%)r7rDr9r9r: set_decoderrUz)JanusForConditionalGeneration.set_decodercCrrrhrr9r9r: get_decoderrz)JanusForConditionalGeneration.get_decoderNrrZrxrrWr&r[rlabelsr\rrr]c  Ks| dur| n|jj} | dur| n|jj} |jd||||||| | | |d | }|j}t| tr5t| dn| }||dd|ddf}d}|durV|j |||jj j d}t |||j |j|j|jd}|S)a labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. N) rZrxrrWr&rr\rrr[)rPrkr`)rOrPr&rKrLrMr9)r-rrr%rJr.rsliceraZ loss_functionrSr`rNr&rKrLrM)r7rZrxrrWr&r[rrkr\rrr]rrrKZ slice_indicesrPrOrr9r9r:r~s@  z%JanusForConditionalGeneration.forwardc  s8tj|f|||||d|} |ddkr|| d<| S)N)r&rrr[r]rrx)rZprepare_inputs_for_generation) r7rZrxr&rrr[r]r model_inputsrfr9r:rms  z;JanusForConditionalGeneration.prepare_inputs_for_generationrcCs"|jj|}|dddd}|S)a, Decodes generated image tokens from language model to continuous pixel values with VQGAN module via upsampling. Args: image_tokens (`torch.LongTensor` of shape `(batch_size, num_of_tokens)`): The tensors corresponding to the input images. rrr r )r%rOrHrr)r7rZ decoded_imager9r9r:decode_image_tokenssz1JanusForConditionalGeneration.decode_image_tokenslogits_processorc% sz|d|j}t|}|dd}|dkr$tjd|||dd|S|jdi|}|tj tj fvr:t d| | ||durK|nt}d|d<|jdur_tdd |_|j|d <|||j|\}} }|j|j} } t|jd krt d |jd |du} |j|| |jd|jr|jdkr|t|jd|_|j||jd|d|| d}|jd|||jd|\}}|jjj j!} |j\}}|"d d}|dd}|"d d}||d<||dddf|jk||dddf|j#dk@}||dddf$||j%|&|}|'|| |}|(dddur=|j)|j*p,d|d t+|j,| || |d|d<t-j.|| f| | d}|j/}|j0}|j1}|j2}|j3}|r^|r^dnd}|rh|rhdnd}|rr|rrdnd}|r||r|dnd}t4| D]}|j5d||d|}|d6|j|d<|d6|j|d<|jj7di|||d}|8||}|j9dddddf:} |j;| }!|||!}"|j<rt-j=|"dd}#t-j>|#dd?d}$nt-j@|"dd}$|$|dd|f<t-A|$|$g}$|$Bd}$|C|$}q|r-|r||!f7}|r|| Df7}|r%||jE7}|r-||jF7}|r;tG||!||||jHdS|S) Ngeneration_configgeneration_modetext)rbrrqguidance_scalezGot incompatible mode for Image Generation, should be one of greedy or sampling. Ensure that beam search is de-activated by setting `num_beams=1` and `num_beam_groups=1`.Tr\zU`guidance_scale` is required for CFG but not provided. Setting to default value of 5.rtrz;Expected input ids of shape (batch_size, seq_len), but got z3Passing `inputs embeds` is not supported currently.)r^r )rqZinput_ids_seq_lengthZencoder_input_idsZprefix_allowed_tokens_fnrpr^)rZrZ expand_sizerZ boi_token_idr&Zstatic)cache_implementationrZ max_cache_lenr^ model_kwargs)ryr^r9)rrZr[)rrrXr)Z num_samples) sequencesscoresrPrLrKr&)IpoprqcopydeepcopyrZgenerateupdateZget_generation_moderZSAMPLEZ GREEDY_SEARCHrvalidateZ_validate_model_kwargsrrtrwarningZ_prepare_model_inputsZ bos_token_idryr^r$rnZ_prepare_special_tokensr/r Z_get_logits_processorZ_expand_inputs_for_generationZnum_return_sequencesr%rMr-num_image_tokensrepeatZgeneration_kwargsZ masked_fill_Z pad_token_idrZ_get_initial_cache_positionrZ _get_cachervmax max_lengthrFZzerosrr output_scores output_logitsreturn_dict_in_generaterrmrzrTZ#_update_model_kwargs_for_generationrJclonerRZ do_samplerZ multinomialZsqueezeZargmaxcatrprcfloatrLrKrr&)%r7rbrrprrqrrrwrZZmodel_input_nameryr^Zkwargs_has_attention_maskrrrZ input_tokensmaskrZgenerated_tokensrrrrrZ raw_scoresZ raw_logitsZdecoder_hidden_statesZdecoder_attentionsirnrrryZnext_token_scoresZprobsZ next_tokenrfr9r:r}s                           z&JanusForConditionalGeneration.generate) NNNNNNNNNNNr)NNNNNNr) r<r=r>Z_tied_weights_keysrAr!r[rrVrFrrcrergrirjrrrrGrr rrrr~rmroZno_gradrr}rr9r9rfr:r_s      > r_)r$r_rLr?r)r))Tr{ dataclassesrtypingrrrrrrFrZ activationsr Z cache_utilsr Z generationr r rrZgeneration.utilsrZmodeling_flash_attention_utilsrZmodeling_layersrZmodeling_outputsrrrZmodeling_utilsrrZprocessing_utilsrutilsrrrrrautorZconfiguration_janusr!r"r#Ztorch.nn.functionalrsrZ get_loggerr<rr$rBrIrNr0rQrrrrrrrrrrrrrrrrrr#r;r?rIrJrLr___all__r9r9r9r:s           ,)K U1P>?,#  MD>n N