

'''simple docstring'''



import copy
import os
from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Union


if TYPE_CHECKING:
     from ...processing_utils import ProcessorMixin
     from ...utils import TensorType

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_UpperCamelCase						:  Any							       =				logging.get_logger(__name__)

_UpperCamelCase						:  Tuple							       =				{
    "google/owlvit-base-patch32": "https://huggingface.co/google/owlvit-base-patch32/resolve/main/config.json",
    "google/owlvit-base-patch16": "https://huggingface.co/google/owlvit-base-patch16/resolve/main/config.json",
    "google/owlvit-large-patch14": "https://huggingface.co/google/owlvit-large-patch14/resolve/main/config.json",
}



class 				_snake_case					(   a_			):
  SCREAMING_SNAKE_CASE			:       str				=     '''owlvit_text_model'''

  def __init__(				self      , _SCREAMING_SNAKE_CASE=4_94_08      , _SCREAMING_SNAKE_CASE=5_12      , _SCREAMING_SNAKE_CASE=20_48      , _SCREAMING_SNAKE_CASE=12      , _SCREAMING_SNAKE_CASE=8      , _SCREAMING_SNAKE_CASE=16      , _SCREAMING_SNAKE_CASE="quick_gelu"      , _SCREAMING_SNAKE_CASE=1e-5      , _SCREAMING_SNAKE_CASE=0.0      , _SCREAMING_SNAKE_CASE=0.02      , _SCREAMING_SNAKE_CASE=1.0      , _SCREAMING_SNAKE_CASE=0      , _SCREAMING_SNAKE_CASE=4_94_06      , _SCREAMING_SNAKE_CASE=4_94_07      , **_SCREAMING_SNAKE_CASE      , ):





      '''simple docstring'''

      super().__init__(pad_token_id=_SCREAMING_SNAKE_CASE      , bos_token_id=_SCREAMING_SNAKE_CASE      , eos_token_id=_SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							vocab_size
      lowerCAmelCase 						=							hidden_size
      lowerCAmelCase 						=							intermediate_size
      lowerCAmelCase 						=							num_hidden_layers
      lowerCAmelCase 						=							num_attention_heads
      lowerCAmelCase 						=							max_position_embeddings
      lowerCAmelCase 						=							hidden_act
      lowerCAmelCase 						=							layer_norm_eps
      lowerCAmelCase 						=							attention_dropout
      lowerCAmelCase 						=							initializer_range
      lowerCAmelCase 						=							initializer_factor


  @classmethod
  def   _SCREAMING_SNAKE_CASE				(				cls      , _SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      cls._set_token_in_kwargs(_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase						,   lowerCAmelCase 						=							cls.get_config_dict(_SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	)

      # get the text config dict if we are loading from OwlViTConfig
      if config_dict.get('model_type'	) == "owlvit":
          lowerCAmelCase 						=							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(_SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	)



class 				_snake_case					(   a_			):
  SCREAMING_SNAKE_CASE			:       Tuple				=     '''owlvit_vision_model'''

  def __init__(				self      , _SCREAMING_SNAKE_CASE=7_68      , _SCREAMING_SNAKE_CASE=30_72      , _SCREAMING_SNAKE_CASE=12      , _SCREAMING_SNAKE_CASE=12      , _SCREAMING_SNAKE_CASE=3      , _SCREAMING_SNAKE_CASE=7_68      , _SCREAMING_SNAKE_CASE=32      , _SCREAMING_SNAKE_CASE="quick_gelu"      , _SCREAMING_SNAKE_CASE=1e-5      , _SCREAMING_SNAKE_CASE=0.0      , _SCREAMING_SNAKE_CASE=0.02      , _SCREAMING_SNAKE_CASE=1.0      , **_SCREAMING_SNAKE_CASE      , ):





      '''simple docstring'''

      super().__init__(**_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							hidden_size
      lowerCAmelCase 						=							intermediate_size
      lowerCAmelCase 						=							num_hidden_layers
      lowerCAmelCase 						=							num_attention_heads
      lowerCAmelCase 						=							num_channels
      lowerCAmelCase 						=							image_size
      lowerCAmelCase 						=							patch_size
      lowerCAmelCase 						=							hidden_act
      lowerCAmelCase 						=							layer_norm_eps
      lowerCAmelCase 						=							attention_dropout
      lowerCAmelCase 						=							initializer_range
      lowerCAmelCase 						=							initializer_factor


  @classmethod
  def   _SCREAMING_SNAKE_CASE				(				cls      , _SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      cls._set_token_in_kwargs(_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase						,   lowerCAmelCase 						=							cls.get_config_dict(_SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	)

      # get the vision config dict if we are loading from OwlViTConfig
      if config_dict.get('model_type'	) == "owlvit":
          lowerCAmelCase 						=							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(_SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	)



class 				_snake_case					(   a_			):
  SCREAMING_SNAKE_CASE			:       List[str]				=     '''owlvit'''
  SCREAMING_SNAKE_CASE			:       Any				=     True

  def __init__(				self      , _SCREAMING_SNAKE_CASE=None      , _SCREAMING_SNAKE_CASE=None      , _SCREAMING_SNAKE_CASE=5_12      , _SCREAMING_SNAKE_CASE=2.6_592      , _SCREAMING_SNAKE_CASE=True      , **_SCREAMING_SNAKE_CASE      , ):





      '''simple docstring'''

      super().__init__(**_SCREAMING_SNAKE_CASE	)

      if text_config is None:
          lowerCAmelCase 						=							{}
          logger.info('text_config is None. Initializing the OwlViTTextConfig with default values.'	)

      if vision_config is None:
          lowerCAmelCase 						=							{}
          logger.info('vision_config is None. initializing the OwlViTVisionConfig with default values.'	)

      lowerCAmelCase 						=							OwlViTTextConfig(**_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							OwlViTVisionConfig(**_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							projection_dim
      lowerCAmelCase 						=							logit_scale_init_value
      lowerCAmelCase 						=							return_dict
      lowerCAmelCase 						=							1.0

  @classmethod
  def   _SCREAMING_SNAKE_CASE				(				cls      , _SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      cls._set_token_in_kwargs(_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase						,   lowerCAmelCase 						=							cls.get_config_dict(_SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	)

      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(_SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	)

  @classmethod
  def   _SCREAMING_SNAKE_CASE				(				cls      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      lowerCAmelCase 						=							{}
      lowerCAmelCase 						=							text_config
      lowerCAmelCase 						=							vision_config

      return cls.from_dict(_SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	)


  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase 						=							copy.deepcopy(self.__dict__	)
      lowerCAmelCase 						=							self.text_config.to_dict()
      lowerCAmelCase 						=							self.vision_config.to_dict()
      lowerCAmelCase 						=							self.__class__.model_type
      return output



class 				_snake_case					(   a_			):

  @property
  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      return OrderedDict(
          [
              ('input_ids', {0: 'batch', 1: 'sequence'}),
              ('pixel_values', {0: 'batch', 1: 'num_channels', 2: 'height', 3: 'width'}),
              ('attention_mask', {0: 'batch', 1: 'sequence'}),
          ]	)

  @property
  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      return OrderedDict(
          [
              ('logits_per_image', {0: 'batch'}),
              ('logits_per_text', {0: 'batch'}),
              ('text_embeds', {0: 'batch'}),
              ('image_embeds', {0: 'batch'}),
          ]	)

  @property
  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      return 1e-4

  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE = -1      , _SCREAMING_SNAKE_CASE = -1      , _SCREAMING_SNAKE_CASE = None      , ):





      '''simple docstring'''

      lowerCAmelCase 						=							super().generate_dummy_inputs(
          processor.tokenizer      , batch_size=_SCREAMING_SNAKE_CASE      , seq_length=_SCREAMING_SNAKE_CASE      , framework=_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							super().generate_dummy_inputs(
          processor.image_processor      , batch_size=_SCREAMING_SNAKE_CASE      , framework=_SCREAMING_SNAKE_CASE	)
      return {**text_input_dict, **image_input_dict}


  @property
  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      return 14








'''simple docstring'''



import torch

from diffusers import DDPMParallelScheduler

from .test_schedulers import SchedulerCommonTest



class 				_snake_case					(   a_			):
  SCREAMING_SNAKE_CASE			:       Union[str, Any]				=     (DDPMParallelScheduler,)

  def   _SCREAMING_SNAKE_CASE				(				self      , **_SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      lowerCAmelCase 						=							{
          'num_train_timesteps': 10_00,
          'beta_start': 0.0_001,
          'beta_end': 0.02,
          'beta_schedule': 'linear',
          'variance_type': 'fixed_small',
          'clip_sample': True,
      }

      config.update(**_SCREAMING_SNAKE_CASE	)
      return config

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      for timesteps in [1, 5, 1_00, 10_00]:
          self.check_over_configs(num_train_timesteps=_SCREAMING_SNAKE_CASE	)

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      for beta_start, beta_end in zip([0.0_001, 0.001, 0.01, 0.1]      , [0.002, 0.02, 0.2, 2]	):
          self.check_over_configs(beta_start=_SCREAMING_SNAKE_CASE      , beta_end=_SCREAMING_SNAKE_CASE	)

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      for schedule in ["linear", "squaredcos_cap_v2"]:
          self.check_over_configs(beta_schedule=_SCREAMING_SNAKE_CASE	)

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      for variance in ["fixed_small", "fixed_large", "other"]:
          self.check_over_configs(variance_type=_SCREAMING_SNAKE_CASE	)

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      for clip_sample in [True, False]:
          self.check_over_configs(clip_sample=_SCREAMING_SNAKE_CASE	)

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      self.check_over_configs(thresholding=_SCREAMING_SNAKE_CASE	)
      for threshold in [0.5, 1.0, 2.0]:
          for prediction_type in ["epsilon", "sample", "v_prediction"]:
              self.check_over_configs(
                  thresholding=_SCREAMING_SNAKE_CASE      , prediction_type=_SCREAMING_SNAKE_CASE      , sample_max_value=_SCREAMING_SNAKE_CASE      , )

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      for prediction_type in ["epsilon", "sample", "v_prediction"]:
          self.check_over_configs(prediction_type=_SCREAMING_SNAKE_CASE	)

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      for t in [0, 5_00, 9_99]:
          self.check_over_forward(time_step=_SCREAMING_SNAKE_CASE	)

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase 						=							self.scheduler_classes[0]
      lowerCAmelCase 						=							self.get_scheduler_config()
      lowerCAmelCase 						=							scheduler_class(**_SCREAMING_SNAKE_CASE	)

      assert torch.sum(torch.abs(scheduler._get_variance(0	) - 0.0	)	) < 1e-5
      assert torch.sum(torch.abs(scheduler._get_variance(4_87	) - 0.00_979	)	) < 1e-5
      assert torch.sum(torch.abs(scheduler._get_variance(9_99	) - 0.02	)	) < 1e-5

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase 						=							self.scheduler_classes[0]
      lowerCAmelCase 						=							self.get_scheduler_config()
      lowerCAmelCase 						=							scheduler_class(**_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							len(_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							self.dummy_model()
      lowerCAmelCase 						=							self.dummy_sample_deter
      lowerCAmelCase 						=							self.dummy_sample_deter + 0.1
      lowerCAmelCase 						=							self.dummy_sample_deter - 0.1

      lowerCAmelCase 						=							samplea.shape[0]
      lowerCAmelCase 						=							torch.stack([samplea, samplea, samplea]      , dim=0	)
      lowerCAmelCase 						=							torch.arange(_SCREAMING_SNAKE_CASE	)[0:3, None].repeat(1      , _SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							model(samples.flatten(0      , 1	)      , timesteps.flatten(0      , 1	)	)
      lowerCAmelCase 						=							scheduler.batch_step_no_noise(_SCREAMING_SNAKE_CASE      , timesteps.flatten(0      , 1	)      , samples.flatten(0      , 1	)	)

      lowerCAmelCase 						=							torch.sum(torch.abs(_SCREAMING_SNAKE_CASE	)	)
      lowerCAmelCase 						=							torch.mean(torch.abs(_SCREAMING_SNAKE_CASE	)	)

      assert abs(result_sum.item() - 1_153.1_833	) < 1e-2
      assert abs(result_mean.item() - 0.5_005	) < 1e-3

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase 						=							self.scheduler_classes[0]
      lowerCAmelCase 						=							self.get_scheduler_config()
      lowerCAmelCase 						=							scheduler_class(**_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							len(_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							self.dummy_model()
      lowerCAmelCase 						=							self.dummy_sample_deter
      lowerCAmelCase 						=							torch.manual_seed(0	)

      for t in reversed(range(_SCREAMING_SNAKE_CASE	)	):
          # 1. predict noise residual
          lowerCAmelCase 						=							model(_SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE	)

          # 2. predict previous mean of sample x_t-1
          lowerCAmelCase 						=							scheduler.step(_SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE      , generator=_SCREAMING_SNAKE_CASE	).prev_sample

          lowerCAmelCase 						=							pred_prev_sample

      lowerCAmelCase 						=							torch.sum(torch.abs(_SCREAMING_SNAKE_CASE	)	)
      lowerCAmelCase 						=							torch.mean(torch.abs(_SCREAMING_SNAKE_CASE	)	)

      assert abs(result_sum.item() - 258.9_606	) < 1e-2
      assert abs(result_mean.item() - 0.3_372	) < 1e-3

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase 						=							self.scheduler_classes[0]
      lowerCAmelCase 						=							self.get_scheduler_config(prediction_type='v_prediction'	)
      lowerCAmelCase 						=							scheduler_class(**_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							len(_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							self.dummy_model()
      lowerCAmelCase 						=							self.dummy_sample_deter
      lowerCAmelCase 						=							torch.manual_seed(0	)

      for t in reversed(range(_SCREAMING_SNAKE_CASE	)	):
          # 1. predict noise residual
          lowerCAmelCase 						=							model(_SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE	)

          # 2. predict previous mean of sample x_t-1
          lowerCAmelCase 						=							scheduler.step(_SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE      , generator=_SCREAMING_SNAKE_CASE	).prev_sample

          lowerCAmelCase 						=							pred_prev_sample

      lowerCAmelCase 						=							torch.sum(torch.abs(_SCREAMING_SNAKE_CASE	)	)
      lowerCAmelCase 						=							torch.mean(torch.abs(_SCREAMING_SNAKE_CASE	)	)

      assert abs(result_sum.item() - 202.0_296	) < 1e-2
      assert abs(result_mean.item() - 0.2_631	) < 1e-3

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase 						=							self.scheduler_classes[0]
      lowerCAmelCase 						=							self.get_scheduler_config()
      lowerCAmelCase 						=							scheduler_class(**_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							[1_00, 87, 50, 1, 0]

      scheduler.set_timesteps(timesteps=_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							scheduler.timesteps

      for i, timestep in enumerate(_SCREAMING_SNAKE_CASE	):
          if i == len(_SCREAMING_SNAKE_CASE	) - 1:
              lowerCAmelCase 						=							-1
          else:
              lowerCAmelCase 						=							timesteps[i + 1]

          lowerCAmelCase 						=							scheduler.previous_timestep(_SCREAMING_SNAKE_CASE	)
          lowerCAmelCase 						=							prev_t.item()

          self.assertEqual(_SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE	)

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase 						=							self.scheduler_classes[0]
      lowerCAmelCase 						=							self.get_scheduler_config()
      lowerCAmelCase 						=							scheduler_class(**_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							[1_00, 87, 50, 51, 0]

      with self.assertRaises(_SCREAMING_SNAKE_CASE      , msg='`custom_timesteps` must be in descending order.'	):
          scheduler.set_timesteps(timesteps=_SCREAMING_SNAKE_CASE	)

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase 						=							self.scheduler_classes[0]
      lowerCAmelCase 						=							self.get_scheduler_config()
      lowerCAmelCase 						=							scheduler_class(**_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							[1_00, 87, 50, 1, 0]
      lowerCAmelCase 						=							len(_SCREAMING_SNAKE_CASE	)

      with self.assertRaises(_SCREAMING_SNAKE_CASE      , msg='Can only pass one of `num_inference_steps` or `custom_timesteps`.'	):
          scheduler.set_timesteps(num_inference_steps=_SCREAMING_SNAKE_CASE      , timesteps=_SCREAMING_SNAKE_CASE	)


  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase 						=							self.scheduler_classes[0]
      lowerCAmelCase 						=							self.get_scheduler_config()
      lowerCAmelCase 						=							scheduler_class(**_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							[scheduler.config.num_train_timesteps]

      with self.assertRaises(
          _SCREAMING_SNAKE_CASE      , msg='`timesteps` must start before `self.config.train_timesteps`: {scheduler.config.num_train_timesteps}}'      , ):
          scheduler.set_timesteps(timesteps=_SCREAMING_SNAKE_CASE	)





