from logging import getLogger
from typing import Any, Callable, List, Optional, Union

import numpy as np
import PIL
import torch

from ...schedulers import DDPMScheduler
from ..onnx_utils import ORT_TO_NP_TYPE, OnnxRuntimeModel
from ..pipeline_utils import ImagePipelineOutput
from . import StableDiffusionUpscalePipeline


lowerCAmelCase__  	=      getLogger(__name__)


lowerCAmelCase__  	=      4
lowerCAmelCase__  	=      7

lowerCAmelCase__  	=      {
    "float16": torch.floataa,
    "float32": torch.floataa,
}

def 	__lowerCamelCase							(				__a    : List[str]		)      ->       Optional[int]:
	if isinstance(__a      ,							torch.Tensor		):
		return image
	elif isinstance(__a      ,							PIL.Image.Image		):
		_lowercase   =[image]

	if isinstance(image[0]      ,							PIL.Image.Image		):
		_lowercase   ,  _lowercase   =image[0].size
		_lowercase   ,  _lowercase   =(x - x % 64 for x in (w, h))  # resize to integer multiple of 32

		_lowercase   =[np.array(i.resize((w, h)		)		)[None, :] for i in image]
		_lowercase   =np.concatenate(__a      ,							axis=0		)
		_lowercase   =np.array(__a		).astype(np.floataa		) / 255.0
		_lowercase   =image.transpose(0      ,							3      ,							1      ,							2		)
		_lowercase   =2.0 * image - 1.0
		_lowercase   =torch.from_numpy(__a		)
	elif isinstance(image[0]      ,							torch.Tensor		):
		_lowercase   =torch.cat(__a      ,							dim=0		)

	return image
