



from typing import Callable, List, Optional, Union

import PIL
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPSegForImageSegmentation,
    CLIPSegProcessor,
    CLIPTextModel,
    CLIPTokenizer,
)

from diffusers import DiffusionPipeline
from diffusers.configuration_utils import FrozenDict
from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion import StableDiffusionInpaintPipeline
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
from diffusers.utils import deprecate, is_accelerate_available, logging


a_						:     List[Any]	     =		logging.get_logger(__name__)  # pylint: disable=invalid-name





class    lowerCamelCase__  (						UpperCAmelCase_):
     """simple docstring"""



     def __init__(self				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a				,			):






           '''simple docstring'''



           super().__init__()

           if hasattr(scheduler.config				,			"steps_offset"    ) and scheduler.config.steps_offset != 1:
                 lowerCamelCase     :   List[str]								=		(
                     F"""The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"""
                     F""" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure """
                     "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
                     " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
                     " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
                     " file"
                 )
                 deprecate("steps_offset!=1"				,			"1.0.0"				,			__a				,			standard_warn=__a    )
                 lowerCamelCase     :   str								=		dict(scheduler.config    )
                 lowerCamelCase     :   str								=		1
                 lowerCamelCase     :   List[Any]								=		FrozenDict(__a    )

           if hasattr(scheduler.config				,			"skip_prk_steps"    ) and scheduler.config.skip_prk_steps is False:
                 lowerCamelCase     :   str								=		(
                     F"""The configuration file of this scheduler: {scheduler} has not set the configuration"""
                     " `skip_prk_steps`. `skip_prk_steps` should be set to True in the configuration file. Please make"
                     " sure to update the config accordingly as not setting `skip_prk_steps` in the config might lead to"
                     " incorrect results in future versions. If you have downloaded this checkpoint from the Hugging Face"
                     " Hub, it would be very nice if you could open a Pull request for the"
                     " `scheduler/scheduler_config.json` file"
                 )
                 deprecate("skip_prk_steps not set"				,			"1.0.0"				,			__a				,			standard_warn=__a    )
                 lowerCamelCase     :   Any								=		dict(scheduler.config    )
                 lowerCamelCase     :   str								=		True
                 lowerCamelCase     :   Dict								=		FrozenDict(__a    )

           if safety_checker is None:
                 logger.warning(
                     F"""You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"""
                     " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
                     " results in services or applications open to the public. Both the diffusers team and Hugging Face"
                     " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
                     " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
                     " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."    )

           self.register_modules(
               segmentation_model=__a				,			segmentation_processor=__a				,			vae=__a				,			text_encoder=__a				,			tokenizer=__a				,			unet=__a				,			scheduler=__a				,			safety_checker=__a				,			feature_extractor=__a				,			)







     def 						_a      (self				,			__a = "auto"    ):






           '''simple docstring'''



           if slice_size == "auto":
                 # half the attention head size is usually a good trade-off between
                 # speed and memory
                 lowerCamelCase     :   Dict								=		self.unet.config.attention_head_dim // 2
           self.unet.set_attention_slice(__a    )







     def 						_a      (self    ):






           '''simple docstring'''



           self.enable_attention_slicing(__a    )







     def 						_a      (self    ):






           '''simple docstring'''



           if is_accelerate_available():
                 from accelerate import cpu_offload
           else:
                 raise ImportError("Please install accelerate via `pip install accelerate`"    )

           lowerCamelCase     :   Optional[int]								=		torch.device("cuda"    )

           for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae, self.safety_checker]:
                 if cpu_offloaded_model is not None:
                       cpu_offload(__a				,			__a    )







     @property
     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
     def 						_a      (self    ):






           '''simple docstring'''



           if self.device != torch.device("meta"    ) or not hasattr(self.unet				,			"_hf_hook"    ):
                 return self.device
           for module in self.unet.modules():
                 if (
                     hasattr(__a				,			"_hf_hook"    )
                     and hasattr(module._hf_hook				,			"execution_device"    )
                     and module._hf_hook.execution_device is not None
                 ):
                       return torch.device(module._hf_hook.execution_device    )
           return self.device







     @torch.no_grad()
     def __call__(self				,			__a				,			__a				,			__a				,			__a = 5_12				,			__a = 5_12				,			__a = 50				,			__a = 7.5				,			__a = None				,			__a = 1				,			__a = 0.0				,			__a = None				,			__a = None				,			__a = "pil"				,			__a = True				,			__a = None				,			__a = 1				,			**__a				,			):






           '''simple docstring'''



           lowerCamelCase     :   int								=		self.segmentation_processor(
               text=[text]				,			images=[image]				,			padding="max_length"				,			return_tensors="pt"    ).to(self.device    )
           lowerCamelCase     :   int								=		self.segmentation_model(**__a    )
           lowerCamelCase     :   str								=		torch.sigmoid(outputs.logits    ).cpu().detach().unsqueeze(-1    ).numpy()
           lowerCamelCase     :   Any								=		self.numpy_to_pil(__a    )[0].resize(image.size    )

           # Run inpainting pipeline with the generated mask
           lowerCamelCase     :   List[str]								=		StableDiffusionInpaintPipeline(
               vae=self.vae				,			text_encoder=self.text_encoder				,			tokenizer=self.tokenizer				,			unet=self.unet				,			scheduler=self.scheduler				,			safety_checker=self.safety_checker				,			feature_extractor=self.feature_extractor				,			)
           return inpainting_pipeline(
               prompt=__a				,			image=__a				,			mask_image=__a				,			height=__a				,			width=__a				,			num_inference_steps=__a				,			guidance_scale=__a				,			negative_prompt=__a				,			num_images_per_prompt=__a				,			eta=__a				,			generator=__a				,			latents=__a				,			output_type=__a				,			return_dict=__a				,			callback=__a				,			callback_steps=__a				,			)




