import collections
import inspect
import unittest

from transformers import DonutSwinConfig
from transformers.testing_utils import require_torch, slow, torch_device
from transformers.utils import is_torch_available

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


if is_torch_available():
				import torch
				from torch import nn

				from transformers import DonutSwinModel
				from transformers.models.donut.modeling_donut_swin import DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST





class 					__snake_case    :




					def __init__(  self       ,			__UpperCamelCase       ,			__UpperCamelCase=13       ,			__UpperCamelCase=32       ,			__UpperCamelCase=2       ,			__UpperCamelCase=3       ,			__UpperCamelCase=16       ,			__UpperCamelCase=[1, 2, 1]       ,			__UpperCamelCase=[2, 2, 4]       ,			__UpperCamelCase=2       ,			__UpperCamelCase=2.0       ,			__UpperCamelCase=True       ,			__UpperCamelCase=0.0       ,			__UpperCamelCase=0.0       ,			__UpperCamelCase=0.1       ,			__UpperCamelCase="gelu"       ,			__UpperCamelCase=False       ,			__UpperCamelCase=True       ,			__UpperCamelCase=0.0_2       ,			__UpperCamelCase=1E-5       ,			__UpperCamelCase=True       ,			__UpperCamelCase=None       ,			__UpperCamelCase=True       ,			__UpperCamelCase=10       ,			__UpperCamelCase=8       ,			)	->    Dict:



									'''simple docstring'''







									snake_case__   :      Tuple   				= parent
									snake_case__   :      Optional[int]   				= batch_size
									snake_case__   :      Optional[int]   				= image_size
									snake_case__   :      Dict   				= patch_size
									snake_case__   :      List[Any]   				= num_channels
									snake_case__   :      List[str]   				= embed_dim
									snake_case__   :      Dict   				= depths
									snake_case__   :      Any   				= num_heads
									snake_case__   :      List[str]   				= window_size
									snake_case__   :      Tuple   				= mlp_ratio
									snake_case__   :      List[Any]   				= qkv_bias
									snake_case__   :      Optional[int]   				= hidden_dropout_prob
									snake_case__   :      int   				= attention_probs_dropout_prob
									snake_case__   :      Union[str, Any]   				= drop_path_rate
									snake_case__   :      Tuple   				= hidden_act
									snake_case__   :      Optional[int]   				= use_absolute_embeddings
									snake_case__   :      Optional[int]   				= patch_norm
									snake_case__   :      List[Any]   				= layer_norm_eps
									snake_case__   :      Any   				= initializer_range
									snake_case__   :      Dict   				= is_training
									snake_case__   :      int   				= scope
									snake_case__   :      Any   				= use_labels
									snake_case__   :      Any   				= type_sequence_label_size
									snake_case__   :      Optional[Any]   				= encoder_stride




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



									'''simple docstring'''







									snake_case__   :      str   				= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

									snake_case__   :      int   				= None
									if self.use_labels:
													snake_case__   :      Union[str, Any]   				= ids_tensor([self.batch_size]       ,			self.type_sequence_label_size						)

									snake_case__   :      Optional[Any]   				= self.get_config()

									return config, pixel_values, labels




					def      __a     (  self						)	->    int:



									'''simple docstring'''







									return DonutSwinConfig(
									    image_size=self.image_size       ,			patch_size=self.patch_size       ,			num_channels=self.num_channels       ,			embed_dim=self.embed_dim       ,			depths=self.depths       ,			num_heads=self.num_heads       ,			window_size=self.window_size       ,			mlp_ratio=self.mlp_ratio       ,			qkv_bias=self.qkv_bias       ,			hidden_dropout_prob=self.hidden_dropout_prob       ,			attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,			drop_path_rate=self.drop_path_rate       ,			hidden_act=self.hidden_act       ,			use_absolute_embeddings=self.use_absolute_embeddings       ,			path_norm=self.patch_norm       ,			layer_norm_eps=self.layer_norm_eps       ,			initializer_range=self.initializer_range       ,			encoder_stride=self.encoder_stride       ,			)




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



									'''simple docstring'''







									snake_case__   :      Any   				= DonutSwinModel(config=__UpperCamelCase						)
									model.to(__UpperCamelCase						)
									model.eval()
									snake_case__   :      Optional[Any]   				= model(__UpperCamelCase						)

									snake_case__   :      List[str]   				= ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths						) - 1))
									snake_case__   :      str   				= int(config.embed_dim * 2 ** (len(config.depths						) - 1)						)

									self.parent.assertEqual(result.last_hidden_state.shape       ,			(self.batch_size, expected_seq_len, expected_dim)						)




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



									'''simple docstring'''







									snake_case__   :      Optional[Any]   				= self.prepare_config_and_inputs()
									(
									    (
									    snake_case__
									)				,    (
									    snake_case__
									)				,    (
									    snake_case__
									)				,    
									)   :      Optional[Any]   				= config_and_inputs
									snake_case__   :      Tuple   				= {'pixel_values': pixel_values}
									return config, inputs_dict