class 	_a  (							lowerCamelCase_					):



	"""simple docstring"""



	def __init__(   self       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_ = 350       ,  ):
		super().__init__(
		    vae=lowerCAmelCase_       ,  text_encoder=lowerCAmelCase_       ,  tokenizer=lowerCAmelCase_       ,  unet=lowerCAmelCase_       ,  low_res_scheduler=lowerCAmelCase_       ,  scheduler=lowerCAmelCase_       ,  safety_checker=lowerCAmelCase_       ,  feature_extractor=lowerCAmelCase_       ,  watermarker=lowerCAmelCase_       ,  max_noise_level=lowerCAmelCase_       ,  )


	def __call__(   self       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_ = 75       ,  lowerCAmelCase_ = 9.0       ,  lowerCAmelCase_ = 20       ,  lowerCAmelCase_ = None       ,  lowerCAmelCase_ = 1       ,  lowerCAmelCase_ = 0.0       ,  lowerCAmelCase_ = None       ,  lowerCAmelCase_ = None       ,  lowerCAmelCase_ = None       ,  lowerCAmelCase_ = None       ,  lowerCAmelCase_ = "pil"       ,  lowerCAmelCase_ = True       ,  lowerCAmelCase_ = None       ,  lowerCAmelCase_ = 1       ,  ):
		self.check_inputs(lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       )

		# 2. Define call parameters
		if prompt is not None and isinstance(lowerCAmelCase_       ,  lowerCAmelCase_       ):
			_lowercase   =1
		elif prompt is not None and isinstance(lowerCAmelCase_       ,  lowerCAmelCase_       ):
			_lowercase   =len(lowerCAmelCase_       )
		else:
			_lowercase   =prompt_embeds.shape[0]

		_lowercase   =self._execution_device
		# 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.
		_lowercase   =guidance_scale > 1.0

		# 3. Encode input prompt
		_lowercase   =self._encode_prompt(
		    lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  prompt_embeds=lowerCAmelCase_       ,  negative_prompt_embeds=lowerCAmelCase_       ,  )

		_lowercase   =ORT_TO_PT_TYPE[str(text_embeddings.dtype       )]

		# 4. Preprocess image
		_lowercase   =preprocess(lowerCAmelCase_       )
		_lowercase   =image.cpu()

		# 5. set timesteps
		self.scheduler.set_timesteps(lowerCAmelCase_       ,  device=lowerCAmelCase_       )
		_lowercase   =self.scheduler.timesteps

		# 5. Add noise to image
		_lowercase   =torch.tensor([noise_level]       ,  dtype=torch.long       ,  device=lowerCAmelCase_       )
		_lowercase   =torch.randn(image.shape       ,  generator=lowerCAmelCase_       ,  device=lowerCAmelCase_       ,  dtype=lowerCAmelCase_       )
		_lowercase   =self.low_res_scheduler.add_noise(lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       )

		_lowercase   =2 if do_classifier_free_guidance else 1
		_lowercase   =np.concatenate([image] * batch_multiplier * num_images_per_prompt       )
		_lowercase   =np.concatenate([noise_level] * image.shape[0]       )

		# 6. Prepare latent variables
		_lowercase   ,  _lowercase   =image.shape[2:]
		_lowercase   =self.prepare_latents(
		    batch_size * num_images_per_prompt       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  )

		# 7. Check that sizes of image and latents match
		_lowercase   =image.shape[1]
		if NUM_LATENT_CHANNELS + num_channels_image != NUM_UNET_INPUT_CHANNELS:
			raise ValueError(
			    "Incorrect configuration settings! The config of `pipeline.unet` expects"
			    F''' {NUM_UNET_INPUT_CHANNELS} but received `num_channels_latents`: {NUM_LATENT_CHANNELS} +'''
			    F''' `num_channels_image`: {num_channels_image} '''
			    F''' = {NUM_LATENT_CHANNELS+num_channels_image}. Please verify the config of'''
			    " `pipeline.unet` or your `image` input."       )

		# 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
		_lowercase   =self.prepare_extra_step_kwargs(lowerCAmelCase_       ,  lowerCAmelCase_       )

		_lowercase   =next(
		    (input.type for input in self.unet.model.get_inputs() if input.name == "timestep")       ,  "tensor(float)"       )
		_lowercase   =ORT_TO_NP_TYPE[timestep_dtype]

		# 9. Denoising loop
		_lowercase   =len(lowerCAmelCase_       ) - num_inference_steps * self.scheduler.order
		with self.progress_bar(total=lowerCAmelCase_       ) as progress_bar:
			for i, t in enumerate(lowerCAmelCase_       ):
				# expand the latents if we are doing classifier free guidance
				_lowercase   =np.concatenate([latents] * 2       ) if do_classifier_free_guidance else latents

				# concat latents, mask, masked_image_latents in the channel dimension
				_lowercase   =self.scheduler.scale_model_input(lowerCAmelCase_       ,  lowerCAmelCase_       )
				_lowercase   =np.concatenate([latent_model_input, image]       ,  axis=1       )

				# timestep to tensor
				_lowercase   =np.array([t]       ,  dtype=lowerCAmelCase_       )

				# predict the noise residual
				_lowercase   =self.unet(
				    sample=lowerCAmelCase_       ,  timestep=lowerCAmelCase_       ,  encoder_hidden_states=lowerCAmelCase_       ,  class_labels=noise_level.astype(np.intaa       )       ,  )[0]

				# perform guidance
				if do_classifier_free_guidance:
					_lowercase   ,  _lowercase   =np.split(lowerCAmelCase_       ,  2       )
				_lowercase   =noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

				# compute the previous noisy sample x_t -> x_t-1
				_lowercase   =self.scheduler.step(
				    torch.from_numpy(lowerCAmelCase_       )       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  **lowerCAmelCase_       ).prev_sample

				# call the callback, if provided
				if i == len(lowerCAmelCase_       ) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
					progress_bar.update()
					if callback is not None and i % callback_steps == 0:
						callback(lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       )

        # 10. Post-processing
		_lowercase   =self.decode_latents(latents.float()       )

		# 11. Convert to PIL
		if output_type == "pil":
			_lowercase   =self.numpy_to_pil(lowerCAmelCase_       )

		if not return_dict:
			return (image,)

		return ImagePipelineOutput(images=lowerCAmelCase_       )


	def __lowerCAmelCase      (   self       ,  lowerCAmelCase_       ):
		_lowercase   =1 / 0.0_8_3_3_3 * latents
		_lowercase   =self.vae(latent_sample=lowerCAmelCase_       )[0]
		_lowercase   =np.clip(image / 2 + 0.5       ,  0       ,  1       )
		_lowercase   =image.transpose((0, 2, 3, 1)       )
		return image


	def __lowerCAmelCase      (   self       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_ = None       ,  lowerCAmelCase_ = None       ,  ):
		if prompt is not None and isinstance(lowerCAmelCase_       ,  lowerCAmelCase_       ):
			_lowercase   =1
		elif prompt is not None and isinstance(lowerCAmelCase_       ,  lowerCAmelCase_       ):
			_lowercase   =len(lowerCAmelCase_       )
		else:
			_lowercase   =prompt_embeds.shape[0]

		if prompt_embeds is None:
			_lowercase   =self.tokenizer(
			    lowerCAmelCase_       ,  padding="max_length"       ,  max_length=self.tokenizer.model_max_length       ,  truncation=lowerCAmelCase_       ,  return_tensors="pt"       ,  )
			_lowercase   =text_inputs.input_ids
			_lowercase   =self.tokenizer(lowerCAmelCase_       ,  padding="longest"       ,  return_tensors="pt"       ).input_ids

			if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
			    lowerCAmelCase_       ,  lowerCAmelCase_       ):
				_lowercase   =self.tokenizer.batch_decode(
				    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]       )
				logger.warning(
				    "The following part of your input was truncated because CLIP can only handle sequences up to"
				    F''' {self.tokenizer.model_max_length} tokens: {removed_text}'''       )

			# no positional arguments to text_encoder
			_lowercase   =self.text_encoder(
			    input_ids=text_input_ids.int().to(lowerCAmelCase_       )       ,  )
			_lowercase   =prompt_embeds[0]

		_lowercase   ,  _lowercase   ,  _lowercase   =prompt_embeds.shape
		# duplicate text embeddings for each generation per prompt, using mps friendly method
		_lowercase   =prompt_embeds.repeat(1       ,  lowerCAmelCase_       )
		_lowercase   =prompt_embeds.reshape(bs_embed * num_images_per_prompt       ,  lowerCAmelCase_       ,  -1       )

		# get unconditional embeddings for classifier free guidance
		if do_classifier_free_guidance and negative_prompt_embeds is None:
			_lowercase   =42
			if negative_prompt is None:
				_lowercase   =[""] * batch_size
			elif type(lowerCAmelCase_       ) is not type(lowerCAmelCase_       ):
				raise TypeError(
				    F'''`negative_prompt` should be the same type to `prompt`, but got {type(lowerCAmelCase_       )} !='''
				    F''' {type(lowerCAmelCase_       )}.'''       )
			elif isinstance(lowerCAmelCase_       ,  lowerCAmelCase_       ):
				_lowercase   =[negative_prompt]
			elif batch_size != len(lowerCAmelCase_       ):
				raise ValueError(
				    F'''`negative_prompt`: {negative_prompt} has batch size {len(lowerCAmelCase_       )}, but `prompt`:'''
				    F''' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'''
				    " the batch size of `prompt`."       )
			else:
				_lowercase   =negative_prompt

			_lowercase   =text_input_ids.shape[-1]
			_lowercase   =self.tokenizer(
			    lowerCAmelCase_       ,  padding="max_length"       ,  max_length=lowerCAmelCase_       ,  truncation=lowerCAmelCase_       ,  return_tensors="pt"       ,  )

			# if hasattr(uncond_input, "attention_mask"):
			#     attention_mask = uncond_input.attention_mask.to(device)
			# else:
			#     attention_mask = None

			_lowercase   =self.text_encoder(
			    input_ids=uncond_input.input_ids.int().to(lowerCAmelCase_       )       ,  )
			_lowercase   =uncond_embeddings[0]

		if do_classifier_free_guidance:
			_lowercase   =uncond_embeddings.shape[1]
			# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
			_lowercase   =uncond_embeddings.repeat(1       ,  lowerCAmelCase_       )
			_lowercase   =uncond_embeddings.reshape(batch_size * num_images_per_prompt       ,  lowerCAmelCase_       ,  -1       )

			# 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
			_lowercase   =np.concatenate([uncond_embeddings, prompt_embeds]       )

		return prompt_embeds







