
"""simple docstring"""





from __future__ import annotations

import copy
import inspect
import unittest

import numpy as np

from transformers import is_tf_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow
from transformers.utils import cached_property

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


if is_tf_available():
				import tensorflow as tf

				from transformers import (
				    TF_LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
				    TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
				    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
				    TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
				    LayoutLMvaConfig,
				    TFLayoutLMvaForQuestionAnswering,
				    TFLayoutLMvaForSequenceClassification,
				    TFLayoutLMvaForTokenClassification,
				    TFLayoutLMvaModel,
				)

if is_vision_available():
				from PIL import Image

				from transformers import LayoutLMvaImageProcessor






class     _UpperCAmelCase:



		def __init__(					self				,		__a				,		__a=2				,		__a=3				,		__a=4				,		__a=2				,		__a=7				,		__a=True				,		__a=True				,		__a=True				,		__a=True				,		__a=99				,		__a=36				,		__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=6				,		__a=6				,		__a=3				,		__a=4				,		__a=None				,		__a=10_00				,		)       ->       Any:






							'''simple docstring'''
							_UpperCamelCase							 =  parent
							_UpperCamelCase							 =  batch_size
							_UpperCamelCase							 =  num_channels
							_UpperCamelCase							 =  image_size
							_UpperCamelCase							 =  patch_size
							_UpperCamelCase							 =  is_training
							_UpperCamelCase							 =  use_input_mask
							_UpperCamelCase							 =  use_token_type_ids
							_UpperCamelCase							 =  use_labels
							_UpperCamelCase							 =  vocab_size
							_UpperCamelCase							 =  hidden_size
							_UpperCamelCase							 =  num_hidden_layers
							_UpperCamelCase							 =  num_attention_heads
							_UpperCamelCase							 =  intermediate_size
							_UpperCamelCase							 =  hidden_act
							_UpperCamelCase							 =  hidden_dropout_prob
							_UpperCamelCase							 =  attention_probs_dropout_prob
							_UpperCamelCase							 =  max_position_embeddings
							_UpperCamelCase							 =  type_vocab_size
							_UpperCamelCase							 =  type_sequence_label_size
							_UpperCamelCase							 =  initializer_range
							_UpperCamelCase							 =  coordinate_size
							_UpperCamelCase							 =  shape_size
							_UpperCamelCase							 =  num_labels
							_UpperCamelCase							 =  num_choices
							_UpperCamelCase							 =  scope
							_UpperCamelCase							 =  range_bbox

							# LayoutLMv3's sequence length equals the number of text tokens + number of patches + 1 (we add 1 for the CLS token)
							_UpperCamelCase							 =  text_seq_length
							_UpperCamelCase							 =  (image_size // patch_size) ** 2 + 1
							_UpperCamelCase							 =  self.text_seq_length + self.image_seq_length



		def 						UpperCAmelCase							(					self)       ->       Union[str, Any]:






							'''simple docstring'''
							_UpperCamelCase							 =  ids_tensor([self.batch_size, self.text_seq_length]				,		self.vocab_size)

							_UpperCamelCase							 =  ids_tensor([self.batch_size, self.text_seq_length, 4]				,		self.range_bbox)
							_UpperCamelCase							 =  bbox.numpy()
							# Ensure that bbox is legal
							for i in range(bbox.shape[0]):
												for j in range(bbox.shape[1]):
																	if bbox[i, j, 3] < bbox[i, j, 1]:
																						_UpperCamelCase							 =  bbox[i, j, 3]
																						_UpperCamelCase							 =  bbox[i, j, 1]
																						_UpperCamelCase							 =  tmp_coordinate
																	if bbox[i, j, 2] < bbox[i, j, 0]:
																						_UpperCamelCase							 =  bbox[i, j, 2]
																						_UpperCamelCase							 =  bbox[i, j, 0]
																						_UpperCamelCase							 =  tmp_coordinate
							_UpperCamelCase							 =  tf.constant(__a)

							_UpperCamelCase							 =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

							_UpperCamelCase							 =  None
							if self.use_input_mask:
												_UpperCamelCase							 =  random_attention_mask([self.batch_size, self.text_seq_length])

							_UpperCamelCase							 =  None
							if self.use_token_type_ids:
												_UpperCamelCase							 =  ids_tensor([self.batch_size, self.text_seq_length]				,		self.type_vocab_size)

							_UpperCamelCase							 =  None
							_UpperCamelCase							 =  None
							if self.use_labels:
												_UpperCamelCase							 =  ids_tensor([self.batch_size]				,		self.type_sequence_label_size)
												_UpperCamelCase							 =  ids_tensor([self.batch_size, self.text_seq_length]				,		self.num_labels)

							_UpperCamelCase							 =  LayoutLMvaConfig(
							    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				,		initializer_range=self.initializer_range				,		coordinate_size=self.coordinate_size				,		shape_size=self.shape_size				,		input_size=self.image_size				,		patch_size=self.patch_size				,		)

							return config, input_ids, bbox, pixel_values, token_type_ids, input_mask, sequence_labels, token_labels



		def 						UpperCAmelCase							(					self				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a)       ->       Optional[Any]:






							'''simple docstring'''
							_UpperCamelCase							 =  TFLayoutLMvaModel(config=__a)

							# text + image
							_UpperCamelCase							 =  model(__a				,		pixel_values=__a				,		training=__a)
							_UpperCamelCase							 =  model(
							    __a				,		bbox=__a				,		pixel_values=__a				,		attention_mask=__a				,		token_type_ids=__a				,		training=__a				,		)
							_UpperCamelCase							 =  model(__a				,		bbox=__a				,		pixel_values=__a				,		training=__a)

							self.parent.assertEqual(result.last_hidden_state.shape				,		(self.batch_size, self.seq_length, self.hidden_size))

							# text only
							_UpperCamelCase							 =  model(__a				,		training=__a)
							self.parent.assertEqual(
							    result.last_hidden_state.shape				,		(self.batch_size, self.text_seq_length, self.hidden_size))

							# image only
							_UpperCamelCase							 =  model({'''pixel_values''': pixel_values}				,		training=__a)
							self.parent.assertEqual(
							    result.last_hidden_state.shape				,		(self.batch_size, self.image_seq_length, self.hidden_size))



		def 						UpperCAmelCase							(					self				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a)       ->       Optional[int]:






							'''simple docstring'''
							_UpperCamelCase							 =  self.num_labels
							_UpperCamelCase							 =  TFLayoutLMvaForSequenceClassification(config=__a)
							_UpperCamelCase							 =  model(
							    __a				,		bbox=__a				,		pixel_values=__a				,		attention_mask=__a				,		token_type_ids=__a				,		labels=__a				,		training=__a				,		)
							self.parent.assertEqual(result.logits.shape				,		(self.batch_size, self.num_labels))



		def 						UpperCAmelCase							(					self				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a)       ->       Any:






							'''simple docstring'''
							_UpperCamelCase							 =  self.num_labels
							_UpperCamelCase							 =  TFLayoutLMvaForTokenClassification(config=__a)
							_UpperCamelCase							 =  model(
							    __a				,		bbox=__a				,		pixel_values=__a				,		attention_mask=__a				,		token_type_ids=__a				,		labels=__a				,		training=__a				,		)
							self.parent.assertEqual(result.logits.shape				,		(self.batch_size, self.text_seq_length, self.num_labels))



		def 						UpperCAmelCase							(					self				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a)       ->       str:






							'''simple docstring'''
							_UpperCamelCase							 =  2
							_UpperCamelCase							 =  TFLayoutLMvaForQuestionAnswering(config=__a)
							_UpperCamelCase							 =  model(
							    __a				,		bbox=__a				,		pixel_values=__a				,		attention_mask=__a				,		token_type_ids=__a				,		start_positions=__a				,		end_positions=__a				,		training=__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))




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






							'''simple docstring'''
							_UpperCamelCase							 =  self.prepare_config_and_inputs()
							((_UpperCamelCase) ,  (_UpperCamelCase) ,  (_UpperCamelCase) ,  (_UpperCamelCase) ,  (_UpperCamelCase) ,  (_UpperCamelCase) ,  (_UpperCamelCase) ,  (_UpperCamelCase))							 =  config_and_inputs
							_UpperCamelCase							 =  {
							    '''input_ids''': input_ids,
							    '''bbox''': bbox,
							    '''pixel_values''': pixel_values,
							    '''token_type_ids''': token_type_ids,
							    '''attention_mask''': input_mask,
							}
							return config, inputs_dict






