


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

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


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







class 							lowerCamelCase__						(				__SCREAMING_SNAKE_CASE   ):






			"""simple docstring"""






			def __init__(				self   ,    snake_case   ,    snake_case=768  ):

						'''simple docstring'''


						super().__init__(lowerCAmelCase__  )
						UpperCamelCase__       =							proj_size

						UpperCamelCase__       =							CLIPVisionModel(lowerCAmelCase__  )
						UpperCamelCase__       =							PaintByExampleMapper(lowerCAmelCase__  )
						UpperCamelCase__       =							nn.LayerNorm(config.hidden_size  )
						UpperCamelCase__       =							nn.Linear(config.hidden_size   ,    self.proj_size  )

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







			def 					snake_case__					(				self   ,    snake_case   ,    snake_case=False  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model(pixel_values=lowerCAmelCase__  )
						UpperCamelCase__       =							clip_output.pooler_output
						UpperCamelCase__       =							self.mapper(latent_states[:, None]  )
						UpperCamelCase__       =							self.final_layer_norm(lowerCAmelCase__  )
						UpperCamelCase__       =							self.proj_out(lowerCAmelCase__  )
						if return_uncond_vector:
									return latent_states, self.uncond_vector

						return latent_states







class 							lowerCamelCase__						(				nn.Module   ):






			"""simple docstring"""






			def __init__(				self   ,    snake_case  ):

						'''simple docstring'''


						super().__init__()
						UpperCamelCase__       =							(config.num_hidden_layers + 1) // 5
						UpperCamelCase__       =							config.hidden_size
						UpperCamelCase__       =							1
						UpperCamelCase__       =							nn.ModuleList(
						    [
						        BasicTransformerBlock(lowerCAmelCase__   ,    lowerCAmelCase__   ,    lowerCAmelCase__   ,    activation_fn="gelu"   ,    attention_bias=lowerCAmelCase__  )
						        for _ in range(lowerCAmelCase__  )
						    ]  )







			def 					snake_case__					(				self   ,    snake_case  ):

						'''simple docstring'''


						for block in self.blocks:
									UpperCamelCase__       =							block(lowerCAmelCase__  )

						return hidden_states






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowerCAmelCase__			:      Union[str, Any]								={
    'configuration_blip_2': [
        'BLIP_2_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'Blip2Config',
        'Blip2QFormerConfig',
        'Blip2VisionConfig',
    ],
    'processing_blip_2': ['Blip2Processor'],
}

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				lowerCAmelCase__			:      Dict								=[
				    'BLIP_2_PRETRAINED_MODEL_ARCHIVE_LIST',
				    'Blip2Model',
				    'Blip2QFormerModel',
				    'Blip2PreTrainedModel',
				    'Blip2ForConditionalGeneration',
				    'Blip2VisionModel',
				]

if TYPE_CHECKING:
				from .configuration_blip_a import (
				    BLIP_2_PRETRAINED_CONFIG_ARCHIVE_MAP,
				    BlipaConfig,
				    BlipaQFormerConfig,
				    BlipaVisionConfig,
				)
				from .processing_blip_a import BlipaProcessor

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_blip_a import (
								    BLIP_2_PRETRAINED_MODEL_ARCHIVE_LIST,
								    BlipaForConditionalGeneration,
								    BlipaModel,
								    BlipaPreTrainedModel,
								    BlipaQFormerModel,
								    BlipaVisionModel,
								)

else:
				import sys

				lowerCAmelCase__			:      Union[str, Any]								=_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