import unittest

from transformers import LayoutLMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      LayoutLMForMaskedLM,
      LayoutLMForQuestionAnswering,
      LayoutLMForSequenceClassification,
      LayoutLMForTokenClassification,
      LayoutLMModel,
  )







class  UpperCamelCase_      :


  """simple docstring"""



  def __init__(						self     :					Tuple  ,	A_     :					List[str]  ,	A_     :					List[Any]=13  ,	A_     :					Optional[int]=7  ,	A_     :					List[Any]=True  ,	A_     :					str=True  ,	A_     :					str=True  ,	A_     :					str=True  ,	A_     :					List[Any]=99  ,	A_     :					str=32  ,	A_     :					Dict=5  ,	A_     :					Tuple=4  ,	A_     :					Optional[Any]=37  ,	A_     :					List[Any]="gelu"  ,	A_     :					Any=0.1  ,	A_     :					Union[str, Any]=0.1  ,	A_     :					List[Any]=512  ,	A_     :					Dict=16  ,	A_     :					List[Any]=2  ,	A_     :					Optional[Any]=0.02  ,	A_     :					List[Any]=3  ,	A_     :					Optional[Any]=4  ,	A_     :					Dict=None  ,	A_     :					str=1000  ,	)       ->			str:
      SCREAMING_SNAKE_CASE__         =    parent
      SCREAMING_SNAKE_CASE__         =    batch_size
      SCREAMING_SNAKE_CASE__         =    seq_length
      SCREAMING_SNAKE_CASE__         =    is_training
      SCREAMING_SNAKE_CASE__         =    use_input_mask
      SCREAMING_SNAKE_CASE__         =    use_token_type_ids
      SCREAMING_SNAKE_CASE__         =    use_labels
      SCREAMING_SNAKE_CASE__         =    vocab_size
      SCREAMING_SNAKE_CASE__         =    hidden_size
      SCREAMING_SNAKE_CASE__         =    num_hidden_layers
      SCREAMING_SNAKE_CASE__         =    num_attention_heads
      SCREAMING_SNAKE_CASE__         =    intermediate_size
      SCREAMING_SNAKE_CASE__         =    hidden_act
      SCREAMING_SNAKE_CASE__         =    hidden_dropout_prob
      SCREAMING_SNAKE_CASE__         =    attention_probs_dropout_prob
      SCREAMING_SNAKE_CASE__         =    max_position_embeddings
      SCREAMING_SNAKE_CASE__         =    type_vocab_size
      SCREAMING_SNAKE_CASE__         =    type_sequence_label_size
      SCREAMING_SNAKE_CASE__         =    initializer_range
      SCREAMING_SNAKE_CASE__         =    num_labels
      SCREAMING_SNAKE_CASE__         =    num_choices
      SCREAMING_SNAKE_CASE__         =    scope
      SCREAMING_SNAKE_CASE__         =    range_bbox



  def   _UpperCamelCase     (						self     :					List[str]						)       ->			int:
      SCREAMING_SNAKE_CASE__         =    ids_tensor([self.batch_size, self.seq_length]  ,	self.vocab_size						)

      SCREAMING_SNAKE_CASE__         =    ids_tensor([self.batch_size, self.seq_length, 4]  ,	self.range_bbox						)
      # 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]:
                  SCREAMING_SNAKE_CASE__         =    bbox[i, j, 3]
                  SCREAMING_SNAKE_CASE__         =    bbox[i, j, 1]
                  SCREAMING_SNAKE_CASE__         =    t
              if bbox[i, j, 2] < bbox[i, j, 0]:
                  SCREAMING_SNAKE_CASE__         =    bbox[i, j, 2]
                  SCREAMING_SNAKE_CASE__         =    bbox[i, j, 0]
                  SCREAMING_SNAKE_CASE__         =    t

      SCREAMING_SNAKE_CASE__         =    None
      if self.use_input_mask:
          SCREAMING_SNAKE_CASE__         =    ids_tensor([self.batch_size, self.seq_length]  ,	vocab_size=2						)

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

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

      SCREAMING_SNAKE_CASE__         =    self.get_config()

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



  def   _UpperCamelCase     (						self     :					List[str]						)       ->			int:
      return LayoutLMConfig(
          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  ,	)



  def   _UpperCamelCase     (						self     :					List[Any]  ,	A_     :					Any  ,	A_     :					List[str]  ,	A_     :					Optional[int]  ,	A_     :					Dict  ,	A_     :					int  ,	A_     :					List[str]  ,	A_     :					Any  ,	A_     :					Tuple						)       ->			List[str]:
      SCREAMING_SNAKE_CASE__         =    LayoutLMModel(config=A_						)
      model.to(A_						)
      model.eval()
      SCREAMING_SNAKE_CASE__         =    model(A_  ,	A_  ,	attention_mask=A_  ,	token_type_ids=A_						)
      SCREAMING_SNAKE_CASE__         =    model(A_  ,	A_  ,	token_type_ids=A_						)
      SCREAMING_SNAKE_CASE__         =    model(A_  ,	A_						)
      self.parent.assertEqual(result.last_hidden_state.shape  ,	(self.batch_size, self.seq_length, self.hidden_size)						)
      self.parent.assertEqual(result.pooler_output.shape  ,	(self.batch_size, self.hidden_size)						)



  def   _UpperCamelCase     (						self     :					int  ,	A_     :					int  ,	A_     :					Any  ,	A_     :					Tuple  ,	A_     :					Dict  ,	A_     :					int  ,	A_     :					Any  ,	A_     :					int  ,	A_     :					Union[str, Any]						)       ->			Tuple:
      SCREAMING_SNAKE_CASE__         =    LayoutLMForMaskedLM(config=A_						)
      model.to(A_						)
      model.eval()
      SCREAMING_SNAKE_CASE__         =    model(A_  ,	A_  ,	attention_mask=A_  ,	token_type_ids=A_  ,	labels=A_						)
      self.parent.assertEqual(result.logits.shape  ,	(self.batch_size, self.seq_length, self.vocab_size)						)



  def   _UpperCamelCase     (						self     :					Dict  ,	A_     :					Union[str, Any]  ,	A_     :					Union[str, Any]  ,	A_     :					Dict  ,	A_     :					Optional[int]  ,	A_     :					Tuple  ,	A_     :					Optional[int]  ,	A_     :					Optional[int]  ,	A_     :					Optional[Any]						)       ->			Optional[Any]:
      SCREAMING_SNAKE_CASE__         =    self.num_labels
      SCREAMING_SNAKE_CASE__         =    LayoutLMForSequenceClassification(A_						)
      model.to(A_						)
      model.eval()
      SCREAMING_SNAKE_CASE__         =    model(
          A_  ,	A_  ,	attention_mask=A_  ,	token_type_ids=A_  ,	labels=A_						)
      self.parent.assertEqual(result.logits.shape  ,	(self.batch_size, self.num_labels)						)



  def   _UpperCamelCase     (						self     :					Any  ,	A_     :					Any  ,	A_     :					Any  ,	A_     :					str  ,	A_     :					Any  ,	A_     :					Optional[Any]  ,	A_     :					Tuple  ,	A_     :					Tuple  ,	A_     :					Any						)       ->			Any:
      SCREAMING_SNAKE_CASE__         =    self.num_labels
      SCREAMING_SNAKE_CASE__         =    LayoutLMForTokenClassification(config=A_						)
      model.to(A_						)
      model.eval()
      SCREAMING_SNAKE_CASE__         =    model(A_  ,	A_  ,	attention_mask=A_  ,	token_type_ids=A_  ,	labels=A_						)
      self.parent.assertEqual(result.logits.shape  ,	(self.batch_size, self.seq_length, self.num_labels)						)



  def   _UpperCamelCase     (						self     :					Optional[int]  ,	A_     :					str  ,	A_     :					Optional[Any]  ,	A_     :					int  ,	A_     :					str  ,	A_     :					Optional[Any]  ,	A_     :					List[str]  ,	A_     :					Union[str, Any]  ,	A_     :					Optional[Any]						)       ->			Any:
      SCREAMING_SNAKE_CASE__         =    LayoutLMForQuestionAnswering(config=A_						)
      model.to(A_						)
      model.eval()
      SCREAMING_SNAKE_CASE__         =    model(
          A_  ,	bbox=A_  ,	attention_mask=A_  ,	token_type_ids=A_  ,	start_positions=A_  ,	end_positions=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]						)       ->			List[str]:
      SCREAMING_SNAKE_CASE__         =    self.prepare_config_and_inputs()
      (
          (
          SCREAMING_SNAKE_CASE__
      )  ,		(
          SCREAMING_SNAKE_CASE__
      )  ,		(
          SCREAMING_SNAKE_CASE__
      )  ,		(
          SCREAMING_SNAKE_CASE__
      )  ,		(
          SCREAMING_SNAKE_CASE__
      )  ,		(
          SCREAMING_SNAKE_CASE__
      )  ,		(
          SCREAMING_SNAKE_CASE__
      )  ,		(
          SCREAMING_SNAKE_CASE__
      )  ,		
      )         =    config_and_inputs
      SCREAMING_SNAKE_CASE__         =    {
          "input_ids": input_ids,
          "bbox": bbox,
          "token_type_ids": token_type_ids,
          "attention_mask": input_mask,
      }
      return config, inputs_dict







