



"""simple docstring"""




import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__snake_case    :		Any  =logging.get_logger(__name__)

__snake_case    :		Optional[Any]  ={
    'Salesforce/blip-vqa-base': 'https://huggingface.co/Salesforce/blip-vqa-base/resolve/main/config.json',
    'Salesforce/blip-vqa-capfit-large': (
        'https://huggingface.co/Salesforce/blip-vqa-base-capfit/resolve/main/config.json'
    ),
    'Salesforce/blip-image-captioning-base': (
        'https://huggingface.co/Salesforce/blip-image-captioning-base/resolve/main/config.json'
    ),
    'Salesforce/blip-image-captioning-large': (
        'https://huggingface.co/Salesforce/blip-image-captioning-large/resolve/main/config.json'
    ),
    'Salesforce/blip-itm-base-coco': 'https://huggingface.co/Salesforce/blip-itm-base-coco/resolve/main/config.json',
    'Salesforce/blip-itm-large-coco': 'https://huggingface.co/Salesforce/blip-itm-large-coco/resolve/main/config.json',
    'Salesforce/blip-itm-base-flikr': 'https://huggingface.co/Salesforce/blip-itm-base-flikr/resolve/main/config.json',
    'Salesforce/blip-itm-large-flikr': (
        'https://huggingface.co/Salesforce/blip-itm-large-flikr/resolve/main/config.json'
    ),
}



class 	lowerCamelCase__      (							lowerCamelCase__):




		'''simple docstring'''



		snake_case_     ="""blip_text_model"""



		def __init__(self		,__lowerCamelCase=3_05_24		,__lowerCamelCase=7_68		,__lowerCamelCase=7_68		,__lowerCamelCase=30_72		,__lowerCamelCase=7_68		,__lowerCamelCase=12		,__lowerCamelCase=8		,__lowerCamelCase=5_12		,__lowerCamelCase="gelu"		,__lowerCamelCase=1e-12		,__lowerCamelCase=0.0		,__lowerCamelCase=0.0		,__lowerCamelCase=0.02		,__lowerCamelCase=3_05_22		,__lowerCamelCase=2		,__lowerCamelCase=0		,__lowerCamelCase=1_02		,__lowerCamelCase=True		,__lowerCamelCase=True		,**__lowerCamelCase		,)					->      Optional[Any]:

						"""simple docstring"""






						super().__init__(
						    pad_token_id=__lowerCamelCase		,bos_token_id=__lowerCamelCase		,eos_token_id=__lowerCamelCase		,sep_token_id=__lowerCamelCase		,**__lowerCamelCase		,)

						lowerCAmelCase__					:    Any												=					vocab_size
						lowerCAmelCase__					:    List[Any]												=					hidden_size
						lowerCAmelCase__					:    Optional[int]												=					encoder_hidden_size
						lowerCAmelCase__					:    Tuple												=					intermediate_size
						lowerCAmelCase__					:    Any												=					projection_dim
						lowerCAmelCase__					:    Optional[Any]												=					hidden_dropout_prob
						lowerCAmelCase__					:    Tuple												=					num_hidden_layers
						lowerCAmelCase__					:    Any												=					num_attention_heads
						lowerCAmelCase__					:    Optional[int]												=					max_position_embeddings
						lowerCAmelCase__					:    Optional[Any]												=					layer_norm_eps
						lowerCAmelCase__					:    Any												=					hidden_act
						lowerCAmelCase__					:    Optional[int]												=					initializer_range
						lowerCAmelCase__					:    Dict												=					attention_probs_dropout_prob
						lowerCAmelCase__					:    Dict												=					is_decoder
						lowerCAmelCase__					:    List[Any]												=					use_cache



		@classmethod
		def 		lowerCAmelCase__				(cls		,__lowerCamelCase		,**__lowerCamelCase			)					->      "PretrainedConfig":

						"""simple docstring"""






						cls._set_token_in_kwargs(__lowerCamelCase			)

						lowerCAmelCase__					:    Dict												=					cls.get_config_dict(__lowerCamelCase		,**__lowerCamelCase			)

						# get the text config dict if we are loading from BlipConfig
						if config_dict.get('''model_type'''			) == "blip":
										lowerCAmelCase__					:    Tuple												=					config_dict['''text_config''']

						if "model_type" in config_dict and hasattr(cls		,'''model_type'''			) and config_dict["model_type"] != cls.model_type:
										logger.warning(
										    f"""You are using a model of type {config_dict['model_type']} to instantiate a model of type """
										    f"""{cls.model_type}. This is not supported for all configurations of models and can yield errors."""			)

						return cls.from_dict(__lowerCamelCase		,**__lowerCamelCase			)



