rhddlZddlmZddlmZmZmZddlZddlmZddl m Z ddl m Z ddl mZmZmZmZdd lmZdd lmZdd lmZmZmZdd lmZmZdd lmZddlm Z m!Z!m"Z"m#Z#m$Z$m%Z%ddl&m'Z'ddl(m)Z)m*Z*m+Z+e#r ddl,mcm-Z.e$j^e0Z1e!GddeZ2ee!dGddeZ3ee!dGddeZ4ee!dGddeZ5Gdd ejlZ7d!ejpd"e9d#ejpfd$Z: dWd%ejld&ejpd'ejpd(ejpd)eejpd*e;d+e;d,ee fd-Z<Gd.d/ejlZ=Gd0d1ejlZ>Gd2d3eZ?Gd4d5ejlZ@e!Gd6d7e2ZAGd8d9ejlZBGd:d;ejlZCGd<d=ejlZDGd>d?ejlZEGd@dAejlZFGdBdCejlZGGdDdEejlZHGdFdGejlZIGdHdIejlZJe!dJGdKdLe2ZKGdMdNejlZLGdOdPejlZMe!dQGdRdSe2ZNGdTdUe2eZOgdVZPy)XN) dataclass)CallableOptionalUnion)nn)ACT2FN)Cache)%ClassifierFreeGuidanceLogitsProcessorGenerationMixinGenerationModeLogitsProcessorList)GenerateDecoderOnlyOutput)GradientCheckpointingLayer)BaseModelOutputBaseModelOutputWithPooling ModelOutput)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)TransformersKwargsauto_docstringcan_return_tupleis_torch_availablelogging torch_int) AutoModel) JanusConfigJanusVisionConfigJanusVQVAEConfigc@eZdZUeed<dZdZddgZddgZdZ dZ dZ dZ y ) JanusPreTrainedModelconfigmodelTLlamaDecoderLayerJanusVisionEncoderLayerpast_key_values causal_maskFN) __name__ __module__ __qualname__r __annotations__base_model_prefixsupports_gradient_checkpointing_no_split_modules_skip_keys_device_placement_supports_flash_attn_supports_sdpa_can_compile_fullgraph!_supports_param_buffer_assignment{/var/www/html/ai-insurance-compliance-backend/venv/lib/python3.12/site-packages/transformers/models/janus/modeling_janus.pyr$r$8sB &*#,.GH#4m"DN!(-%r8r$z9 Base class for Janus VQ-VAE mode model outputs. ) custom_introc\eZdZUdZdZeejed<dZ ejed<y)JanusVQVAEOutputz 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__r=rtorch FloatTensorr.r>r7r8r9r<r<Fs/9=(5#4#45<(,NE%%,r8r<zy Base class for Janus model's outputs that may also contain a past key/values (to speed up sequential decoding). ceZdZUdZdZeejed<dZ ee e ejed<dZ ee ejed<dZ ee ejed<dZ ee ejed<y)JanusBaseModelOutputWithPastal 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 (`Cache`, *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. 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-r?rDrr@rAr.r)tuplerErFrGr7r8r9rCrCXs,6:x 1 129AEOXeE%*;*;$<=>E8B%(9(9":;Br8rCzQ Base class for Janus causal language model (or autoregressive) outputs. ceZdZUdZdZeejed<dZ eejed<dZ ee ejed<dZ ee ejed<dZee ejed<dZee ejed<y) JanusCausalLMOutputWithPasta 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 (`Cache`, *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. 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)rErFrG)r+r,r-r?rKrr@rAr.rLr)listrErHrFrGr7r8r9rJrJ|s$)-D(5$$ %,*.FHU&& '.9=OXd5#4#456=8B%(9(9":;Br8rJceZdZdeffd Zdej dededej fdZd dej d e dej fd Z xZ S) JanusVisionEmbeddingsr%cft|||_|j|_|j |_|j |_tj|j|j|j |j d|_ |j |j zdz|_ |j|_ tj|j|j|_|jdt!j"|jj%ddy)Nvalid) in_channels out_channels kernel_sizestridepaddingr position_ids)rF) persistent)super__init__r% hidden_size embed_dim image_size patch_sizerConv2d num_channelspatch_embedding num_patches num_positions Embeddingposition_embeddingregister_bufferr@arangeexpandselfr% __class__s r9r[zJanusVisionEmbeddings.__init__s  ++ ++ ++!yy++??  !OOt>1D!--"$,,t/A/A4>>"R ^U\\$:L:L-M-T-TU\-]jopr8 embeddingsheightwidthreturnc|jd}|jjjd}tjj s%||k(r ||k(r|j|j S|jjjd}|jd}||jz}||jz} t|dz} |jd| | |}|jdddd}tjj||| fdd }|jddddjdd|}|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 rrrXg?rrbicubicF)sizemode align_corners)shaperfweightr@jit is_tracingrW unsqueezer_rreshapepermuter functional interpolateview) rkrmrnrorcrdpatch_pos_embeddim new_height new_widthsqrt_num_positionss r9interpolate_pos_encodingz.JanusVisionEmbeddings.interpolate_pos_encodingsE!&&q) //66<