@require_torch
class  UpperCamelCase_      (	_snake_case		,  _snake_case		,  unittest.TestCase      ):


  """simple docstring"""
  SCREAMING_SNAKE_CASE__					:			Optional[int]					    =   (
      (
          LayoutLMModel,
          LayoutLMForMaskedLM,
          LayoutLMForSequenceClassification,
          LayoutLMForTokenClassification,
          LayoutLMForQuestionAnswering,
      )
      if is_torch_available()
      else None
  )
  SCREAMING_SNAKE_CASE__					:			Union[str, Any]					    =   (
      {
          """document-question-answering""": LayoutLMForQuestionAnswering,
          """feature-extraction""": LayoutLMModel,
          """fill-mask""": LayoutLMForMaskedLM,
          """text-classification""": LayoutLMForSequenceClassification,
          """token-classification""": LayoutLMForTokenClassification,
          """zero-shot""": LayoutLMForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  SCREAMING_SNAKE_CASE__					:			Any					    =   True



  def   _UpperCamelCase     (						self     :					int						)       ->			Optional[Any]:
      SCREAMING_SNAKE_CASE__         =    LayoutLMModelTester(self						)
      SCREAMING_SNAKE_CASE__         =    ConfigTester(self  ,	config_class=A_  ,	hidden_size=37						)



  def   _UpperCamelCase     (						self     :					Tuple						)       ->			Tuple:
      self.config_tester.run_common_tests()



  def   _UpperCamelCase     (						self     :					Dict						)       ->			Tuple:
      SCREAMING_SNAKE_CASE__         =    self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*A_						)



  def   _UpperCamelCase     (						self     :					Optional[Any]						)       ->			Dict:
      SCREAMING_SNAKE_CASE__         =    self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
          SCREAMING_SNAKE_CASE__         =    type
          self.model_tester.create_and_check_model(*A_						)



  def   _UpperCamelCase     (						self     :					Tuple						)       ->			Tuple:
      SCREAMING_SNAKE_CASE__         =    self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*A_						)



  def   _UpperCamelCase     (						self     :					Dict						)       ->			Optional[Any]:
      SCREAMING_SNAKE_CASE__         =    self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_sequence_classification(*A_						)



  def   _UpperCamelCase     (						self     :					List[str]						)       ->			List[str]:
      SCREAMING_SNAKE_CASE__         =    self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*A_						)



  def   _UpperCamelCase     (						self     :					int						)       ->			Optional[Any]:
      SCREAMING_SNAKE_CASE__         =    self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_question_answering(*A_						)







  @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."						)
  def   _UpperCamelCase     (						self     :					List[Any]						)       ->			str:
      pass





