






"""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    snake_case						:


   '''simple docstring'''

   def __init__(      self					:	Optional[int],			_lowerCamelCase					:	Optional[Any],			_lowerCamelCase					:	str=False,			_lowerCamelCase					:	List[Any]=13,			_lowerCamelCase					:	str=(32, 1_28),			_lowerCamelCase					:	str=4,			_lowerCamelCase					:	int=3,			_lowerCamelCase					:	Union[str, Any]=27,			_lowerCamelCase					:	Tuple=38,			_lowerCamelCase					:	Any=99,			_lowerCamelCase					:	Any=99,			_lowerCamelCase					:	List[str]=32,			_lowerCamelCase					:	int=5,			_lowerCamelCase					:	Union[str, Any]=4,			_lowerCamelCase					:	str=4.0,			_lowerCamelCase					:	Optional[Any]=2_57,			_lowerCamelCase					:	Union[str, Any]=None,			):

     '''simple docstring'''
     __A		     =					parent
     __A		     =					is_training
     __A		     =					batch_size
     __A		     =					image_size
     __A		     =					patch_size
     __A		     =					num_channels
     __A		     =					max_token_length
     __A		     =					num_character_labels
     __A		     =					num_bpe_labels
     __A		     =					num_wordpiece_labels
     __A		     =					hidden_size
     __A		     =					num_hidden_layers
     __A		     =					num_attention_heads
     __A		     =					mlp_ratio
     __A		     =					patch_embeds_hidden_size
     __A		     =					output_hidden_states


   def 							_SCREAMING_SNAKE_CASE							(      self					:	int    ):

     '''simple docstring'''
     __A		     =					floats_tensor([self.batch_size, self.num_channels, self.image_size[0], self.image_size[1]]    )
     __A		     =					self.get_config()
     return config, pixel_values


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Dict    ):

     '''simple docstring'''
     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 							_SCREAMING_SNAKE_CASE							(      self					:	Optional[Any],			_lowerCamelCase					:	List[str],			_lowerCamelCase					:	Dict    ):

     '''simple docstring'''
     __A		     =					MgpstrForSceneTextRecognition(__snake_case    )
     model.to(__snake_case    )
     model.eval()
     with torch.no_grad():
       __A		     =					model(__snake_case    )
     self.parent.assertEqual(
         generated_ids[0][0].shape,			(self.batch_size, self.max_token_length, self.num_character_labels)    )






   def 							_SCREAMING_SNAKE_CASE							(      self					:	Tuple    ):

     '''simple docstring'''
     __A		     =					self.prepare_config_and_inputs()
     __A		     =					config_and_inputs
     __A		     =					{'pixel_values': pixel_values}
     return config, inputs_dict





@require_torch
class    snake_case						(				__UpperCAmelCase  ,      __UpperCAmelCase  ,      unittest.TestCase		):


   '''simple docstring'''

   A_		:      Tuple							=					(MgpstrForSceneTextRecognition,) if is_torch_available() else ()
   A_		:      Optional[Any]							=					{"feature-extraction": MgpstrForSceneTextRecognition} if is_torch_available() else {}
   A_		:      int							=					False

   A_		:      Tuple							=					False
   A_		:      Optional[Any]							=					False
   A_		:      List[Any]							=					False
   A_		:      Any							=					False


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Union[str, Any]    ):

     '''simple docstring'''
     __A		     =					MgpstrModelTester(self    )
     __A		     =					ConfigTester(self,			config_class=__snake_case,			has_text_modality=__snake_case    )


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Tuple    ):

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


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Tuple    ):

     '''simple docstring'''
     __A		     =					self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*__snake_case    )


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

     '''simple docstring'''
     pass


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Union[str, Any]    ):

     '''simple docstring'''
     __A		     =					self.model_tester.prepare_config_and_inputs_for_common()

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


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Dict    ):

     '''simple docstring'''
     __A		     =					self.model_tester.prepare_config_and_inputs_for_common()

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

       __A		     =					['pixel_values']
       self.assertListEqual(arg_names[:1],			__snake_case    )


   @unittest.skip(reason='''MgpstrModel does not support feedforward chunking'''    )
   def 							_SCREAMING_SNAKE_CASE							(      self					:	List[Any]    ):

     '''simple docstring'''
     pass


   def 							_SCREAMING_SNAKE_CASE							(      self					:	str    ):

     '''simple docstring'''
     __A		     =					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

       __A		     =					True
       __A		     =					model_class(__snake_case    )
       self.assertTrue(model.is_gradient_checkpointing    )


   def 							_SCREAMING_SNAKE_CASE							(      self					:	List[Any]    ):

     '''simple docstring'''





     def check_hidden_states_output(_lowerCamelCase					:	Optional[int],			_lowerCamelCase					:	int,			_lowerCamelCase					:	Tuple    ):
       __A		     =					model_class(__snake_case    )
       model.to(__snake_case    )
       model.eval()

       with torch.no_grad():
         __A		     =					model(**self._prepare_for_class(__snake_case,			__snake_case    )    )

       __A		     =					outputs.hidden_states

       __A		     =					getattr(
           self.model_tester,			'''expected_num_hidden_layers''',			self.model_tester.num_hidden_layers + 1    )
       self.assertEqual(len(__snake_case    ),			__snake_case    )

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

     __A		     =					self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
       __A		     =					True
       check_hidden_states_output(__snake_case,			__snake_case,			__snake_case    )

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

       check_hidden_states_output(__snake_case,			__snake_case,			__snake_case    )


   def 							_SCREAMING_SNAKE_CASE							(      self					:	int    ):

     '''simple docstring'''
     __A		     =					self.model_tester.prepare_config_and_inputs_for_common()

     __A		     =					_config_zero_init(__snake_case    )
     for model_class in self.all_model_classes:
       __A		     =					model_class(config=__snake_case    )
       for name, param in model.named_parameters():
         if isinstance(__snake_case,			(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 							_SCREAMING_SNAKE_CASE							(      self					:	List[str]    ):

     '''simple docstring'''
     pass





def    lowerCAmelCase							(       ):




  """simple docstring"""







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





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


   '''simple docstring'''

   @slow
   def 							_SCREAMING_SNAKE_CASE							(      self					:	List[str]    ):

     '''simple docstring'''
     __A		     =					'alibaba-damo/mgp-str-base'
     __A		     =					MgpstrForSceneTextRecognition.from_pretrained(__snake_case    ).to(__snake_case    )
     __A		     =					MgpstrProcessor.from_pretrained(__snake_case    )

     __A		     =					prepare_img()
     __A		     =					processor(images=__snake_case,			return_tensors='''pt'''    ).pixel_values.to(__snake_case    )

     # forward pass
     with torch.no_grad():
       __A		     =					model(__snake_case    )

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

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

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

     __A		     =					torch.tensor(
         [[[-39.53_97, -44.40_24, -36.18_44], [-61.47_09, -63.86_39, -58.34_54], [-74.02_25, -68.54_94, -71.21_64]]],			device=__snake_case,			)

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











'''simple docstring'''






import tempfile
import unittest

import numpy as np
import torch

from diffusers import DiffusionPipeline
from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor







class       lowerCamelCase       (    unittest.TestCase	):


					def  SCREAMING_SNAKE_CASE_		(       self						:	List[Any]							,  __snake_case						:	bool = False							):





											'''simple docstring'''





											_snake_case: Any	       =		10

											if only_cross_attention:
																	_snake_case: Optional[int]	       =		12
											else:
																	# when only cross attention is not set, the cross attention dim must be the same as the query dim
																	_snake_case: Optional[int]	       =		query_dim

											return {
											    "query_dim": query_dim,
											    "cross_attention_dim": cross_attention_dim,
											    "heads": 2,
											    "dim_head": 4,
											    "added_kv_proj_dim": 6,
											    "norm_num_groups": 1,
											    "only_cross_attention": only_cross_attention,
											    "processor": AttnAddedKVProcessor(),
											}


					def  SCREAMING_SNAKE_CASE_		(       self						:	List[Any]							,  __snake_case						:	List[Any]							,  __snake_case						:	Tuple							):





											'''simple docstring'''





											_snake_case: Optional[int]	       =		2

											_snake_case: Union[str, Any]	       =		torch.rand(__snake_case							,  __snake_case							,  3							,  2							)
											_snake_case: Optional[int]	       =		torch.rand(__snake_case							,  4							,  __snake_case							)
											_snake_case: int	       =		None

											return {
											    "hidden_states": hidden_states,
											    "encoder_hidden_states": encoder_hidden_states,
											    "attention_mask": attention_mask,
											}




					def  SCREAMING_SNAKE_CASE_		(       self						:	List[Any]							):





											'''simple docstring'''





											torch.manual_seed(0							)

											_snake_case: int	       =		self.get_constructor_arguments(only_cross_attention=__snake_case							)
											_snake_case: Union[str, Any]	       =		Attention(**__snake_case							)

											self.assertTrue(attn.to_k is not None							)
											self.assertTrue(attn.to_v is not None							)

											_snake_case: Optional[int]	       =		self.get_forward_arguments(
											    query_dim=constructor_args['query_dim']							,  added_kv_proj_dim=constructor_args['added_kv_proj_dim']							)

											_snake_case: str	       =		attn(**__snake_case							)

											# only self attention

											torch.manual_seed(0							)

											_snake_case: str	       =		self.get_constructor_arguments(only_cross_attention=__snake_case							)
											_snake_case: Optional[Any]	       =		Attention(**__snake_case							)

											self.assertTrue(attn.to_k is None							)
											self.assertTrue(attn.to_v is None							)

											_snake_case: Dict	       =		self.get_forward_arguments(
											    query_dim=constructor_args['query_dim']							,  added_kv_proj_dim=constructor_args['added_kv_proj_dim']							)

											_snake_case: str	       =		attn(**__snake_case							)

											self.assertTrue((only_cross_attn_out != self_and_cross_attn_out).all()							)







class       lowerCamelCase       (    unittest.TestCase	):


					def  SCREAMING_SNAKE_CASE_		(       self						:	Tuple							):





											'''simple docstring'''





											_snake_case: str	       =		DiffusionPipeline.from_pretrained('hf-internal-testing/tiny-stable-diffusion-pipe'							,  safety_checker=__snake_case							)

											_snake_case: int	       =		pipe(
											    'foo'							,  num_inference_steps=2							,  generator=torch.Generator('cpu'							).manual_seed(0							)							,  output_type='np'							,  ).images

											# the initial conversion succeeds
											_snake_case: Dict	       =		DiffusionPipeline.from_pretrained(
											    'hf-internal-testing/tiny-stable-diffusion-pipe'							,  device_map='sequential'							,  safety_checker=__snake_case							)

											_snake_case: Any	       =		pipe(
											    'foo'							,  num_inference_steps=2							,  generator=torch.Generator('cpu'							).manual_seed(0							)							,  output_type='np'							,  ).images

											with tempfile.TemporaryDirectory() as tmpdir:
																	# save the converted model
																	pipe.save_pretrained(__snake_case							)

																	# can also load the converted weights
																	_snake_case: Union[str, Any]	       =		DiffusionPipeline.from_pretrained(__snake_case							,  device_map='sequential'							,  safety_checker=__snake_case							)

											_snake_case: Union[str, Any]	       =		pipe(
											    'foo'							,  num_inference_steps=2							,  generator=torch.Generator('cpu'							).manual_seed(0							)							,  output_type='np'							,  ).images

											self.assertTrue(np.allclose(__snake_case							,  __snake_case							,  atol=1e-5							)							)
											self.assertTrue(np.allclose(__snake_case							,  __snake_case							,  atol=1e-5							)							)