class 	lowerCamelCase__      (							lowerCamelCase__):




		'''simple docstring'''



		snake_case_     ="""blip_vision_model"""



		def __init__(self		,__lowerCamelCase=7_68		,__lowerCamelCase=30_72		,__lowerCamelCase=5_12		,__lowerCamelCase=12		,__lowerCamelCase=12		,__lowerCamelCase=3_84		,__lowerCamelCase=16		,__lowerCamelCase="gelu"		,__lowerCamelCase=1e-5		,__lowerCamelCase=0.0		,__lowerCamelCase=1e-10		,**__lowerCamelCase		,)					->      Optional[Any]:

						"""simple docstring"""






						super().__init__(**__lowerCamelCase			)

						lowerCAmelCase__					:    str												=					hidden_size
						lowerCAmelCase__					:    List[Any]												=					intermediate_size
						lowerCAmelCase__					:    Optional[int]												=					projection_dim
						lowerCAmelCase__					:    Union[str, Any]												=					num_hidden_layers
						lowerCAmelCase__					:    Union[str, Any]												=					num_attention_heads
						lowerCAmelCase__					:    Optional[int]												=					patch_size
						lowerCAmelCase__					:    Tuple												=					image_size
						lowerCAmelCase__					:    Union[str, Any]												=					initializer_range
						lowerCAmelCase__					:    List[str]												=					attention_dropout
						lowerCAmelCase__					:    int												=					layer_norm_eps
						lowerCAmelCase__					:    Dict												=					hidden_act



		@classmethod
		def 		lowerCAmelCase__				(cls		,__lowerCamelCase		,**__lowerCamelCase			)					->      "PretrainedConfig":

						"""simple docstring"""






						cls._set_token_in_kwargs(__lowerCamelCase			)

						lowerCAmelCase__					:    Any												=					cls.get_config_dict(__lowerCamelCase		,**__lowerCamelCase			)

						# get the vision config dict if we are loading from BlipConfig
						if config_dict.get('''model_type'''			) == "blip":
										lowerCAmelCase__					:    int												=					config_dict['''vision_config''']

						if "model_type" in config_dict and hasattr(cls		,'''model_type'''			) and config_dict["model_type"] != cls.model_type:
										logger.warning(
										    f"""You are using a model of type {config_dict['model_type']} to instantiate a model of type """
										    f"""{cls.model_type}. This is not supported for all configurations of models and can yield errors."""			)

						return cls.from_dict(__lowerCamelCase		,**__lowerCamelCase			)



class 	lowerCamelCase__      (							lowerCamelCase__):




		'''simple docstring'''



		snake_case_     ="""blip"""
		snake_case_     =True



		def __init__(self		,__lowerCamelCase=None		,__lowerCamelCase=None		,__lowerCamelCase=5_12		,__lowerCamelCase=2.6592		,__lowerCamelCase=2_56		,**__lowerCamelCase		,)					->      List[str]:

						"""simple docstring"""






						super().__init__(**__lowerCamelCase			)

						if text_config is None:
										lowerCAmelCase__					:    List[Any]												=					{}
										logger.info('''`text_config` is `None`. Initializing the `BlipTextConfig` with default values.'''			)

						if vision_config is None:
										lowerCAmelCase__					:    List[Any]												=					{}
										logger.info('''`vision_config` is `None`. Initializing the `BlipVisionConfig` with default values.'''			)

						lowerCAmelCase__					:    List[Any]												=					BlipTextConfig(**__lowerCamelCase			)
						lowerCAmelCase__					:    int												=					BlipVisionConfig(**__lowerCamelCase			)

						lowerCAmelCase__					:    Optional[int]												=					self.vision_config.hidden_size

						lowerCAmelCase__					:    Any												=					projection_dim
						lowerCAmelCase__					:    Dict												=					logit_scale_init_value
						lowerCAmelCase__					:    int												=					1.0
						lowerCAmelCase__					:    Tuple												=					0.02
						lowerCAmelCase__					:    Union[str, Any]												=					image_text_hidden_size



		@classmethod
		def 		lowerCAmelCase__				(cls		,__lowerCamelCase		,__lowerCamelCase		,**__lowerCamelCase			)					->      List[Any]:

						"""simple docstring"""






						return cls(text_config=text_config.to_dict()		,vision_config=vision_config.to_dict()		,**__lowerCamelCase			)



		def 		lowerCAmelCase__				(self			)					->      Dict:

						"""simple docstring"""






						lowerCAmelCase__					:    Tuple												=					copy.deepcopy(self.__dict__			)
						lowerCAmelCase__					:    int												=					self.text_config.to_dict()
						lowerCAmelCase__					:    Tuple												=					self.vision_config.to_dict()
						lowerCAmelCase__					:    List[Any]												=					self.__class__.model_type
						return output






import flax.linen as nn
import jax
import jax.numpy as jnp



