



'''simple docstring'''





import copy
import os
from typing import Union

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


__lowerCamelCase         =			logging.get_logger(__name__)

__lowerCamelCase         =			{
    '''BridgeTower/bridgetower-base''': '''https://huggingface.co/BridgeTower/bridgetower-base/blob/main/config.json''',
    '''BridgeTower/bridgetower-base-itm-mlm''': (
        '''https://huggingface.co/BridgeTower/bridgetower-base-itm-mlm/blob/main/config.json'''
    ),
}


class   snake_case_ (lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				"""bridgetower_vision_model"""



       def __init__(     self     ,lowercase=768     ,lowercase=12     ,lowercase=3     ,lowercase=16     ,lowercase=288     ,lowercase=1     ,lowercase=1E-05     ,lowercase=False     ,lowercase=True     ,lowercase=False     ,**lowercase     ,):

             """simple docstring"""





             super().__init__(**lowercase)
             UpperCAmelCase_  :	Optional[int]				      =		hidden_size
             UpperCAmelCase_  :	List[str]				      =		num_hidden_layers
             UpperCAmelCase_  :	Dict				      =		num_channels
             UpperCAmelCase_  :	List[str]				      =		patch_size
             UpperCAmelCase_  :	str				      =		image_size
             UpperCAmelCase_  :	Tuple				      =		initializer_factor
             UpperCAmelCase_  :	str				      =		layer_norm_eps
             UpperCAmelCase_  :	Optional[int]				      =		stop_gradient
             UpperCAmelCase_  :	Dict				      =		share_layernorm
             UpperCAmelCase_  :	Tuple				      =		remove_last_layer



       @classmethod
       def  A_   (     cls     ,lowercase     ,**lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	List[Any]				      =		cls.get_config_dict(lowercase     ,**lowercase)

             if config_dict.get("model_type") == "bridgetower":
                   UpperCAmelCase_  :	Dict				      =		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(lowercase     ,**lowercase)


class   snake_case_ (lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				"""bridgetower_text_model"""



       def __init__(     self     ,lowercase=50265     ,lowercase=768     ,lowercase=12     ,lowercase=12     ,lowercase=1     ,lowercase=3072     ,lowercase="gelu"     ,lowercase=0.1     ,lowercase=0.1     ,lowercase=514     ,lowercase=1     ,lowercase=1E-05     ,lowercase=1     ,lowercase=0     ,lowercase=2     ,lowercase="absolute"     ,lowercase=True     ,**lowercase     ,):

             """simple docstring"""





             super().__init__(**lowercase)

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



       @classmethod
       def  A_   (     cls     ,lowercase     ,**lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	Union[str, Any]				      =		cls.get_config_dict(lowercase     ,**lowercase)

             if config_dict.get("model_type") == "bridgetower":
                   UpperCAmelCase_  :	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(lowercase     ,**lowercase)


class   snake_case_ (lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				"""bridgetower"""



       def __init__(     self     ,lowercase=True     ,lowercase="gelu"     ,lowercase=768     ,lowercase=1     ,lowercase=1E-05     ,lowercase=False     ,lowercase="add"     ,lowercase=12     ,lowercase=6     ,lowercase=False     ,lowercase=False     ,lowercase=None     ,lowercase=None     ,**lowercase     ,):

             """simple docstring"""





             UpperCAmelCase_  :	Dict				      =		kwargs.pop("text_config_dict"     ,lowercase)
             UpperCAmelCase_  :	str				      =		kwargs.pop("vision_config_dict"     ,lowercase)

             super().__init__(**lowercase)
             UpperCAmelCase_  :	Optional[int]				      =		share_cross_modal_transformer_layers
             UpperCAmelCase_  :	List[str]				      =		hidden_act
             UpperCAmelCase_  :	Union[str, Any]				      =		hidden_size
             UpperCAmelCase_  :	Any				      =		initializer_factor
             UpperCAmelCase_  :	int				      =		layer_norm_eps
             UpperCAmelCase_  :	Any				      =		share_link_tower_layers
             UpperCAmelCase_  :	List[Any]				      =		link_tower_type
             UpperCAmelCase_  :	Optional[int]				      =		num_attention_heads
             UpperCAmelCase_  :	int				      =		num_hidden_layers
             UpperCAmelCase_  :	List[Any]				      =		tie_word_embeddings
             UpperCAmelCase_  :	Optional[int]				      =		init_layernorm_from_vision_encoder

             if text_config is None:
                   UpperCAmelCase_  :	Union[str, Any]				      =		{}
                   logger.info("`text_config` is `None`. Initializing the `BridgeTowerTextConfig` with default values.")

             if vision_config is None:
                   UpperCAmelCase_  :	Any				      =		{}
                   logger.info("`vision_config` is `None`. Initializing the `BridgeTowerVisionConfig` with default values.")

             UpperCAmelCase_  :	List[str]				      =		BridgeTowerTextConfig(**lowercase)
             UpperCAmelCase_  :	List[str]				      =		BridgeTowerVisionConfig(**lowercase)



       @classmethod
       def  A_   (     cls     ,lowercase     ,lowercase     ,**lowercase):

             """simple docstring"""





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



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	List[str]				      =		copy.deepcopy(self.__dict__)
             UpperCAmelCase_  :	List[Any]				      =		self.text_config.to_dict()
             UpperCAmelCase_  :	Union[str, Any]				      =		self.vision_config.to_dict()
             UpperCAmelCase_  :	List[str]				      =		self.__class__.model_type
             return output





# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch

import math
from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, randn_tensor
from .scheduling_utils import SchedulerMixin, SchedulerOutput


@dataclass
class   snake_case_ (lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				42
       _lowerCamelCase			    =				42


class   snake_case_ (lowercase__ , lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				1



       @register_to_config
       def __init__(     self     ,lowercase = 2000     ,lowercase = 0.15     ,lowercase = 0.01     ,lowercase = 1348.0     ,lowercase = 1E-5     ,lowercase = 1     ,):

             """simple docstring"""





             UpperCAmelCase_  :	Optional[int]				      =		sigma_max

             # setable values
             UpperCAmelCase_  :	Optional[int]				      =		None

             self.set_sigmas(lowercase     ,lowercase     ,lowercase     ,lowercase)



       def  A_   (     self     ,lowercase     ,lowercase = None):

             """simple docstring"""





             return sample



       def  A_   (     self     ,lowercase     ,lowercase = None     ,lowercase = None):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		sampling_eps if sampling_eps is not None else self.config.sampling_eps

             UpperCAmelCase_  :	List[Any]				      =		torch.linspace(1     ,lowercase     ,lowercase     ,device=lowercase)



       def  A_   (     self     ,lowercase     ,lowercase = None     ,lowercase = None     ,lowercase = None):

             """simple docstring"""





             UpperCAmelCase_  :	Any				      =		sigma_min if sigma_min is not None else self.config.sigma_min
             UpperCAmelCase_  :	int				      =		sigma_max if sigma_max is not None else self.config.sigma_max
             UpperCAmelCase_  :	Union[str, Any]				      =		sampling_eps if sampling_eps is not None else self.config.sampling_eps
             if self.timesteps is None:
                   self.set_timesteps(lowercase     ,lowercase)

             UpperCAmelCase_  :	Union[str, Any]				      =		sigma_min * (sigma_max / sigma_min) ** (self.timesteps / sampling_eps)
             UpperCAmelCase_  :	Optional[int]				      =		torch.exp(torch.linspace(math.log(lowercase)     ,math.log(lowercase)     ,lowercase))
             UpperCAmelCase_  :	Any				      =		torch.tensor([sigma_min * (sigma_max / sigma_min) ** t for t in self.timesteps])



       def  A_   (     self     ,lowercase     ,lowercase):

             """simple docstring"""





             return torch.where(
                 timesteps == 0     ,torch.zeros_like(t.to(timesteps.device))     ,self.discrete_sigmas[timesteps - 1].to(timesteps.device)     ,)



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase = None     ,lowercase = True     ,):

             """simple docstring"""





             if self.timesteps is None:
                   raise ValueError(
                       "`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler")

             UpperCAmelCase_  :	Optional[int]				      =		timestep * torch.ones(
                 sample.shape[0]     ,device=sample.device)  # torch.repeat_interleave(timestep, sample.shape[0])
             UpperCAmelCase_  :	Tuple				      =		(timestep * (len(self.timesteps) - 1)).long()

             # mps requires indices to be in the same device, so we use cpu as is the default with cuda
             UpperCAmelCase_  :	Optional[int]				      =		timesteps.to(self.discrete_sigmas.device)

             UpperCAmelCase_  :	Optional[Any]				      =		self.discrete_sigmas[timesteps].to(sample.device)
             UpperCAmelCase_  :	Optional[Any]				      =		self.get_adjacent_sigma(lowercase     ,lowercase).to(sample.device)
             UpperCAmelCase_  :	Any				      =		torch.zeros_like(lowercase)
             UpperCAmelCase_  :	Dict				      =		(sigma**2 - adjacent_sigma**2) ** 0.5

             # equation 6 in the paper: the model_output modeled by the network is grad_x log pt(x)
             # also equation 47 shows the analog from SDE models to ancestral sampling methods
             UpperCAmelCase_  :	Dict				      =		diffusion.flatten()
             while len(diffusion.shape) < len(sample.shape):
                   UpperCAmelCase_  :	List[str]				      =		diffusion.unsqueeze(-1)
             UpperCAmelCase_  :	List[Any]				      =		drift - diffusion**2 * model_output

             #  equation 6: sample noise for the diffusion term of
             UpperCAmelCase_  :	Union[str, Any]				      =		randn_tensor(
                 sample.shape     ,layout=sample.layout     ,generator=lowercase     ,device=sample.device     ,dtype=sample.dtype)
             UpperCAmelCase_  :	Any				      =		sample - drift  # subtract because `dt` is a small negative timestep
             # TODO is the variable diffusion the correct scaling term for the noise?
             UpperCAmelCase_  :	Tuple				      =		prev_sample_mean + diffusion * noise  # add impact of diffusion field g

             if not return_dict:
                   return (prev_sample, prev_sample_mean)

             return SdeVeOutput(prev_sample=lowercase     ,prev_sample_mean=lowercase)



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase = None     ,lowercase = True     ,):

             """simple docstring"""





             if self.timesteps is None:
                   raise ValueError(
                       "`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler")

             # For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z"
             # sample noise for correction
             UpperCAmelCase_  :	int				      =		randn_tensor(sample.shape     ,layout=sample.layout     ,generator=lowercase).to(sample.device)

             # compute step size from the model_output, the noise, and the snr
             UpperCAmelCase_  :	Union[str, Any]				      =		torch.norm(model_output.reshape(model_output.shape[0]     ,-1)     ,dim=-1).mean()
             UpperCAmelCase_  :	Optional[Any]				      =		torch.norm(noise.reshape(noise.shape[0]     ,-1)     ,dim=-1).mean()
             UpperCAmelCase_  :	List[Any]				      =		(self.config.snr * noise_norm / grad_norm) ** 2 * 2
             UpperCAmelCase_  :	Optional[Any]				      =		step_size * torch.ones(sample.shape[0]).to(sample.device)
             # self.repeat_scalar(step_size, sample.shape[0])

             # compute corrected sample: model_output term and noise term
             UpperCAmelCase_  :	Any				      =		step_size.flatten()
             while len(step_size.shape) < len(sample.shape):
                   UpperCAmelCase_  :	Tuple				      =		step_size.unsqueeze(-1)
             UpperCAmelCase_  :	Dict				      =		sample + step_size * model_output
             UpperCAmelCase_  :	int				      =		prev_sample_mean + ((step_size * 2) ** 0.5) * noise

             if not return_dict:
                   return (prev_sample,)

             return SchedulerOutput(prev_sample=lowercase)



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,):

             """simple docstring"""





             UpperCAmelCase_  :	Any				      =		timesteps.to(original_samples.device)
             UpperCAmelCase_  :	List[str]				      =		self.discrete_sigmas.to(original_samples.device)[timesteps]
             UpperCAmelCase_  :	Tuple				      =		(
                 noise * sigmas[:, None, None, None]
                 if noise is not None
                 else torch.randn_like(lowercase) * sigmas[:, None, None, None]
             )
             UpperCAmelCase_  :	Tuple				      =		noise + original_samples
             return noisy_samples



       def __len__(     self):

             """simple docstring"""





             return self.config.num_train_timesteps