from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import numpy
     import tensorflow as tf

     from transformers import (
         TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
         TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
         TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
         BertConfig,
         DPRConfig,
         TFDPRContextEncoder,
         TFDPRQuestionEncoder,
         TFDPRReader,
     )





class    lowerCamelCase__  :
     """simple docstring"""



     def __init__(self				,			__a				,			__a=13				,			__a=7				,			__a=True				,			__a=True				,			__a=True				,			__a=True				,			__a=99				,			__a=32				,			__a=2				,			__a=4				,			__a=37				,			__a="gelu"				,			__a=0.1				,			__a=0.1				,			__a=5_12				,			__a=16				,			__a=2				,			__a=0.02				,			__a=3				,			__a=4				,			__a=None				,			__a=0				,			):






           '''simple docstring'''



           lowerCamelCase								=		parent
           lowerCamelCase								=		batch_size
           lowerCamelCase								=		seq_length
           lowerCamelCase								=		is_training
           lowerCamelCase								=		use_input_mask
           lowerCamelCase								=		use_token_type_ids
           lowerCamelCase								=		use_labels
           lowerCamelCase								=		vocab_size
           lowerCamelCase								=		hidden_size
           lowerCamelCase								=		num_hidden_layers
           lowerCamelCase								=		num_attention_heads
           lowerCamelCase								=		intermediate_size
           lowerCamelCase								=		hidden_act
           lowerCamelCase								=		hidden_dropout_prob
           lowerCamelCase								=		attention_probs_dropout_prob
           lowerCamelCase								=		max_position_embeddings
           lowerCamelCase								=		type_vocab_size
           lowerCamelCase								=		type_sequence_label_size
           lowerCamelCase								=		initializer_range
           lowerCamelCase								=		num_labels
           lowerCamelCase								=		num_choices
           lowerCamelCase								=		scope
           lowerCamelCase								=		projection_dim







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		ids_tensor([self.batch_size, self.seq_length]				,			self.vocab_size    )

           lowerCamelCase								=		None
           if self.use_input_mask:
                 # follow test_modeling_tf_ctrl.py
                 lowerCamelCase								=		random_attention_mask([self.batch_size, self.seq_length]    )

           lowerCamelCase								=		None
           if self.use_token_type_ids:
                 lowerCamelCase								=		ids_tensor([self.batch_size, self.seq_length]				,			self.type_vocab_size    )

           lowerCamelCase								=		None
           lowerCamelCase								=		None
           lowerCamelCase								=		None
           if self.use_labels:
                 lowerCamelCase								=		ids_tensor([self.batch_size]				,			self.type_sequence_label_size    )
                 lowerCamelCase								=		ids_tensor([self.batch_size, self.seq_length]				,			self.num_labels    )
                 lowerCamelCase								=		ids_tensor([self.batch_size]				,			self.num_choices    )

           lowerCamelCase								=		BertConfig(
               vocab_size=self.vocab_size				,			hidden_size=self.hidden_size				,			num_hidden_layers=self.num_hidden_layers				,			num_attention_heads=self.num_attention_heads				,			intermediate_size=self.intermediate_size				,			hidden_act=self.hidden_act				,			hidden_dropout_prob=self.hidden_dropout_prob				,			attention_probs_dropout_prob=self.attention_probs_dropout_prob				,			max_position_embeddings=self.max_position_embeddings				,			type_vocab_size=self.type_vocab_size				,			is_decoder=__a				,			initializer_range=self.initializer_range				,			)
           lowerCamelCase								=		DPRConfig(projection_dim=self.projection_dim				,			**config.to_dict()    )

           return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







     def 						_a      (self				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		TFDPRContextEncoder(config=__a    )
           lowerCamelCase								=		model(__a				,			attention_mask=__a				,			token_type_ids=__a    )
           lowerCamelCase								=		model(__a				,			token_type_ids=__a    )
           lowerCamelCase								=		model(__a    )
           self.parent.assertEqual(result.pooler_output.shape				,			(self.batch_size, self.projection_dim or self.hidden_size)    )







     def 						_a      (self				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		TFDPRQuestionEncoder(config=__a    )
           lowerCamelCase								=		model(__a				,			attention_mask=__a				,			token_type_ids=__a    )
           lowerCamelCase								=		model(__a				,			token_type_ids=__a    )
           lowerCamelCase								=		model(__a    )
           self.parent.assertEqual(result.pooler_output.shape				,			(self.batch_size, self.projection_dim or self.hidden_size)    )







     def 						_a      (self				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		TFDPRReader(config=__a    )
           lowerCamelCase								=		model(__a				,			attention_mask=__a    )

           self.parent.assertEqual(result.start_logits.shape				,			(self.batch_size, self.seq_length)    )
           self.parent.assertEqual(result.end_logits.shape				,			(self.batch_size, self.seq_length)    )
           self.parent.assertEqual(result.relevance_logits.shape				,			(self.batch_size,)    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.prepare_config_and_inputs()
           (
               (
               lowerCamelCase
           )			,							(
               lowerCamelCase
           )			,							(
               lowerCamelCase
           )			,							(
               lowerCamelCase
           )			,							(
               lowerCamelCase
           )			,							(
               lowerCamelCase
           )			,							(
               lowerCamelCase
           )			,							
           )								=		config_and_inputs
           lowerCamelCase								=		{"input_ids": input_ids}
           return config, inputs_dict





@require_tf
class    lowerCamelCase__  (						UpperCAmelCase_       ,    UpperCAmelCase_       ,    unittest.TestCase):
     """simple docstring"""



     _A       =			(
         (
             TFDPRContextEncoder,
             TFDPRQuestionEncoder,
             TFDPRReader,
         )
         if is_tf_available()
         else ()
     )
     _A       =			{'feature-extraction': TFDPRQuestionEncoder} if is_tf_available() else {}

     _A       =			False
     _A       =			False
     _A       =			False
     _A       =			False
     _A       =			False







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		TFDPRModelTester(self    )
           lowerCamelCase								=		ConfigTester(self				,			config_class=__a				,			hidden_size=37    )







     def 						_a      (self    ):






           '''simple docstring'''



           self.config_tester.run_common_tests()







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_dpr_context_encoder(*__a    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_dpr_question_encoder(*__a    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_dpr_reader(*__a    )







     @slow
     def 						_a      (self    ):






           '''simple docstring'''



           for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCamelCase								=		TFDPRContextEncoder.from_pretrained(__a    )
                 self.assertIsNotNone(__a    )

           for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCamelCase								=		TFDPRContextEncoder.from_pretrained(__a    )
                 self.assertIsNotNone(__a    )

           for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCamelCase								=		TFDPRQuestionEncoder.from_pretrained(__a    )
                 self.assertIsNotNone(__a    )

           for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCamelCase								=		TFDPRReader.from_pretrained(__a    )
                 self.assertIsNotNone(__a    )





@require_tf
class    lowerCamelCase__  (						unittest.TestCase):
     """simple docstring"""



     @slow
     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		TFDPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base"    )

           lowerCamelCase								=		tf.constant(
               [[1_01, 75_92, 10_10, 20_03, 20_26, 38_99, 1_01_40, 10_29, 1_02]]    )  # [CLS] hello, is my dog cute? [SEP]
           lowerCamelCase								=		model(__a    )[0]  # embedding shape = (1, 768)
           # compare the actual values for a slice.
           lowerCamelCase								=		tf.constant(
               [
                   [
                       0.0323_6253,
                       0.1275_3335,
                       0.1681_8509,
                       0.0027_9786,
                       0.389_6933,
                       0.2426_4945,
                       0.217_8971,
                       -0.0233_5227,
                       -0.0848_1959,
                       -0.1432_4117,
                   ]
               ]    )
           self.assertTrue(numpy.allclose(output[:, :10].numpy()				,			expected_slice.numpy()				,			atol=1E-4    )    )