@require_torch
class 					__snake_case    (     _lowerCamelCase       ,_lowerCamelCase       ,unittest.TestCase      ):
					__lowerCamelCase			   =  (DonutSwinModel,) if is_torch_available() else ()
					__lowerCamelCase			   =  {"""feature-extraction""": DonutSwinModel} if is_torch_available() else {}
					__lowerCamelCase			   =  True

					__lowerCamelCase			   =  False
					__lowerCamelCase			   =  False
					__lowerCamelCase			   =  False




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



									'''simple docstring'''







									snake_case__   :      Tuple   				= DonutSwinModelTester(self						)
									snake_case__   :      int   				= ConfigTester(self       ,			config_class=__UpperCamelCase       ,			embed_dim=37						)




					def      __a     (  self						)	->    str:



									'''simple docstring'''







									self.create_and_test_config_common_properties()
									self.config_tester.create_and_test_config_to_json_string()
									self.config_tester.create_and_test_config_to_json_file()
									self.config_tester.create_and_test_config_from_and_save_pretrained()
									self.config_tester.create_and_test_config_with_num_labels()
									self.config_tester.check_config_can_be_init_without_params()
									self.config_tester.check_config_arguments_init()




					def      __a     (  self						)	->    int:



									'''simple docstring'''







									return




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



									'''simple docstring'''







									snake_case__   :      str   				= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*__UpperCamelCase						)




					def      __a     (  self						)	->    Tuple:



									'''simple docstring'''







									pass




					def      __a     (  self						)	->    Tuple:



									'''simple docstring'''







									snake_case__				,    snake_case__   :      Optional[Any]   				= self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
													snake_case__   :      List[Any]   				= model_class(__UpperCamelCase						)
													self.assertIsInstance(model.get_input_embeddings()       ,			(nn.Module)						)
													snake_case__   :      Any   				= model.get_output_embeddings()
													self.assertTrue(x is None or isinstance(__UpperCamelCase       ,			nn.Linear						)						)




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



									'''simple docstring'''







									snake_case__				,    snake_case__   :      List[str]   				= self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
													snake_case__   :      Any   				= model_class(__UpperCamelCase						)
													snake_case__   :      Optional[int]   				= inspect.signature(model.forward						)
													# signature.parameters is an OrderedDict => so arg_names order is deterministic
													snake_case__   :      int   				= [*signature.parameters.keys()]

													snake_case__   :      str   				= ['pixel_values']
													self.assertListEqual(arg_names[:1]       ,			__UpperCamelCase						)




					def      __a     (  self						)	->    Any:



									'''simple docstring'''







									snake_case__				,    snake_case__   :      Union[str, Any]   				= self.model_tester.prepare_config_and_inputs_for_common()
									snake_case__   :      List[str]   				= True

									for model_class in self.all_model_classes:
													snake_case__   :      int   				= True
													snake_case__   :      List[str]   				= False
													snake_case__   :      Optional[int]   				= True
													snake_case__   :      int   				= model_class(__UpperCamelCase						)
													model.to(__UpperCamelCase						)
													model.eval()
													with torch.no_grad():
																	snake_case__   :      str   				= model(**self._prepare_for_class(__UpperCamelCase       ,			__UpperCamelCase						)						)
													snake_case__   :      Any   				= outputs.attentions
													snake_case__   :      Optional[int]   				= len(self.model_tester.depths						)
													self.assertEqual(len(__UpperCamelCase						)       ,			__UpperCamelCase						)

													# check that output_attentions also work using config
													del inputs_dict["output_attentions"]
													snake_case__   :      str   				= True
													snake_case__   :      List[str]   				= config.window_size**2
													snake_case__   :      Optional[Any]   				= model_class(__UpperCamelCase						)
													model.to(__UpperCamelCase						)
													model.eval()
													with torch.no_grad():
																	snake_case__   :      List[str]   				= model(**self._prepare_for_class(__UpperCamelCase       ,			__UpperCamelCase						)						)
													snake_case__   :      Optional[int]   				= outputs.attentions
													self.assertEqual(len(__UpperCamelCase						)       ,			__UpperCamelCase						)

													self.assertListEqual(
													    list(attentions[0].shape[-3:]						)       ,			[self.model_tester.num_heads[0], window_size_squared, window_size_squared]       ,			)
													snake_case__   :      Union[str, Any]   				= len(__UpperCamelCase						)

													# Check attention is always last and order is fine
													snake_case__   :      List[str]   				= True
													snake_case__   :      str   				= True
													snake_case__   :      Dict   				= model_class(__UpperCamelCase						)
													model.to(__UpperCamelCase						)
													model.eval()
													with torch.no_grad():
																	snake_case__   :      Optional[int]   				= model(**self._prepare_for_class(__UpperCamelCase       ,			__UpperCamelCase						)						)

													if hasattr(self.model_tester       ,			'num_hidden_states_types'						):
																	snake_case__   :      Tuple   				= self.model_tester.num_hidden_states_types
													else:
																	# also another +1 for reshaped_hidden_states
																	snake_case__   :      Optional[Any]   				= 2
													self.assertEqual(out_len + added_hidden_states       ,			len(__UpperCamelCase						)						)

													snake_case__   :      Union[str, Any]   				= outputs.attentions

													self.assertEqual(len(__UpperCamelCase						)       ,			__UpperCamelCase						)

													self.assertListEqual(
													    list(self_attentions[0].shape[-3:]						)       ,			[self.model_tester.num_heads[0], window_size_squared, window_size_squared]       ,			)




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



									'''simple docstring'''







									snake_case__   :      Dict   				= model_class(__UpperCamelCase						)
									model.to(__UpperCamelCase						)
									model.eval()

									with torch.no_grad():
													snake_case__   :      Tuple   				= model(**self._prepare_for_class(__UpperCamelCase       ,			__UpperCamelCase						)						)

									snake_case__   :      List[str]   				= outputs.hidden_states

									snake_case__   :      Any   				= getattr(
									    self.model_tester       ,			'expected_num_hidden_layers'       ,			len(self.model_tester.depths						) + 1						)
									self.assertEqual(len(__UpperCamelCase						)       ,			__UpperCamelCase						)

									# DonutSwin has a different seq_length
									snake_case__   :      List[str]   				= (
									    config.patch_size
									    if isinstance(config.patch_size       ,			collections.abc.Iterable						)
									    else (config.patch_size, config.patch_size)
									)

									snake_case__   :      Optional[Any]   				= (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

									self.assertListEqual(
									    list(hidden_states[0].shape[-2:]						)       ,			[num_patches, self.model_tester.embed_dim]       ,			)

									snake_case__   :      int   				= outputs.reshaped_hidden_states
									self.assertEqual(len(__UpperCamelCase						)       ,			__UpperCamelCase						)

									snake_case__				,    snake_case__				,    snake_case__				,    snake_case__   :      Dict   				= reshaped_hidden_states[0].shape
									snake_case__   :      int   				= (
									    reshaped_hidden_states[0].view(__UpperCamelCase       ,			__UpperCamelCase       ,			height * width						).permute(0       ,			2       ,			1						)
									)
									self.assertListEqual(
									    list(reshaped_hidden_states.shape[-2:]						)       ,			[num_patches, self.model_tester.embed_dim]       ,			)




					def      __a     (  self						)	->    List[str]:



									'''simple docstring'''







									snake_case__				,    snake_case__   :      str   				= self.model_tester.prepare_config_and_inputs_for_common()

									snake_case__   :      Optional[int]   				= (
									    self.model_tester.image_size
									    if isinstance(self.model_tester.image_size       ,			collections.abc.Iterable						)
									    else (self.model_tester.image_size, self.model_tester.image_size)
									)

									for model_class in self.all_model_classes:
													snake_case__   :      Optional[Any]   				= True
													self.check_hidden_states_output(__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase						)

													# check that output_hidden_states also work using config
													del inputs_dict["output_hidden_states"]
													snake_case__   :      List[Any]   				= True

													self.check_hidden_states_output(__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase						)




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



									'''simple docstring'''







									snake_case__				,    snake_case__   :      int   				= self.model_tester.prepare_config_and_inputs_for_common()
									snake_case__   :      Any   				= 3

									snake_case__   :      Dict   				= (
									    self.model_tester.image_size
									    if isinstance(self.model_tester.image_size       ,			collections.abc.Iterable						)
									    else (self.model_tester.image_size, self.model_tester.image_size)
									)
									snake_case__   :      Optional[int]   				= (
									    config.patch_size
									    if isinstance(config.patch_size       ,			collections.abc.Iterable						)
									    else (config.patch_size, config.patch_size)
									)

									snake_case__   :      Optional[int]   				= image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
									snake_case__   :      List[str]   				= image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

									for model_class in self.all_model_classes:
													snake_case__   :      str   				= True
													self.check_hidden_states_output(__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			(padded_height, padded_width)						)

													# check that output_hidden_states also work using config
													del inputs_dict["output_hidden_states"]
													snake_case__   :      List[str]   				= True
													self.check_hidden_states_output(__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			(padded_height, padded_width)						)




					@slow
					def      __a     (  self						)	->    Optional[Any]:



									'''simple docstring'''







									for model_name in DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													snake_case__   :      str   				= DonutSwinModel.from_pretrained(__UpperCamelCase						)
													self.assertIsNotNone(__UpperCamelCase						)




					def      __a     (  self						)	->    Dict:



									'''simple docstring'''







									snake_case__				,    snake_case__   :      List[str]   				= self.model_tester.prepare_config_and_inputs_for_common()

									snake_case__   :      List[Any]   				= _config_zero_init(__UpperCamelCase						)
									for model_class in self.all_model_classes:
													snake_case__   :      Union[str, Any]   				= model_class(config=__UpperCamelCase						)
													for name, param in model.named_parameters():
																	if "embeddings" not in name and param.requires_grad:
																					self.assertIn(
																					    ((param.data.mean() * 1E9).round() / 1E9).item()       ,			[0.0, 1.0]       ,			msg=F"""Parameter {name} of model {model_class} seems not properly initialized"""       ,			)



import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionDualGuidedPipeline
from diffusers.utils.testing_utils import load_image, nightly, require_torch_gpu, torch_device


lowerCAmelCase__	:  Union[str, Any]						=  False





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




					def      __a     (  self						)	->    str:



									'''simple docstring'''







									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()




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



									'''simple docstring'''







									snake_case__   :      str   				= VersatileDiffusionDualGuidedPipeline.from_pretrained('shi-labs/versatile-diffusion'						)
									# remove text_unet
									pipe.remove_unused_weights()
									pipe.to(__UpperCamelCase						)
									pipe.set_progress_bar_config(disable=__UpperCamelCase						)

									snake_case__   :      Union[str, Any]   				= load_image(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'						)

									snake_case__   :      Dict   				= torch.manual_seed(0						)
									snake_case__   :      Tuple   				= pipe(
									    prompt='first prompt'       ,			image=__UpperCamelCase       ,			text_to_image_strength=0.7_5       ,			generator=__UpperCamelCase       ,			guidance_scale=7.5       ,			num_inference_steps=2       ,			output_type='numpy'       ,			).images

									with tempfile.TemporaryDirectory() as tmpdirname:
													pipe.save_pretrained(__UpperCamelCase						)
													snake_case__   :      List[Any]   				= VersatileDiffusionDualGuidedPipeline.from_pretrained(__UpperCamelCase						)

									pipe.to(__UpperCamelCase						)
									pipe.set_progress_bar_config(disable=__UpperCamelCase						)

									snake_case__   :      str   				= generator.manual_seed(0						)
									snake_case__   :      List[str]   				= pipe(
									    prompt='first prompt'       ,			image=__UpperCamelCase       ,			text_to_image_strength=0.7_5       ,			generator=__UpperCamelCase       ,			guidance_scale=7.5       ,			num_inference_steps=2       ,			output_type='numpy'       ,			).images

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




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



									'''simple docstring'''







									snake_case__   :      Union[str, Any]   				= VersatileDiffusionDualGuidedPipeline.from_pretrained('shi-labs/versatile-diffusion'						)
									pipe.remove_unused_weights()
									pipe.to(__UpperCamelCase						)
									pipe.set_progress_bar_config(disable=__UpperCamelCase						)

									snake_case__   :      str   				= 'cyberpunk 2077'
									snake_case__   :      Dict   				= load_image(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'						)
									snake_case__   :      Union[str, Any]   				= torch.manual_seed(0						)
									snake_case__   :      Union[str, Any]   				= pipe(
									    prompt=__UpperCamelCase       ,			image=__UpperCamelCase       ,			text_to_image_strength=0.7_5       ,			generator=__UpperCamelCase       ,			guidance_scale=7.5       ,			num_inference_steps=50       ,			output_type='numpy'       ,			).images

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

									assert image.shape == (1, 512, 512, 3)
									snake_case__   :      List[str]   				= np.array([0.0_7_8_7, 0.0_8_4_9, 0.0_8_2_6, 0.0_8_1_2, 0.0_8_0_7, 0.0_7_9_5, 0.0_8_1_8, 0.0_7_9_8, 0.0_7_7_9]						)

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