def 				__a							(       ):
    # Here we prepare a batch of 2 sequences to test a LayoutLM forward pass on:
    # fmt: off
    SCREAMING_SNAKE_CASE__         =    torch.tensor([[101,1019,1014,1016,1037,12849,4747,1004,14246,2278,5439,4524,5002,2930,2193,2930,4341,3208,1005,1055,2171,2848,11300,3531,102],[101,4070,4034,7020,1024,3058,1015,1013,2861,1013,6070,19274,2772,6205,27814,16147,16147,4343,2047,10283,10969,14389,1012,2338,102]]			,							device=__A						)  # noqa: E231
    SCREAMING_SNAKE_CASE__         =    torch.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],]			,							device=__A						)  # noqa: E231
    SCREAMING_SNAKE_CASE__         =    torch.tensor([[[0,0,0,0],[423,237,440,251],[427,272,441,287],[419,115,437,129],[961,885,992,912],[256,38,330,58],[256,38,330,58],[336,42,353,57],[360,39,401,56],[360,39,401,56],[411,39,471,59],[479,41,528,59],[533,39,630,60],[67,113,134,131],[141,115,209,132],[68,149,133,166],[141,149,187,164],[195,148,287,165],[195,148,287,165],[195,148,287,165],[295,148,349,165],[441,149,492,166],[497,149,546,164],[64,201,125,218],[1000,1000,1000,1000]],[[0,0,0,0],[662,150,754,166],[665,199,742,211],[519,213,554,228],[519,213,554,228],[134,433,187,454],[130,467,204,480],[130,467,204,480],[130,467,204,480],[130,467,204,480],[130,467,204,480],[314,469,376,482],[504,684,582,706],[941,825,973,900],[941,825,973,900],[941,825,973,900],[941,825,973,900],[610,749,652,765],[130,659,168,672],[176,657,237,672],[238,657,312,672],[443,653,628,672],[443,653,628,672],[716,301,825,317],[1000,1000,1000,1000]]]			,							device=__A						)  # noqa: E231
    SCREAMING_SNAKE_CASE__         =    torch.tensor([[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]]			,							device=__A						)  # noqa: E231
    # these are sequence labels (i.e. at the token level)
    SCREAMING_SNAKE_CASE__         =    torch.tensor([[-100,10,10,10,9,1,-100,7,7,-100,7,7,4,2,5,2,8,8,-100,-100,5,0,3,2,-100],[-100,12,12,12,-100,12,10,-100,-100,-100,-100,10,12,9,-100,-100,-100,10,10,10,9,12,-100,10,-100]]			,							device=__A						)  # noqa: E231
    # fmt: on

    return input_ids, attention_mask, bbox, token_type_ids, labels







