





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_  							=				logging.get_logger(__name__)

UpperCAmelCase_  							=				{
    '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 					lowercase__ (			__lowerCamelCase       ):

   '''simple docstring'''


   a		:							int     	=	"owlvit_text_model"




   def __init__(					self,	__magic_name__=49408,	__magic_name__=512,	__magic_name__=2048,	__magic_name__=12,	__magic_name__=8,	__magic_name__=16,	__magic_name__="quick_gelu",	__magic_name__=1E-5,	__magic_name__=0.0,	__magic_name__=0.02,	__magic_name__=1.0,	__magic_name__=0,	__magic_name__=49406,	__magic_name__=49407,	**__magic_name__,	)				->   Tuple:






       """simple docstring"""
       super().__init__(pad_token_id=__magic_name__,	bos_token_id=__magic_name__,	eos_token_id=__magic_name__,	**__magic_name__							)

       UpperCamelCase__							:       Optional[int]      =   vocab_size
       UpperCamelCase__							:       List[str]      =   hidden_size
       UpperCamelCase__							:       Optional[int]      =   intermediate_size
       UpperCamelCase__							:       int      =   num_hidden_layers
       UpperCamelCase__							:       Optional[int]      =   num_attention_heads
       UpperCamelCase__							:       str      =   max_position_embeddings
       UpperCamelCase__							:       str      =   hidden_act
       UpperCamelCase__							:       int      =   layer_norm_eps
       UpperCamelCase__							:       Union[str, Any]      =   attention_dropout
       UpperCamelCase__							:       List[Any]      =   initializer_range
       UpperCamelCase__							:       Dict      =   initializer_factor




   @classmethod
   def    UpperCamelCase__						(					cls,	__magic_name__,	**__magic_name__							)				->   "PretrainedConfig":






       """simple docstring"""
       cls._set_token_in_kwargs(__magic_name__							)

       UpperCamelCase__     ,UpperCamelCase__							:       List[str]      =   cls.get_config_dict(__magic_name__,	**__magic_name__							)

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






class 					lowercase__ (			__lowerCamelCase       ):

   '''simple docstring'''


   a		:							Optional[int]     	=	"owlvit_vision_model"




   def __init__(					self,	__magic_name__=768,	__magic_name__=3072,	__magic_name__=12,	__magic_name__=12,	__magic_name__=3,	__magic_name__=768,	__magic_name__=32,	__magic_name__="quick_gelu",	__magic_name__=1E-5,	__magic_name__=0.0,	__magic_name__=0.02,	__magic_name__=1.0,	**__magic_name__,	)				->   Optional[Any]:






       """simple docstring"""
       super().__init__(**__magic_name__							)

       UpperCamelCase__							:       Union[str, Any]      =   hidden_size
       UpperCamelCase__							:       Optional[int]      =   intermediate_size
       UpperCamelCase__							:       Optional[Any]      =   num_hidden_layers
       UpperCamelCase__							:       str      =   num_attention_heads
       UpperCamelCase__							:       List[Any]      =   num_channels
       UpperCamelCase__							:       int      =   image_size
       UpperCamelCase__							:       List[str]      =   patch_size
       UpperCamelCase__							:       Optional[Any]      =   hidden_act
       UpperCamelCase__							:       List[str]      =   layer_norm_eps
       UpperCamelCase__							:       Optional[Any]      =   attention_dropout
       UpperCamelCase__							:       Any      =   initializer_range
       UpperCamelCase__							:       Tuple      =   initializer_factor




   @classmethod
   def    UpperCamelCase__						(					cls,	__magic_name__,	**__magic_name__							)				->   "PretrainedConfig":






       """simple docstring"""
       cls._set_token_in_kwargs(__magic_name__							)

       UpperCamelCase__     ,UpperCamelCase__							:       Union[str, Any]      =   cls.get_config_dict(__magic_name__,	**__magic_name__							)

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






class 					lowercase__ (			__lowerCamelCase       ):

   '''simple docstring'''


   a		:							int     	=	"owlvit"
   a		:							Optional[Any]     	=	True




   def __init__(					self,	__magic_name__=None,	__magic_name__=None,	__magic_name__=512,	__magic_name__=2.6592,	__magic_name__=True,	**__magic_name__,	)				->   List[Any]:






       """simple docstring"""
       super().__init__(**__magic_name__							)

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

       if vision_config is None:
           UpperCamelCase__							:       Union[str, Any]      =   {}
           logger.info('''vision_config is None. initializing the OwlViTVisionConfig with default values.'''							)

       UpperCamelCase__							:       List[str]      =   OwlViTTextConfig(**__magic_name__							)
       UpperCamelCase__							:       Union[str, Any]      =   OwlViTVisionConfig(**__magic_name__							)

       UpperCamelCase__							:       Optional[int]      =   projection_dim
       UpperCamelCase__							:       Union[str, Any]      =   logit_scale_init_value
       UpperCamelCase__							:       int      =   return_dict
       UpperCamelCase__							:       Dict      =   1.0




   @classmethod
   def    UpperCamelCase__						(					cls,	__magic_name__,	**__magic_name__							)				->   "PretrainedConfig":






       """simple docstring"""
       cls._set_token_in_kwargs(__magic_name__							)

       UpperCamelCase__     ,UpperCamelCase__							:       Optional[Any]      =   cls.get_config_dict(__magic_name__,	**__magic_name__							)

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




   @classmethod
   def    UpperCamelCase__						(					cls,	__magic_name__,	__magic_name__,	**__magic_name__							)				->   Union[str, Any]:






       """simple docstring"""
       UpperCamelCase__							:       Optional[Any]      =   {}
       UpperCamelCase__							:       List[str]      =   text_config
       UpperCamelCase__							:       List[str]      =   vision_config

       return cls.from_dict(__magic_name__,	**__magic_name__							)




   def    UpperCamelCase__						(					self							)				->   Dict:






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






class 					lowercase__ (			__lowerCamelCase       ):

   '''simple docstring'''


   @property
   def    UpperCamelCase__						(					self							)				->   Mapping[str, Mapping[int, str]]:






       """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    UpperCamelCase__						(					self							)				->   Mapping[str, Mapping[int, str]]:






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




   @property
   def    UpperCamelCase__						(					self							)				->   float:






       """simple docstring"""
       return 1E-4




   def    UpperCamelCase__						(					self,	__magic_name__,	__magic_name__ = -1,	__magic_name__ = -1,	__magic_name__ = None,	)				->   Mapping[str, Any]:






       """simple docstring"""
       UpperCamelCase__							:       Union[str, Any]      =   super().generate_dummy_inputs(
           processor.tokenizer,	batch_size=__magic_name__,	seq_length=__magic_name__,	framework=__magic_name__							)
       UpperCamelCase__							:       Optional[Any]      =   super().generate_dummy_inputs(
           processor.image_processor,	batch_size=__magic_name__,	framework=__magic_name__							)
       return {**text_input_dict, **image_input_dict}




   @property
   def    UpperCamelCase__						(					self							)				->   int:






       """simple docstring"""
       return 14








import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionTextToImagePipeline
from diffusers.utils.testing_utils import nightly, require_torch_gpu, torch_device


UpperCAmelCase_  							=				False






class 					lowercase__ (			unittest.TestCase       ):

   '''simple docstring'''


   pass






@nightly
@require_torch_gpu
class 					lowercase__ (			unittest.TestCase       ):

   '''simple docstring'''


   def    UpperCamelCase__						(					self							)				->   Optional[Any]:






       """simple docstring"""
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()




   def    UpperCamelCase__						(					self							)				->   List[Any]:






       """simple docstring"""
       UpperCamelCase__							:       Dict      =   VersatileDiffusionTextToImagePipeline.from_pretrained('''shi-labs/versatile-diffusion'''							)
       # remove text_unet
       pipe.remove_unused_weights()
       pipe.to(__magic_name__							)
       pipe.set_progress_bar_config(disable=__magic_name__							)

       UpperCamelCase__							:       int      =   '''A painting of a squirrel eating a burger '''
       UpperCamelCase__							:       str      =   torch.manual_seed(0							)
       UpperCamelCase__							:       int      =   pipe(
           prompt=__magic_name__,	generator=__magic_name__,	guidance_scale=7.5,	num_inference_steps=2,	output_type='''numpy'''							).images

       with tempfile.TemporaryDirectory() as tmpdirname:
           pipe.save_pretrained(__magic_name__							)
           UpperCamelCase__							:       Union[str, Any]      =   VersatileDiffusionTextToImagePipeline.from_pretrained(__magic_name__							)
       pipe.to(__magic_name__							)
       pipe.set_progress_bar_config(disable=__magic_name__							)

       UpperCamelCase__							:       Optional[Any]      =   generator.manual_seed(0							)
       UpperCamelCase__							:       str      =   pipe(
           prompt=__magic_name__,	generator=__magic_name__,	guidance_scale=7.5,	num_inference_steps=2,	output_type='''numpy'''							).images

       assert np.abs(image - new_image							).sum() < 1E-5, "Models don't have the same forward pass"




   def    UpperCamelCase__						(					self							)				->   Any:






       """simple docstring"""
       UpperCamelCase__							:       str      =   VersatileDiffusionTextToImagePipeline.from_pretrained(
           '''shi-labs/versatile-diffusion''',	torch_dtype=torch.floataa							)
       pipe.to(__magic_name__							)
       pipe.set_progress_bar_config(disable=__magic_name__							)

       UpperCamelCase__							:       Optional[int]      =   '''A painting of a squirrel eating a burger '''
       UpperCamelCase__							:       str      =   torch.manual_seed(0							)
       UpperCamelCase__							:       Optional[int]      =   pipe(
           prompt=__magic_name__,	generator=__magic_name__,	guidance_scale=7.5,	num_inference_steps=50,	output_type='''numpy'''							).images

       UpperCamelCase__							:       int      =   image[0, 253:256, 253:256, -1]

       assert image.shape == (1, 512, 512, 3)
       UpperCamelCase__							:       int      =   np.array([0.3367, 0.3169, 0.2656, 0.3870, 0.4790, 0.3796, 0.4009, 0.4878, 0.4778]							)

       assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2

