





from dataclasses import dataclass
from typing import Tuple

import numpy as np
import torch





@dataclass
class 			lowerCAmelCase__:






       '''simple docstring'''


       __snake_case 			=    42  # [batch_size x 3]
       __snake_case 			=    42  # [batch_size x 3]
       __snake_case 			=    42  # [batch_size x 3]
       __snake_case 			=    42  # [batch_size x 3]
       __snake_case 			=    42
       __snake_case 			=    42
       __snake_case 			=    42
       __snake_case 			=    42
       __snake_case 			=    42




       def     UpperCamelCase_			(      self							)			->     Dict:
            assert self.x.shape[0] == self.y.shape[0] == self.z.shape[0] == self.origin.shape[0]
            assert self.x.shape[1] == self.y.shape[1] == self.z.shape[1] == self.origin.shape[1] == 3
            assert len(self.x.shape							) == len(self.y.shape							) == len(self.z.shape							) == len(self.origin.shape							) == 2




       def     UpperCamelCase_			(      self							)			->     int:
            return torch.from_numpy(np.array([self.width, self.height]    ,						dtype=np.floataa							)							)




       def     UpperCamelCase_			(      self							)			->     Optional[int]:
            return torch.from_numpy(np.array([self.x_fov, self.y_fov]    ,						dtype=np.floataa							)							)




       def     UpperCamelCase_			(      self							)			->     torch.Tensor:
            _SCREAMING_SNAKE_CASE       :	int						    =							torch.arange(self.height * self.width							)
            _SCREAMING_SNAKE_CASE       :	str						    =							torch.stack(
                [
                    pixel_indices % self.width,
                    torch.div(__lowerCamelCase    ,						self.width    ,						rounding_mode="trunc"							),
                ]    ,						axis=1    ,						)
            return coords




       @property
       def     UpperCamelCase_			(      self							)			->     List[Any]:
            _SCREAMING_SNAKE_CASE    ,			*_SCREAMING_SNAKE_CASE       :	int						    =							self.shape
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							int(np.prod(__lowerCamelCase							)							)

            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							self.get_image_coords()
            _SCREAMING_SNAKE_CASE       :	Dict						    =							torch.broadcast_to(coords.unsqueeze(0							)    ,						[batch_size * inner_batch_size, *coords.shape]							)
            _SCREAMING_SNAKE_CASE       :	str						    =							self.get_camera_rays(__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	Tuple						    =							rays.view(__lowerCamelCase    ,						inner_batch_size * self.height * self.width    ,						2    ,						3							)

            return rays




       def     UpperCamelCase_			(      self    ,						__lowerCamelCase							)			->     torch.Tensor:
            _SCREAMING_SNAKE_CASE    ,			*_SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	Any						    =							coords.shape
            assert n_coords == 2
            assert batch_size == self.origin.shape[0]

            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							coords.view(__lowerCamelCase    ,						-1    ,						2							)

            _SCREAMING_SNAKE_CASE       :	List[str]						    =							self.resolution()
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							self.fov()

            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							(flat.float() / (res - 1)) * 2 - 1
            _SCREAMING_SNAKE_CASE       :	int						    =							fracs * torch.tan(fov / 2							)

            _SCREAMING_SNAKE_CASE       :	int						    =							fracs.view(__lowerCamelCase    ,						-1    ,						2							)
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							(
                self.z.view(__lowerCamelCase    ,						1    ,						3							)
                + self.x.view(__lowerCamelCase    ,						1    ,						3							) * fracs[:, :, :1]
                + self.y.view(__lowerCamelCase    ,						1    ,						3							) * fracs[:, :, 1:]
            )
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							directions / directions.norm(dim=-1    ,						keepdim=__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							torch.stack(
                [
                    torch.broadcast_to(self.origin.view(__lowerCamelCase    ,						1    ,						3							)    ,						[batch_size, directions.shape[1], 3]							),
                    directions,
                ]    ,						dim=2    ,						)
            return rays.view(__lowerCamelCase    ,						*__lowerCamelCase    ,						2    ,						3							)







       def     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase							)			->     "DifferentiableProjectiveCamera":
            assert width * self.height == height * self.width, "The aspect ratio should not change."
            return DifferentiableProjectiveCamera(
                origin=self.origin    ,						x=self.x    ,						y=self.y    ,						z=self.z    ,						width=__lowerCamelCase    ,						height=__lowerCamelCase    ,						x_fov=self.x_fov    ,						y_fov=self.y_fov    ,						)






def     lowerCamelCase__			(__lowerCamelCase	):
     _SCREAMING_SNAKE_CASE       :	str						    =							[]
     _SCREAMING_SNAKE_CASE       :	List[str]						    =							[]
     _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							[]
     _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							[]
     for theta in np.linspace(0,       2 * np.pi,       num=20	):
          _SCREAMING_SNAKE_CASE       :	List[str]						    =							np.array([np.sin(__lowerCamelCase	), np.cos(__lowerCamelCase	), -0.5]	)
          z /= np.sqrt(np.sum(z**2	)	)
          _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							-z * 4
          _SCREAMING_SNAKE_CASE       :	int						    =							np.array([np.cos(__lowerCamelCase	), -np.sin(__lowerCamelCase	), 0.0]	)
          _SCREAMING_SNAKE_CASE       :	List[Any]						    =							np.cross(__lowerCamelCase,       __lowerCamelCase	)
          origins.append(__lowerCamelCase	)
          xs.append(__lowerCamelCase	)
          ys.append(__lowerCamelCase	)
          zs.append(__lowerCamelCase	)
     return DifferentiableProjectiveCamera(
         origin=torch.from_numpy(np.stack(__lowerCamelCase,       axis=0	)	).float(),       x=torch.from_numpy(np.stack(__lowerCamelCase,       axis=0	)	).float(),       y=torch.from_numpy(np.stack(__lowerCamelCase,       axis=0	)	).float(),       z=torch.from_numpy(np.stack(__lowerCamelCase,       axis=0	)	).float(),       width=__lowerCamelCase,       height=__lowerCamelCase,       x_fov=0.7,       y_fov=0.7,       shape=(1, len(__lowerCamelCase	)),       )






import copy
import os
from typing import TYPE_CHECKING, List, Union


if TYPE_CHECKING:
       pass

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


UpperCamelCase__      =logging.get_logger(__name__)

UpperCamelCase__      ={
    'kakaobrain/align-base': 'https://huggingface.co/kakaobrain/align-base/resolve/main/config.json',
}





class 			lowerCAmelCase__( __lowercase    ):






       '''simple docstring'''


       __snake_case 			=    'align_text_model'




       def __init__(      self    ,						__lowerCamelCase=3_0_5_2_2    ,						__lowerCamelCase=7_6_8    ,						__lowerCamelCase=1_2    ,						__lowerCamelCase=1_2    ,						__lowerCamelCase=3_0_7_2    ,						__lowerCamelCase="gelu"    ,						__lowerCamelCase=0.1    ,						__lowerCamelCase=0.1    ,						__lowerCamelCase=5_1_2    ,						__lowerCamelCase=2    ,						__lowerCamelCase=0.02    ,						__lowerCamelCase=1E-12    ,						__lowerCamelCase=0    ,						__lowerCamelCase="absolute"    ,						__lowerCamelCase=True    ,						**__lowerCamelCase    ,						)			->     List[Any]:
            super().__init__(**__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	Tuple						    =							vocab_size
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							hidden_size
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							num_hidden_layers
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							num_attention_heads
            _SCREAMING_SNAKE_CASE       :	int						    =							hidden_act
            _SCREAMING_SNAKE_CASE       :	Any						    =							intermediate_size
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							hidden_dropout_prob
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							attention_probs_dropout_prob
            _SCREAMING_SNAKE_CASE       :	Dict						    =							max_position_embeddings
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							type_vocab_size
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							initializer_range
            _SCREAMING_SNAKE_CASE       :	Dict						    =							layer_norm_eps
            _SCREAMING_SNAKE_CASE       :	str						    =							position_embedding_type
            _SCREAMING_SNAKE_CASE       :	Dict						    =							use_cache
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							pad_token_id







       @classmethod
       def     UpperCamelCase_			(      cls    ,						__lowerCamelCase    ,						**__lowerCamelCase							)			->     "PretrainedConfig":
            cls._set_token_in_kwargs(__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	List[Any]						    =							cls.get_config_dict(__lowerCamelCase    ,						**__lowerCamelCase							)

            # get the text config dict if we are loading from AlignConfig
            if config_dict.get("model_type"							) == "align":
                 _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							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__( __lowercase    ):






       '''simple docstring'''


       __snake_case 			=    'align_vision_model'




       def __init__(      self    ,						__lowerCamelCase = 3    ,						__lowerCamelCase = 6_0_0    ,						__lowerCamelCase = 2.0    ,						__lowerCamelCase = 3.1    ,						__lowerCamelCase = 8    ,						__lowerCamelCase = [3, 3, 5, 3, 5, 5, 3]    ,						__lowerCamelCase = [3_2, 1_6, 2_4, 4_0, 8_0, 1_1_2, 1_9_2]    ,						__lowerCamelCase = [1_6, 2_4, 4_0, 8_0, 1_1_2, 1_9_2, 3_2_0]    ,						__lowerCamelCase = []    ,						__lowerCamelCase = [1, 2, 2, 2, 1, 2, 1]    ,						__lowerCamelCase = [1, 2, 2, 3, 3, 4, 1]    ,						__lowerCamelCase = [1, 6, 6, 6, 6, 6, 6]    ,						__lowerCamelCase = 0.25    ,						__lowerCamelCase = "swish"    ,						__lowerCamelCase = 2_5_6_0    ,						__lowerCamelCase = "mean"    ,						__lowerCamelCase = 0.02    ,						__lowerCamelCase = 0.001    ,						__lowerCamelCase = 0.99    ,						__lowerCamelCase = 0.2    ,						**__lowerCamelCase    ,						)			->     Dict:
            super().__init__(**__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	int						    =							num_channels
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							image_size
            _SCREAMING_SNAKE_CASE       :	Dict						    =							width_coefficient
            _SCREAMING_SNAKE_CASE       :	str						    =							depth_coefficient
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							depth_divisor
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							kernel_sizes
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							in_channels
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							out_channels
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							depthwise_padding
            _SCREAMING_SNAKE_CASE       :	str						    =							strides
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							num_block_repeats
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							expand_ratios
            _SCREAMING_SNAKE_CASE       :	int						    =							squeeze_expansion_ratio
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							hidden_act
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							hidden_dim
            _SCREAMING_SNAKE_CASE       :	Dict						    =							pooling_type
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							initializer_range
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							batch_norm_eps
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							batch_norm_momentum
            _SCREAMING_SNAKE_CASE       :	int						    =							drop_connect_rate
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							sum(__lowerCamelCase							) * 4







       @classmethod
       def     UpperCamelCase_			(      cls    ,						__lowerCamelCase    ,						**__lowerCamelCase							)			->     "PretrainedConfig":
            cls._set_token_in_kwargs(__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	int						    =							cls.get_config_dict(__lowerCamelCase    ,						**__lowerCamelCase							)

            # get the vision config dict if we are loading from AlignConfig
            if config_dict.get("model_type"							) == "align":
                 _SCREAMING_SNAKE_CASE       :	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__( __lowercase    ):






       '''simple docstring'''


       __snake_case 			=    'align'
       __snake_case 			=    True




       def __init__(      self    ,						__lowerCamelCase=None    ,						__lowerCamelCase=None    ,						__lowerCamelCase=6_4_0    ,						__lowerCamelCase=1.0    ,						__lowerCamelCase=0.02    ,						**__lowerCamelCase    ,						)			->     List[Any]:
            super().__init__(**__lowerCamelCase							)

            if text_config is None:
                 _SCREAMING_SNAKE_CASE       :	List[Any]						    =							{}
                 logger.info("text_config is None. Initializing the AlignTextConfig with default values."							)

            if vision_config is None:
                 _SCREAMING_SNAKE_CASE       :	List[str]						    =							{}
                 logger.info("vision_config is None. Initializing the AlignVisionConfig with default values."							)

            _SCREAMING_SNAKE_CASE       :	Dict						    =							AlignTextConfig(**__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	Any						    =							AlignVisionConfig(**__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	int						    =							projection_dim
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							temperature_init_value
            _SCREAMING_SNAKE_CASE       :	Any						    =							initializer_range




       @classmethod
       def     UpperCamelCase_			(      cls    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						**__lowerCamelCase							)			->     List[str]:
            return cls(text_config=text_config.to_dict()    ,						vision_config=vision_config.to_dict()    ,						**__lowerCamelCase							)







       def     UpperCamelCase_			(      self							)			->     Union[str, Any]:
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							copy.deepcopy(self.__dict__							)
            _SCREAMING_SNAKE_CASE       :	Any						    =							self.text_config.to_dict()
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							self.vision_config.to_dict()
            _SCREAMING_SNAKE_CASE       :	Dict						    =							self.__class__.model_type
            return output