@require_torch
class  UpperCamelCase_      (	unittest.TestCase      ):


  """simple docstring"""



  @slow
  def   _UpperCamelCase     (						self     :					List[str]						)       ->			List[Any]:
      SCREAMING_SNAKE_CASE__         =    LayoutLMModel.from_pretrained("microsoft/layoutlm-base-uncased"						).to(A_						)

      SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__         =    prepare_layoutlm_batch_inputs()

      # forward pass
      SCREAMING_SNAKE_CASE__         =    model(input_ids=A_  ,	bbox=A_  ,	attention_mask=A_  ,	token_type_ids=A_						)

      # test the sequence output on [0, :3, :3]
      SCREAMING_SNAKE_CASE__         =    torch.tensor(
          [[0.1_785, -0.1_947, -0.0_425], [-0.3_254, -0.2_807, 0.2_553], [-0.5_391, -0.3_322, 0.3_364]]  ,	device=A_  ,	)

      self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]  ,	A_  ,	atol=1E-3						)						)

      # test the pooled output on [1, :3]
      SCREAMING_SNAKE_CASE__         =    torch.tensor([-0.6_580, -0.0_214, 0.8_552]  ,	device=A_						)

      self.assertTrue(torch.allclose(outputs.pooler_output[1, :3]  ,	A_  ,	atol=1E-3						)						)



  @slow
  def   _UpperCamelCase     (						self     :					str						)       ->			str:
      # initialize model with randomly initialized sequence classification head
      SCREAMING_SNAKE_CASE__         =    LayoutLMForSequenceClassification.from_pretrained("microsoft/layoutlm-base-uncased"  ,	num_labels=2						).to(
          A_						)

      SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__         =    prepare_layoutlm_batch_inputs()

      # forward pass
      SCREAMING_SNAKE_CASE__         =    model(
          input_ids=A_  ,	bbox=A_  ,	attention_mask=A_  ,	token_type_ids=A_  ,	labels=torch.tensor([1, 1]  ,	device=A_						)  ,	)

      # test whether we get a loss as a scalar
      SCREAMING_SNAKE_CASE__         =    outputs.loss
      SCREAMING_SNAKE_CASE__         =    torch.Size([]						)
      self.assertEqual(loss.shape  ,	A_						)

      # test the shape of the logits
      SCREAMING_SNAKE_CASE__         =    outputs.logits
      SCREAMING_SNAKE_CASE__         =    torch.Size((2, 2)						)
      self.assertEqual(logits.shape  ,	A_						)



  @slow
  def   _UpperCamelCase     (						self     :					List[str]						)       ->			List[Any]:
      # initialize model with randomly initialized token classification head
      SCREAMING_SNAKE_CASE__         =    LayoutLMForTokenClassification.from_pretrained("microsoft/layoutlm-base-uncased"  ,	num_labels=13						).to(
          A_						)

      SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__         =    prepare_layoutlm_batch_inputs()

      # forward pass
      SCREAMING_SNAKE_CASE__         =    model(
          input_ids=A_  ,	bbox=A_  ,	attention_mask=A_  ,	token_type_ids=A_  ,	labels=A_						)

      # test the loss calculation to be around 2.65
      # expected_loss = torch.tensor(2.65, device=torch_device)

      # The loss is currently somewhat random and can vary between 0.1-0.3 atol.
      # self.assertTrue(torch.allclose(outputs.loss, expected_loss, atol=0.1))

      # test the shape of the logits
      SCREAMING_SNAKE_CASE__         =    outputs.logits
      SCREAMING_SNAKE_CASE__         =    torch.Size((2, 25, 13)						)
      self.assertEqual(logits.shape  ,	A_						)







  @slow
  def   _UpperCamelCase     (						self     :					Dict						)       ->			List[Any]:
      # initialize model with randomly initialized token classification head
      SCREAMING_SNAKE_CASE__         =    LayoutLMForQuestionAnswering.from_pretrained("microsoft/layoutlm-base-uncased"						).to(A_						)

      SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__         =    prepare_layoutlm_batch_inputs()

      # forward pass
      SCREAMING_SNAKE_CASE__         =    model(input_ids=A_  ,	bbox=A_  ,	attention_mask=A_  ,	token_type_ids=A_						)

      # test the shape of the logits
      SCREAMING_SNAKE_CASE__         =    torch.Size((2, 25)						)
      self.assertEqual(outputs.start_logits.shape  ,	A_						)
      self.assertEqual(outputs.end_logits.shape  ,	A_						)


