
"""simple docstring"""


import inspect
import unittest

import requests

from transformers import MgpstrConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

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


if is_torch_available():
						import torch
						from torch import nn

						from transformers import MgpstrForSceneTextRecognition


if is_vision_available():
						from PIL import Image

						from transformers import MgpstrProcessor




class    __lowerCamelCase     :




				def __init__(							self       ,	UpperCAmelCase       ,	UpperCAmelCase=False       ,	UpperCAmelCase=13       ,	UpperCAmelCase=(32, 128)       ,	UpperCAmelCase=4       ,	UpperCAmelCase=3       ,	UpperCAmelCase=27       ,	UpperCAmelCase=38       ,	UpperCAmelCase=99       ,	UpperCAmelCase=99       ,	UpperCAmelCase=32       ,	UpperCAmelCase=5       ,	UpperCAmelCase=4       ,	UpperCAmelCase=4.0       ,	UpperCAmelCase=257       ,	UpperCAmelCase=None       ,	):
							lowerCamelCase_ 		=     parent
							lowerCamelCase_ 		=     is_training
							lowerCamelCase_ 		=     batch_size
							lowerCamelCase_ 		=     image_size
							lowerCamelCase_ 		=     patch_size
							lowerCamelCase_ 		=     num_channels
							lowerCamelCase_ 		=     max_token_length
							lowerCamelCase_ 		=     num_character_labels
							lowerCamelCase_ 		=     num_bpe_labels
							lowerCamelCase_ 		=     num_wordpiece_labels
							lowerCamelCase_ 		=     hidden_size
							lowerCamelCase_ 		=     num_hidden_layers
							lowerCamelCase_ 		=     num_attention_heads
							lowerCamelCase_ 		=     mlp_ratio
							lowerCamelCase_ 		=     patch_embeds_hidden_size
							lowerCamelCase_ 		=     output_hidden_states




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     floats_tensor([self.batch_size, self.num_channels, self.image_size[0], self.image_size[1]]      )
							lowerCamelCase_ 		=     self.get_config()
							return config, pixel_values




				def 					UpperCAmelCase__	(							self      ):
							return MgpstrConfig(
							    image_size=self.image_size       ,	patch_size=self.patch_size       ,	num_channels=self.num_channels       ,	max_token_length=self.max_token_length       ,	num_character_labels=self.num_character_labels       ,	num_bpe_labels=self.num_bpe_labels       ,	num_wordpiece_labels=self.num_wordpiece_labels       ,	hidden_size=self.hidden_size       ,	num_hidden_layers=self.num_hidden_layers       ,	num_attention_heads=self.num_attention_heads       ,	mlp_ratio=self.mlp_ratio       ,	output_hidden_states=self.output_hidden_states       ,	)




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     MgpstrForSceneTextRecognition(UpperCAmelCase      )
							model.to(UpperCAmelCase      )
							model.eval()
							with torch.no_grad():
										lowerCamelCase_ 		=     model(UpperCAmelCase      )
							self.parent.assertEqual(
							    generated_ids[0][0].shape       ,	(self.batch_size, self.max_token_length, self.num_character_labels)      )





				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     self.prepare_config_and_inputs()
							lowerCamelCase_		,   lowerCamelCase_ 		=     config_and_inputs
							lowerCamelCase_ 		=     {'''pixel_values''': pixel_values}
							return config, inputs_dict