class 	lowerCamelCase__      (							nn.Module):




		'''simple docstring'''



		snake_case_     =42
		snake_case_     =jnp.floataa



		def 		lowerCAmelCase__				(self			)					->      List[str]:

						"""simple docstring"""






						lowerCAmelCase__					:    Dict												=					nn.Conv(
						    self.out_channels		,kernel_size=(3, 3)		,strides=(1, 1)		,padding=((1, 1), (1, 1))		,dtype=self.dtype		,)



		def __call__(self		,__lowerCamelCase			)					->      Union[str, Any]:

						"""simple docstring"""






						lowerCAmelCase__						,		lowerCAmelCase__						,		lowerCAmelCase__						,		lowerCAmelCase__					:    Dict												=					hidden_states.shape
						lowerCAmelCase__					:    Dict												=					jax.image.resize(
						    __lowerCamelCase		,shape=(batch, height * 2, width * 2, channels)		,method='''nearest'''		,)
						lowerCAmelCase__					:    Dict												=					self.conv(__lowerCamelCase			)
						return hidden_states



class 	lowerCamelCase__      (							nn.Module):




		'''simple docstring'''



		snake_case_     =42
		snake_case_     =jnp.floataa



		def 		lowerCAmelCase__				(self			)					->      int:

						"""simple docstring"""






						lowerCAmelCase__					:    List[Any]												=					nn.Conv(
						    self.out_channels		,kernel_size=(3, 3)		,strides=(2, 2)		,padding=((1, 1), (1, 1))		,dtype=self.dtype		,)



		def __call__(self		,__lowerCamelCase			)					->      Any:

						"""simple docstring"""






						lowerCAmelCase__					:    Optional[int]												=					self.conv(__lowerCamelCase			)
						return hidden_states



class 	lowerCamelCase__      (							nn.Module):




		'''simple docstring'''



		snake_case_     =42
		snake_case_     =None
		snake_case_     =0.0
		snake_case_     =None
		snake_case_     =jnp.floataa



		def 		lowerCAmelCase__				(self			)					->      int:

						"""simple docstring"""






						lowerCAmelCase__					:    List[str]												=					self.in_channels if self.out_channels is None else self.out_channels

						lowerCAmelCase__					:    Union[str, Any]												=					nn.GroupNorm(num_groups=32		,epsilon=1e-5			)
						lowerCAmelCase__					:    Union[str, Any]												=					nn.Conv(
						    __lowerCamelCase		,kernel_size=(3, 3)		,strides=(1, 1)		,padding=((1, 1), (1, 1))		,dtype=self.dtype		,)

						lowerCAmelCase__					:    int												=					nn.Dense(__lowerCamelCase		,dtype=self.dtype			)

						lowerCAmelCase__					:    List[Any]												=					nn.GroupNorm(num_groups=32		,epsilon=1e-5			)
						lowerCAmelCase__					:    Union[str, Any]												=					nn.Dropout(self.dropout_prob			)
						lowerCAmelCase__					:    Optional[int]												=					nn.Conv(
						    __lowerCamelCase		,kernel_size=(3, 3)		,strides=(1, 1)		,padding=((1, 1), (1, 1))		,dtype=self.dtype		,)

						lowerCAmelCase__					:    str												=					self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

						lowerCAmelCase__					:    Union[str, Any]												=					None
						if use_nin_shortcut:
										lowerCAmelCase__					:    Optional[Any]												=					nn.Conv(
										    __lowerCamelCase		,kernel_size=(1, 1)		,strides=(1, 1)		,padding='''VALID'''		,dtype=self.dtype		,)



		def __call__(self		,__lowerCamelCase		,__lowerCamelCase		,__lowerCamelCase=True			)					->      Union[str, Any]:

						"""simple docstring"""






						lowerCAmelCase__					:    Tuple												=					hidden_states
						lowerCAmelCase__					:    Dict												=					self.norma(__lowerCamelCase			)
						lowerCAmelCase__					:    Union[str, Any]												=					nn.swish(__lowerCamelCase			)
						lowerCAmelCase__					:    List[Any]												=					self.conva(__lowerCamelCase			)

						lowerCAmelCase__					:    List[Any]												=					self.time_emb_proj(nn.swish(__lowerCamelCase			)			)
						lowerCAmelCase__					:    Dict												=					jnp.expand_dims(jnp.expand_dims(__lowerCamelCase		,1			)		,1			)
						lowerCAmelCase__					:    Optional[int]												=					hidden_states + temb

						lowerCAmelCase__					:    Optional[int]												=					self.norma(__lowerCamelCase			)
						lowerCAmelCase__					:    Optional[int]												=					nn.swish(__lowerCamelCase			)
						lowerCAmelCase__					:    int												=					self.dropout(__lowerCamelCase		,__lowerCamelCase			)
						lowerCAmelCase__					:    List[Any]												=					self.conva(__lowerCamelCase			)

						if self.conv_shortcut is not None:
										lowerCAmelCase__					:    Optional[int]												=					self.conv_shortcut(__lowerCamelCase			)

						return hidden_states + residual

