'''simple docstring'''

import torch
from torch import nn
from transformers import CLIPPreTrainedModel, CLIPVisionModel

from ...models.attention import BasicTransformerBlock
from ...utils import logging


UpperCamelCase__  :						Any					=      logging.get_logger(__name__)  # pylint: disable=invalid-name


class 					_lowercase						(   lowerCAmelCase ):
      '''simple docstring'''


      def __init__(					self       ,lowerCamelCase_       ,lowerCamelCase_=768				)      ->      int:



         '''simple docstring'''



         super().__init__(lowerCamelCase_				)
         UpperCAmelCase__				:					int 					=	proj_size

         UpperCAmelCase__				:					Dict 					=	CLIPVisionModel(lowerCamelCase_				)
         UpperCAmelCase__				:					List[Any] 					=	PaintByExampleMapper(lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[int] 					=	nn.LayerNorm(config.hidden_size				)
         UpperCAmelCase__				:					str 					=	nn.Linear(config.hidden_size       ,self.proj_size				)

         # uncondition for scaling
         UpperCAmelCase__				:					int 					=	nn.Parameter(torch.randn((1, 1, self.proj_size)				)				)




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_=False				)      ->      Any:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[int] 					=	self.model(pixel_values=lowerCamelCase_				)
         UpperCAmelCase__				:					str 					=	clip_output.pooler_output
         UpperCAmelCase__				:					Dict 					=	self.mapper(latent_states[:, None]				)
         UpperCAmelCase__				:					Optional[Any] 					=	self.final_layer_norm(lowerCamelCase_				)
         UpperCAmelCase__				:					List[str] 					=	self.proj_out(lowerCamelCase_				)
         if return_uncond_vector:
            return latent_states, self.uncond_vector

         return latent_states




class 					_lowercase						(   nn.Module ):
      '''simple docstring'''


      def __init__(					self       ,lowerCamelCase_				)      ->      List[Any]:



         '''simple docstring'''



         super().__init__()
         UpperCAmelCase__				:					int 					=	(config.num_hidden_layers + 1) // 5
         UpperCAmelCase__				:					Any 					=	config.hidden_size
         UpperCAmelCase__				:					Tuple 					=	1
         UpperCAmelCase__				:					Dict 					=	nn.ModuleList(
             [
                 BasicTransformerBlock(lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,activation_fn='''gelu'''       ,attention_bias=lowerCamelCase_				)
                 for _ in range(lowerCamelCase_				)
             ]				)




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_				)      ->      Optional[Any]:



         '''simple docstring'''



         for block in self.blocks:
            UpperCAmelCase__				:					Any 					=	block(lowerCamelCase_				)

         return hidden_states





'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCamelCase__  :						str					=      {
    'configuration_git': ['GIT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'GitConfig', 'GitVisionConfig'],
    'processing_git': ['GitProcessor'],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__  :						str					=      [
      'GIT_PRETRAINED_MODEL_ARCHIVE_LIST',
      'GitForCausalLM',
      'GitModel',
      'GitPreTrainedModel',
      'GitVisionModel',
  ]

if TYPE_CHECKING:
  from .configuration_git import GIT_PRETRAINED_CONFIG_ARCHIVE_MAP, GitConfig, GitVisionConfig
  from .processing_git import GitProcessor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_git import (
        GIT_PRETRAINED_MODEL_ARCHIVE_LIST,
        GitForCausalLM,
        GitModel,
        GitPreTrainedModel,
        GitVisionModel,
    )

else:
  import sys

  UpperCamelCase__  :						Tuple					=      _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