@require_torch
class    __lowerCamelCase     (   lowerCAmelCase	,  lowerCAmelCase	,  unittest.TestCase ):
				a__:			Dict									=						(MgpstrForSceneTextRecognition,) if is_torch_available() else ()
				a__:			Optional[int]									=						{'feature-extraction': MgpstrForSceneTextRecognition} if is_torch_available() else {}
				a__:			int									=						False

				a__:			Any									=						False
				a__:			Union[str, Any]									=						False
				a__:			str									=						False
				a__:			List[str]									=						False




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     MgpstrModelTester(self      )
							lowerCamelCase_ 		=     ConfigTester(self       ,	config_class=UpperCAmelCase       ,	has_text_modality=UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							self.config_tester.run_common_tests()




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*UpperCAmelCase      )




				@unittest.skip(reason='''MgpstrModel does not use inputs_embeds'''      )
				def 					UpperCAmelCase__	(							self      ):
							pass




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
										lowerCamelCase_ 		=     model_class(UpperCAmelCase      )
										self.assertIsInstance(model.get_input_embeddings()       ,	(nn.Module)      )
										lowerCamelCase_ 		=     model.get_output_embeddings()
										self.assertTrue(x is None or isinstance(UpperCAmelCase       ,	nn.Linear      )      )




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs_for_common()

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

										lowerCamelCase_ 		=     ['''pixel_values''']
										self.assertListEqual(arg_names[:1]       ,	UpperCAmelCase      )




				@unittest.skip(reason='''MgpstrModel does not support feedforward chunking'''      )
				def 					UpperCAmelCase__	(							self      ):
							pass




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
										if not model_class.supports_gradient_checkpointing:
													continue

										lowerCamelCase_ 		=     True
										lowerCamelCase_ 		=     model_class(UpperCAmelCase      )
										self.assertTrue(model.is_gradient_checkpointing      )




				def 					UpperCAmelCase__	(							self      ):
							def check_hidden_states_output(UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ):
										lowerCamelCase_ 		=     model_class(UpperCAmelCase      )
										model.to(UpperCAmelCase      )
										model.eval()

										with torch.no_grad():
													lowerCamelCase_ 		=     model(**self._prepare_for_class(UpperCAmelCase       ,	UpperCAmelCase      )      )

										lowerCamelCase_ 		=     outputs.hidden_states

										lowerCamelCase_ 		=     getattr(
										    self.model_tester       ,	'''expected_num_hidden_layers'''       ,	self.model_tester.num_hidden_layers + 1      )
										self.assertEqual(len(UpperCAmelCase      )       ,	UpperCAmelCase      )

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

							lowerCamelCase_		,   lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
										lowerCamelCase_ 		=     True
										check_hidden_states_output(UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      )

										# check that output_hidden_states also work using config
										del inputs_dict["output_hidden_states"]
										lowerCamelCase_ 		=     True

										check_hidden_states_output(UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs_for_common()

							lowerCamelCase_ 		=     _config_zero_init(UpperCAmelCase      )
							for model_class in self.all_model_classes:
										lowerCamelCase_ 		=     model_class(config=UpperCAmelCase      )
										for name, param in model.named_parameters():
													if isinstance(UpperCAmelCase       ,	(nn.Linear, nn.Convad, nn.LayerNorm)      ):
																if 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"       ,	)





				@unittest.skip(reason='''Retain_grad is tested in individual model tests'''      )
				def 					UpperCAmelCase__	(							self      ):
							pass






def    lowercase	(  ):
			lowerCamelCase_ 		=     '''https://i.postimg.cc/ZKwLg2Gw/367-14.png'''
			lowerCamelCase_ 		=     Image.open(requests.get(lowerCAmelCase__  ,stream=lowerCAmelCase__	).raw	).convert('''RGB'''	)
			return im




@require_vision
@require_torch
class    __lowerCamelCase     (   unittest.TestCase ):




				@slow
				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     '''alibaba-damo/mgp-str-base'''
							lowerCamelCase_ 		=     MgpstrForSceneTextRecognition.from_pretrained(UpperCAmelCase      ).to(UpperCAmelCase      )
							lowerCamelCase_ 		=     MgpstrProcessor.from_pretrained(UpperCAmelCase      )

							lowerCamelCase_ 		=     prepare_img()
							lowerCamelCase_ 		=     processor(images=UpperCAmelCase       ,	return_tensors='''pt'''      ).pixel_values.to(UpperCAmelCase      )

							# forward pass
							with torch.no_grad():
										lowerCamelCase_ 		=     model(UpperCAmelCase      )

							# verify the logits
							self.assertEqual(outputs.logits[0].shape       ,	torch.Size((1, 27, 38)      )      )

							lowerCamelCase_ 		=     processor.batch_decode(outputs.logits      )
							lowerCamelCase_ 		=     '''ticket'''

							self.assertEqual(out_strs['''generated_text'''][0]       ,	UpperCAmelCase      )

							lowerCamelCase_ 		=     torch.tensor(
							    [[[-3_9.5_3_9_7, -4_4.4_0_2_4, -3_6.1_8_4_4], [-6_1.4_7_0_9, -6_3.8_6_3_9, -5_8.3_4_5_4], [-7_4.0_2_2_5, -6_8.5_4_9_4, -7_1.2_1_6_4]]]       ,	device=UpperCAmelCase       ,	)

							self.assertTrue(torch.allclose(outputs.logits[0][:, 1:4, 1:4]       ,	UpperCAmelCase       ,	atol=1e-4      )      )





"""simple docstring"""


import tempfile
import unittest

import torch

from diffusers import UNetaDConditionModel
from diffusers.training_utils import EMAModel
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device


enable_full_determinism()




class    __lowerCamelCase     (   unittest.TestCase ):
				a__:			Union[str, Any]									=						'hf-internal-testing/tiny-stable-diffusion-pipe'
				a__:			Tuple									=						1
				a__:			int									=						77
				a__:			List[str]									=						32
				a__:			List[Any]									=						4
				a__:			str									=						64
				a__:			Optional[Any]									=						torch.manual_seed(0 )




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase=0.9_9_9_9      ):
							lowerCamelCase_ 		=     UNetaDConditionModel.from_pretrained(self.model_id       ,	subfolder='''unet'''      )
							lowerCamelCase_ 		=     unet.to(UpperCAmelCase      )
							lowerCamelCase_ 		=     EMAModel(unet.parameters()       ,	decay=UpperCAmelCase       ,	model_cls=UpperCAmelCase       ,	model_config=unet.config      )
							return unet, ema_unet




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     torch.randn(
							    self.batch_size       ,	self.num_in_channels       ,	self.latent_height       ,	self.latent_width       ,	generator=self.generator      ).to(UpperCAmelCase      )
							lowerCamelCase_ 		=     torch.randint(0       ,	1000       ,	size=(self.batch_size,)       ,	generator=self.generator      ).to(UpperCAmelCase      )
							lowerCamelCase_ 		=     torch.randn(
							    self.batch_size       ,	self.prompt_length       ,	self.text_encoder_hidden_dim       ,	generator=self.generator      ).to(UpperCAmelCase      )
							return noisy_latents, timesteps, encoder_hidden_states




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     {}
							for k, param in unet.state_dict().items():
										lowerCamelCase_ 		=     torch.randn_like(UpperCAmelCase      ) + (param * torch.randn_like(UpperCAmelCase      ))
										updated_state_dict.update({k: updated_param}      )
							unet.load_state_dict(UpperCAmelCase      )
							return unet




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models()
							# Take the first (hypothetical) EMA step.
							ema_unet.step(unet.parameters()      )
							assert ema_unet.optimization_step == 1

							# Take two more.
							for _ in range(2      ):
										ema_unet.step(unet.parameters()      )
							assert ema_unet.optimization_step == 3




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models()
							# Since the `unet` is not being updated (i.e., backprop'd)
							# there won't be any difference between the `params` of `unet`
							# and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
							ema_unet.step(unet.parameters()      )
							lowerCamelCase_ 		=     list(unet.parameters()      )
							for s_param, param in zip(ema_unet.shadow_params       ,	UpperCAmelCase      ):
										assert torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )

							# The above holds true even if we call `ema.step()` multiple times since
							# `unet` params are still not being updated.
							for _ in range(4      ):
										ema_unet.step(unet.parameters()      )
							for s_param, param in zip(ema_unet.shadow_params       ,	UpperCAmelCase      ):
										assert torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models()
							# Here we simulate the parameter updates for `unet`. Since there might
							# be some parameters which are initialized to zero we take extra care to
							# initialize their values to something non-zero before the multiplication.
							lowerCamelCase_ 		=     self.simulate_backprop(UpperCAmelCase      )

							# Take the EMA step.
							ema_unet.step(unet_pseudo_updated_step_one.parameters()      )

							# Now the EMA'd parameters won't be equal to the original model parameters.
							lowerCamelCase_ 		=     list(unet_pseudo_updated_step_one.parameters()      )
							for s_param, param in zip(ema_unet.shadow_params       ,	UpperCAmelCase      ):
										assert ~torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )

							# Ensure this is the case when we take multiple EMA steps.
							for _ in range(4      ):
										ema_unet.step(unet.parameters()      )
							for s_param, param in zip(ema_unet.shadow_params       ,	UpperCAmelCase      ):
										assert ~torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							# If we call EMA step after a backpropagation consecutively for two times,
							# the shadow params from those two steps should be different.
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models()

							# First backprop + EMA
							lowerCamelCase_ 		=     self.simulate_backprop(UpperCAmelCase      )
							ema_unet.step(unet_step_one.parameters()      )
							lowerCamelCase_ 		=     ema_unet.shadow_params

							# Second backprop + EMA
							lowerCamelCase_ 		=     self.simulate_backprop(UpperCAmelCase      )
							ema_unet.step(unet_step_two.parameters()      )
							lowerCamelCase_ 		=     ema_unet.shadow_params

							for step_one, step_two in zip(UpperCAmelCase       ,	UpperCAmelCase      ):
										assert ~torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							# If there's no decay even if there are backprops, EMA steps
							# won't take any effect i.e., the shadow params would remain the
							# same.
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models(decay=0.0      )
							lowerCamelCase_ 		=     self.simulate_backprop(UpperCAmelCase      )
							ema_unet.step(unet_step_one.parameters()      )
							lowerCamelCase_ 		=     ema_unet.shadow_params

							lowerCamelCase_ 		=     self.simulate_backprop(UpperCAmelCase      )
							ema_unet.step(unet_step_two.parameters()      )
							lowerCamelCase_ 		=     ema_unet.shadow_params

							for step_one, step_two in zip(UpperCAmelCase       ,	UpperCAmelCase      ):
										assert torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )





				@skip_mps
				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models()
							lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_dummy_inputs()

							with tempfile.TemporaryDirectory() as tmpdir:
										ema_unet.save_pretrained(UpperCAmelCase      )
										lowerCamelCase_ 		=     UNetaDConditionModel.from_pretrained(UpperCAmelCase       ,	model_cls=UpperCAmelCase      )
										lowerCamelCase_ 		=     loaded_unet.to(unet.device      )

							# Since no EMA step has been performed the outputs should match.
							lowerCamelCase_ 		=     unet(UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ).sample
							lowerCamelCase_ 		=     loaded_unet(UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ).sample

							assert torch.allclose(UpperCAmelCase       ,	UpperCAmelCase       ,	atol=1e-4      )