@require_tf
class     _UpperCAmelCase(    lowerCamelCase      ,     lowerCamelCase      ,     unittest.TestCase	):
		lowercase__						=      (
		    (
		        TFLayoutLMvaModel,
		        TFLayoutLMvaForQuestionAnswering,
		        TFLayoutLMvaForSequenceClassification,
		        TFLayoutLMvaForTokenClassification,
		    )
		    if is_tf_available()
		    else ()
		)
		lowercase__						=      (
		    {'document-question-answering': TFLayoutLMvaForQuestionAnswering, 'feature-extraction': TFLayoutLMvaModel}
		    if is_tf_available()
		    else {}
		)

		lowercase__						=      False
		lowercase__						=      False
		lowercase__						=      False



		def 						UpperCAmelCase							(					self				,		__a				,		__a				,		__a				,		__a				,		__a)       ->       int:






							'''simple docstring'''
							return True



		def 						UpperCAmelCase							(					self				,		__a				,		__a				,		__a=False)       ->       dict:






							'''simple docstring'''
							_UpperCamelCase							 =  copy.deepcopy(__a)

							if model_class in get_values(__a):
												_UpperCamelCase							 =  {
												    k: tf.tile(tf.expand_dims(__a				,		1)				,		(1, self.model_tester.num_choices) + (1,) * (v.ndim - 1))
												    if isinstance(__a				,		tf.Tensor) and v.ndim > 0
												    else v
												    for k, v in inputs_dict.items()
												}

							if return_labels:
												if model_class in get_values(__a):
																	_UpperCamelCase							 =  tf.ones(self.model_tester.batch_size				,		dtype=tf.intaa)
												elif model_class in get_values(__a):
																	_UpperCamelCase							 =  tf.zeros(self.model_tester.batch_size				,		dtype=tf.intaa)
																	_UpperCamelCase							 =  tf.zeros(self.model_tester.batch_size				,		dtype=tf.intaa)
												elif model_class in get_values(__a):
																	_UpperCamelCase							 =  tf.zeros(self.model_tester.batch_size				,		dtype=tf.intaa)
												elif model_class in get_values(__a):
																	_UpperCamelCase							 =  tf.zeros(
																	    (self.model_tester.batch_size, self.model_tester.text_seq_length)				,		dtype=tf.intaa)

							return inputs_dict



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






							'''simple docstring'''
							_UpperCamelCase							 =  TFLayoutLMvaModelTester(self)
							_UpperCamelCase							 =  ConfigTester(self				,		config_class=__a				,		hidden_size=37)



		def 						UpperCAmelCase							(					self)       ->       Any:






							'''simple docstring'''
							self.config_tester.run_common_tests()



		def 						UpperCAmelCase							(					self)       ->       str:






							'''simple docstring'''
							_UpperCamelCase ,  _UpperCamelCase							 =  self.model_tester.prepare_config_and_inputs_for_common()
							for model_class in self.all_model_classes:
												_UpperCamelCase							 =  model_class(__a)
												if getattr(__a				,		'''hf_compute_loss'''				,		__a):
																	# The number of elements in the loss should be the same as the number of elements in the label
																	_UpperCamelCase							 =  self._prepare_for_class(inputs_dict.copy()				,		__a				,		return_labels=__a)
																	_UpperCamelCase							 =  prepared_for_class[
																	    sorted(prepared_for_class.keys() - inputs_dict.keys()				,		reverse=__a)[0]
																	]
																	_UpperCamelCase							 =  added_label.shape.as_list()[:1]

																	# Test that model correctly compute the loss with kwargs
																	_UpperCamelCase							 =  self._prepare_for_class(inputs_dict.copy()				,		__a				,		return_labels=__a)
																	_UpperCamelCase							 =  prepared_for_class.pop('''input_ids''')

																	_UpperCamelCase							 =  model(__a				,		**__a)[0]
																	self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1])

																	# Test that model correctly compute the loss when we mask some positions
																	_UpperCamelCase							 =  self._prepare_for_class(inputs_dict.copy()				,		__a				,		return_labels=__a)
																	_UpperCamelCase							 =  prepared_for_class.pop('''input_ids''')
																	if "labels" in prepared_for_class:
																						_UpperCamelCase							 =  prepared_for_class['''labels'''].numpy()
																						if len(labels.shape) > 1 and labels.shape[1] != 1:
																											_UpperCamelCase							 =  -1_00
																											_UpperCamelCase							 =  tf.convert_to_tensor(__a)
																											_UpperCamelCase							 =  model(__a				,		**__a)[0]
																											self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1])
																											self.assertTrue(not np.any(np.isnan(loss.numpy())))

                # Test that model correctly compute the loss with a dict
																	_UpperCamelCase							 =  self._prepare_for_class(inputs_dict.copy()				,		__a				,		return_labels=__a)
																	_UpperCamelCase							 =  model(__a)[0]
																	self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1])

																	# Test that model correctly compute the loss with a tuple
																	_UpperCamelCase							 =  self._prepare_for_class(inputs_dict.copy()				,		__a				,		return_labels=__a)

																	# Get keys that were added with the _prepare_for_class function
																	_UpperCamelCase							 =  prepared_for_class.keys() - inputs_dict.keys()
																	_UpperCamelCase							 =  inspect.signature(model.call).parameters
																	_UpperCamelCase							 =  list(signature.keys())

																	# Create a dictionary holding the location of the tensors in the tuple
																	_UpperCamelCase							 =  {0: '''input_ids'''}
																	for label_key in label_keys:
																						_UpperCamelCase							 =  signature_names.index(__a)
																						_UpperCamelCase							 =  label_key
																	_UpperCamelCase							 =  sorted(tuple_index_mapping.items())
																	# Initialize a list with their default values, update the values and convert to a tuple
																	_UpperCamelCase							 =  []

																	for name in signature_names:
																						if name != "kwargs":
																											list_input.append(signature[name].default)

																	for index, value in sorted_tuple_index_mapping:
																						_UpperCamelCase							 =  prepared_for_class[value]

																	_UpperCamelCase							 =  tuple(__a)

																	# Send to model
																	_UpperCamelCase							 =  model(tuple_input[:-1])[0]

																	self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1])



		def 						UpperCAmelCase							(					self)       ->       int:






							'''simple docstring'''
							(
							    (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  
							)							 =  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(__a				,		__a				,		__a				,		__a				,		__a				,		__a)



		def 						UpperCAmelCase							(					self)       ->       Union[str, Any]:






							'''simple docstring'''
							(
							    (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  
							)							 =  self.model_tester.prepare_config_and_inputs()
							for type in ["absolute", "relative_key", "relative_key_query"]:
												_UpperCamelCase							 =  type
												self.model_tester.create_and_check_model(__a				,		__a				,		__a				,		__a				,		__a				,		__a)



		def 						UpperCAmelCase							(					self)       ->       int:






							'''simple docstring'''
							(
							    (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  
							)							 =  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(
							    __a				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a)



		def 						UpperCAmelCase							(					self)       ->       Tuple:






							'''simple docstring'''
							(
							    (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  
							)							 =  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(
							    __a				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a)



		def 						UpperCAmelCase							(					self)       ->       str:






							'''simple docstring'''
							(
							    (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  (
							    _UpperCamelCase
							) ,  
							)							 =  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(
							    __a				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a)




		@slow
		def 						UpperCAmelCase							(					self)       ->       Optional[int]:






							'''simple docstring'''
							for model_name in TF_LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												_UpperCamelCase							 =  TFLayoutLMvaModel.from_pretrained(__a)
												self.assertIsNotNone(__a)





def       lowerCamelCase__		(			)							->		int:


					"""simple docstring"""





					_UpperCamelCase							 =  Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''	)
					return image






@require_tf
class     _UpperCAmelCase(    unittest.TestCase	):



		@cached_property
		def 						UpperCAmelCase							(					self)       ->       List[str]:






							'''simple docstring'''
							return LayoutLMvaImageProcessor(apply_ocr=__a) if is_vision_available() else None




		@slow
		def 						UpperCAmelCase							(					self)       ->       Tuple:






							'''simple docstring'''
							_UpperCamelCase							 =  TFLayoutLMvaModel.from_pretrained('''microsoft/layoutlmv3-base''')

							_UpperCamelCase							 =  self.default_image_processor
							_UpperCamelCase							 =  prepare_img()
							_UpperCamelCase							 =  image_processor(images=__a				,		return_tensors='''tf''').pixel_values

							_UpperCamelCase							 =  tf.constant([[1, 2]])
							_UpperCamelCase							 =  tf.expand_dims(tf.constant([[1, 2, 3, 4], [5, 6, 7, 8]])				,		axis=0)

							# forward pass
							_UpperCamelCase							 =  model(input_ids=__a				,		bbox=__a				,		pixel_values=__a				,		training=__a)

							# verify the logits
							_UpperCamelCase							 =  (1, 1_99, 7_68)
							self.assertEqual(outputs.last_hidden_state.shape				,		__a)

							_UpperCamelCase							 =  tf.constant(
							    [[-0.0529, 0.3618, 0.1632], [-0.1587, -0.1667, -0.0400], [-0.1557, -0.1671, -0.0505]])

							self.assertTrue(np.allclose(outputs.last_hidden_state[0, :3, :3]				,		__a				,		atol=1e-4))








"""simple docstring"""





import inspect
from typing import Optional, Union

import numpy as np
import PIL
import torch
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPFeatureExtractor, CLIPModel, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    DPMSolverMultistepScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput
from diffusers.utils import (
    PIL_INTERPOLATION,
    randn_tensor,
)



def       lowerCamelCase__		(			__snake_case,     __snake_case,     __snake_case	)							->		int:


					"""simple docstring"""





					if isinstance(__snake_case,     torch.Tensor	):
										return image
					elif isinstance(__snake_case,     PIL.Image.Image	):
										_UpperCamelCase							 =  [image]

					if isinstance(image[0],     PIL.Image.Image	):
										_UpperCamelCase							 =  [np.array(i.resize((w, h),     resample=PIL_INTERPOLATION['''lanczos''']	)	)[None, :] for i in image]
										_UpperCamelCase							 =  np.concatenate(__snake_case,     axis=0	)
										_UpperCamelCase							 =  np.array(__snake_case	).astype(np.floataa	) / 255.0
										_UpperCamelCase							 =  image.transpose(0,     3,     1,     2	)
										_UpperCamelCase							 =  2.0 * image - 1.0
										_UpperCamelCase							 =  torch.from_numpy(__snake_case	)
					elif isinstance(image[0],     torch.Tensor	):
										_UpperCamelCase							 =  torch.cat(__snake_case,     dim=0	)
					return image



def       lowerCamelCase__		(			__snake_case,     __snake_case,     __snake_case,     __snake_case=0.9995	)							->		List[Any]:


					"""simple docstring"""





					if not isinstance(__snake_case,     np.ndarray	):
										_UpperCamelCase							 =  True
										_UpperCamelCase							 =  va.device
										_UpperCamelCase							 =  va.cpu().numpy()
										_UpperCamelCase							 =  va.cpu().numpy()

					_UpperCamelCase							 =  np.sum(va * va / (np.linalg.norm(__snake_case	) * np.linalg.norm(__snake_case	))	)
					if np.abs(__snake_case	) > DOT_THRESHOLD:
										_UpperCamelCase							 =  (1 - t) * va + t * va
					else:
										_UpperCamelCase							 =  np.arccos(__snake_case	)
										_UpperCamelCase							 =  np.sin(__snake_case	)
										_UpperCamelCase							 =  theta_a * t
										_UpperCamelCase							 =  np.sin(__snake_case	)
										_UpperCamelCase							 =  np.sin(theta_a - theta_t	) / sin_theta_a
										_UpperCamelCase							 =  sin_theta_t / sin_theta_a
										_UpperCamelCase							 =  sa * va + sa * va

					if inputs_are_torch:
										_UpperCamelCase							 =  torch.from_numpy(__snake_case	).to(__snake_case	)

					return va



def       lowerCamelCase__		(			__snake_case,     __snake_case	)							->		List[str]:


					"""simple docstring"""





					_UpperCamelCase							 =  F.normalize(__snake_case,     dim=-1	)
					_UpperCamelCase							 =  F.normalize(__snake_case,     dim=-1	)
					return (x - y).norm(dim=-1	).div(2	).arcsin().pow(2	).mul(2	)



def       lowerCamelCase__		(			__snake_case,     __snake_case	)							->		Optional[int]:


					"""simple docstring"""





					for param in model.parameters():
										_UpperCamelCase							 =  value






class     _UpperCAmelCase(    lowerCamelCase	):



		def __init__(					self				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a=None				,		__a=None				,		__a=None				,		)       ->       List[str]:






							'''simple docstring'''
							super().__init__()
							self.register_modules(
							    vae=__a				,		text_encoder=__a				,		clip_model=__a				,		tokenizer=__a				,		unet=__a				,		scheduler=__a				,		feature_extractor=__a				,		coca_model=__a				,		coca_tokenizer=__a				,		coca_transform=__a				,		)
							_UpperCamelCase							 =  (
							    feature_extractor.size
							    if isinstance(feature_extractor.size				,		__a)
							    else feature_extractor.size['''shortest_edge''']
							)
							_UpperCamelCase							 =  transforms.Normalize(mean=feature_extractor.image_mean				,		std=feature_extractor.image_std)
							set_requires_grad(self.text_encoder				,		__a)
							set_requires_grad(self.clip_model				,		__a)



		def 						UpperCAmelCase							(					self				,		__a = "auto")       ->       Union[str, Any]:






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



		def 						UpperCAmelCase							(					self)       ->       Optional[int]:






							'''simple docstring'''
							self.enable_attention_slicing(__a)



		def 						UpperCAmelCase							(					self)       ->       int:






							'''simple docstring'''
							set_requires_grad(self.vae				,		__a)



		def 						UpperCAmelCase							(					self)       ->       Union[str, Any]:






							'''simple docstring'''
							set_requires_grad(self.vae				,		__a)



		def 						UpperCAmelCase							(					self)       ->       Tuple:






							'''simple docstring'''
							set_requires_grad(self.unet				,		__a)



		def 						UpperCAmelCase							(					self)       ->       int:






							'''simple docstring'''
							set_requires_grad(self.unet				,		__a)



		def 						UpperCAmelCase							(					self				,		__a				,		__a				,		__a)       ->       Any:






							'''simple docstring'''
							# get the original timestep using init_timestep
							_UpperCamelCase							 =  min(int(num_inference_steps * strength)				,		__a)

							_UpperCamelCase							 =  max(num_inference_steps - init_timestep				,		0)
							_UpperCamelCase							 =  self.scheduler.timesteps[t_start:]

							return timesteps, num_inference_steps - t_start



		def 						UpperCAmelCase							(					self				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a=None)       ->       Tuple:






							'''simple docstring'''
							if not isinstance(__a				,		torch.Tensor):
												raise ValueError(F'''`image` has to be of type `torch.Tensor` but is {type(__a)}''')

							_UpperCamelCase							 =  image.to(device=__a				,		dtype=__a)

							if isinstance(__a				,		__a):
												_UpperCamelCase							 =  [
												    self.vae.encode(image[i : i + 1]).latent_dist.sample(generator[i]) for i in range(__a)
												]
												_UpperCamelCase							 =  torch.cat(__a				,		dim=0)
							else:
												_UpperCamelCase							 =  self.vae.encode(__a).latent_dist.sample(__a)

							# Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
							_UpperCamelCase							 =  0.1_8215 * init_latents
							_UpperCamelCase							 =  init_latents.repeat_interleave(__a				,		dim=0)

							_UpperCamelCase							 =  randn_tensor(init_latents.shape				,		generator=__a				,		device=__a				,		dtype=__a)

							# get latents
							_UpperCamelCase							 =  self.scheduler.add_noise(__a				,		__a				,		__a)
							_UpperCamelCase							 =  init_latents

							return latents



		def 						UpperCAmelCase							(					self				,		__a)       ->       str:






							'''simple docstring'''
							_UpperCamelCase							 =  self.coca_transform(__a).unsqueeze(0)
							with torch.no_grad(), torch.cuda.amp.autocast():
												_UpperCamelCase							 =  self.coca_model.generate(transformed_image.to(device=self.device				,		dtype=self.coca_model.dtype))
							_UpperCamelCase							 =  self.coca_tokenizer.decode(generated[0].cpu().numpy())
							return generated.split('''<end_of_text>''')[0].replace('''<start_of_text>'''				,		'''''').rstrip(''' .,''')



		def 						UpperCAmelCase							(					self				,		__a				,		__a)       ->       Optional[int]:






							'''simple docstring'''
							_UpperCamelCase							 =  self.feature_extractor.preprocess(__a)
							_UpperCamelCase							 =  torch.from_numpy(clip_image_input['''pixel_values'''][0]).unsqueeze(0).to(self.device).half()
							_UpperCamelCase							 =  self.clip_model.get_image_features(__a)
							_UpperCamelCase							 =  image_embeddings_clip / image_embeddings_clip.norm(p=2				,		dim=-1				,		keepdim=__a)
							_UpperCamelCase							 =  image_embeddings_clip.repeat_interleave(__a				,		dim=0)
							return image_embeddings_clip



		@torch.enable_grad()
		def 						UpperCAmelCase							(					self				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a				,		)       ->       Tuple:






							'''simple docstring'''
							_UpperCamelCase							 =  latents.detach().requires_grad_()

							_UpperCamelCase							 =  self.scheduler.scale_model_input(__a				,		__a)

							# predict the noise residual
							_UpperCamelCase							 =  self.unet(__a				,		__a				,		encoder_hidden_states=__a).sample

							if isinstance(self.scheduler				,		(PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)):
												_UpperCamelCase							 =  self.scheduler.alphas_cumprod[timestep]
												_UpperCamelCase							 =  1 - alpha_prod_t
												# compute predicted original sample from predicted noise also called
												# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
												_UpperCamelCase							 =  (latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

												_UpperCamelCase							 =  torch.sqrt(__a)
												_UpperCamelCase							 =  pred_original_sample * (fac) + latents * (1 - fac)
							elif isinstance(self.scheduler				,		__a):
												_UpperCamelCase							 =  self.scheduler.sigmas[index]
												_UpperCamelCase							 =  latents - sigma * noise_pred
							else:
												raise ValueError(F'''scheduler type {type(self.scheduler)} not supported''')

							# Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
							_UpperCamelCase							 =  1 / 0.1_8215 * sample
							_UpperCamelCase							 =  self.vae.decode(__a).sample
							_UpperCamelCase							 =  (image / 2 + 0.5).clamp(0				,		1)

							_UpperCamelCase							 =  transforms.Resize(self.feature_extractor_size)(__a)
							_UpperCamelCase							 =  self.normalize(__a).to(latents.dtype)

							_UpperCamelCase							 =  self.clip_model.get_image_features(__a)
							_UpperCamelCase							 =  image_embeddings_clip / image_embeddings_clip.norm(p=2				,		dim=-1				,		keepdim=__a)

							_UpperCamelCase							 =  spherical_dist_loss(__a				,		__a).mean() * clip_guidance_scale

							_UpperCamelCase							 =  -torch.autograd.grad(__a				,		__a)[0]

							if isinstance(self.scheduler				,		__a):
												_UpperCamelCase							 =  latents.detach() + grads * (sigma**2)
												_UpperCamelCase							 =  noise_pred_original
							else:
												_UpperCamelCase							 =  noise_pred_original - torch.sqrt(__a) * grads
							return noise_pred, latents




		@torch.no_grad()
		def __call__(					self				,		__a				,		__a				,		__a = None				,		__a = None				,		__a = 5_12				,		__a = 5_12				,		__a = 0.6				,		__a = 50				,		__a = 7.5				,		__a = 1				,		__a = 0.0				,		__a = 1_00				,		__a = None				,		__a = "pil"				,		__a = True				,		__a = 0.8				,		__a = 0.1				,		__a = 0.1				,		)       ->       Dict:






							'''simple docstring'''
							if isinstance(__a				,		__a) and len(__a) != batch_size:
												raise ValueError(F'''You have passed {batch_size} batch_size, but only {len(__a)} generators.''')

							if height % 8 != 0 or width % 8 != 0:
												raise ValueError(F'''`height` and `width` have to be divisible by 8 but are {height} and {width}.''')

							if isinstance(__a				,		torch.Generator) and batch_size > 1:
												_UpperCamelCase							 =  [generator] + [None] * (batch_size - 1)

							_UpperCamelCase							 =  [
							    ('''model''', self.coca_model is None),
							    ('''tokenizer''', self.coca_tokenizer is None),
							    ('''transform''', self.coca_transform is None),
							]
							_UpperCamelCase							 =  [x[0] for x in coca_is_none if x[1]]
							_UpperCamelCase							 =  ''', '''.join(__a)
							# generate prompts with coca model if prompt is None
							if content_prompt is None:
												if len(__a):
																	raise ValueError(
																	    F'''Content prompt is None and CoCa [{coca_is_none_str}] is None.'''
																	    F'''Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.''')
												_UpperCamelCase							 =  self.get_image_description(__a)
							if style_prompt is None:
												if len(__a):
																	raise ValueError(
																	    F'''Style prompt is None and CoCa [{coca_is_none_str}] is None.'''
																	    F''' Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.''')
												_UpperCamelCase							 =  self.get_image_description(__a)

							# get prompt text embeddings for content and style
							_UpperCamelCase							 =  self.tokenizer(
							    __a				,		padding='''max_length'''				,		max_length=self.tokenizer.model_max_length				,		truncation=__a				,		return_tensors='''pt'''				,		)
							_UpperCamelCase							 =  self.text_encoder(content_text_input.input_ids.to(self.device))[0]

							_UpperCamelCase							 =  self.tokenizer(
							    __a				,		padding='''max_length'''				,		max_length=self.tokenizer.model_max_length				,		truncation=__a				,		return_tensors='''pt'''				,		)
							_UpperCamelCase							 =  self.text_encoder(style_text_input.input_ids.to(self.device))[0]

							_UpperCamelCase							 =  slerp(__a				,		__a				,		__a)

							# duplicate text embeddings for each generation per prompt
							_UpperCamelCase							 =  text_embeddings.repeat_interleave(__a				,		dim=0)

							# set timesteps
							_UpperCamelCase							 =  '''offset''' in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys())
							_UpperCamelCase							 =  {}
							if accepts_offset:
												_UpperCamelCase							 =  1

							self.scheduler.set_timesteps(__a				,		**__a)
							# Some schedulers like PNDM have timesteps as arrays
							# It's more optimized to move all timesteps to correct device beforehand
							self.scheduler.timesteps.to(self.device)

							_UpperCamelCase ,  _UpperCamelCase							 =  self.get_timesteps(__a				,		__a				,		self.device)
							_UpperCamelCase							 =  timesteps[:1].repeat(__a)

							# Preprocess image
							_UpperCamelCase							 =  preprocess(__a				,		__a				,		__a)
							_UpperCamelCase							 =  self.prepare_latents(
							    __a				,		__a				,		__a				,		text_embeddings.dtype				,		self.device				,		__a)

							_UpperCamelCase							 =  preprocess(__a				,		__a				,		__a)
							_UpperCamelCase							 =  self.prepare_latents(
							    __a				,		__a				,		__a				,		text_embeddings.dtype				,		self.device				,		__a)

							_UpperCamelCase							 =  slerp(__a				,		__a				,		__a)

							if clip_guidance_scale > 0:
												_UpperCamelCase							 =  self.get_clip_image_embeddings(__a				,		__a)
												_UpperCamelCase							 =  self.get_clip_image_embeddings(__a				,		__a)
												_UpperCamelCase							 =  slerp(
												    __a				,		__a				,		__a)

							# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
							# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
							# corresponds to doing no classifier free guidance.
							_UpperCamelCase							 =  guidance_scale > 1.0
							# get unconditional embeddings for classifier free guidance
							if do_classifier_free_guidance:
												_UpperCamelCase							 =  content_text_input.input_ids.shape[-1]
												_UpperCamelCase							 =  self.tokenizer(['''''']				,		padding='''max_length'''				,		max_length=__a				,		return_tensors='''pt''')
												_UpperCamelCase							 =  self.text_encoder(uncond_input.input_ids.to(self.device))[0]
												# duplicate unconditional embeddings for each generation per prompt
												_UpperCamelCase							 =  uncond_embeddings.repeat_interleave(__a				,		dim=0)

												# For classifier free guidance, we need to do two forward passes.
												# Here we concatenate the unconditional and text embeddings into a single batch
												# to avoid doing two forward passes
												_UpperCamelCase							 =  torch.cat([uncond_embeddings, text_embeddings])

							# get the initial random noise unless the user supplied it

							# Unlike in other pipelines, latents need to be generated in the target device
							# for 1-to-1 results reproducibility with the CompVis implementation.
							# However this currently doesn't work in `mps`.
							_UpperCamelCase							 =  (batch_size, self.unet.config.in_channels, height // 8, width // 8)
							_UpperCamelCase							 =  text_embeddings.dtype
							if latents is None:
												if self.device.type == "mps":
																	# randn does not work reproducibly on mps
																	_UpperCamelCase							 =  torch.randn(__a				,		generator=__a				,		device='''cpu'''				,		dtype=__a).to(
																	    self.device)
												else:
																	_UpperCamelCase							 =  torch.randn(__a				,		generator=__a				,		device=self.device				,		dtype=__a)
							else:
												if latents.shape != latents_shape:
																	raise ValueError(F'''Unexpected latents shape, got {latents.shape}, expected {latents_shape}''')
												_UpperCamelCase							 =  latents.to(self.device)

							# scale the initial noise by the standard deviation required by the scheduler
							_UpperCamelCase							 =  latents * self.scheduler.init_noise_sigma

							# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
							# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
							# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
							# and should be between [0, 1]
							_UpperCamelCase							 =  '''eta''' in set(inspect.signature(self.scheduler.step).parameters.keys())
							_UpperCamelCase							 =  {}
							if accepts_eta:
												_UpperCamelCase							 =  eta

							# check if the scheduler accepts generator
							_UpperCamelCase							 =  '''generator''' in set(inspect.signature(self.scheduler.step).parameters.keys())
							if accepts_generator:
												_UpperCamelCase							 =  generator

							with self.progress_bar(total=__a):
												for i, t in enumerate(__a):
																	# expand the latents if we are doing classifier free guidance
																	_UpperCamelCase							 =  torch.cat([latents] * 2) if do_classifier_free_guidance else latents
																	_UpperCamelCase							 =  self.scheduler.scale_model_input(__a				,		__a)

																	# predict the noise residual
																	_UpperCamelCase							 =  self.unet(__a				,		__a				,		encoder_hidden_states=__a).sample

																	# perform classifier free guidance
																	if do_classifier_free_guidance:
																						_UpperCamelCase ,  _UpperCamelCase							 =  noise_pred.chunk(2)
																						_UpperCamelCase							 =  noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

																	# perform clip guidance
																	if clip_guidance_scale > 0:
																						_UpperCamelCase							 =  (
																						    text_embeddings.chunk(2)[1] if do_classifier_free_guidance else text_embeddings
																						)
																						_UpperCamelCase ,  _UpperCamelCase							 =  self.cond_fn(
																						    __a				,		__a				,		__a				,		__a				,		__a				,		__a				,		__a				,		)

																	# compute the previous noisy sample x_t -> x_t-1
																	_UpperCamelCase							 =  self.scheduler.step(__a				,		__a				,		__a				,		**__a).prev_sample

        # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
							_UpperCamelCase							 =  1 / 0.1_8215 * latents
							_UpperCamelCase							 =  self.vae.decode(__a).sample

							_UpperCamelCase							 =  (image / 2 + 0.5).clamp(0				,		1)
							_UpperCamelCase							 =  image.cpu().permute(0				,		2				,		3				,		1).numpy()

							if output_type == "pil":
												_UpperCamelCase							 =  self.numpy_to_pil(__a)

							if not return_dict:
												return (image, None)

							return StableDiffusionPipelineOutput(images=__a				,		nsfw_content_detected=__a)






