






import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


_SCREAMING_SNAKE_CASE   =  False






class 					SCREAMING_SNAKE_CASE_  (       unittest.TestCase ):
       """simple docstring"""


       def 				UpperCamelCase__						(				self	:Tuple):



          """simple docstring"""

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





       @property
       def 				UpperCamelCase__						(				self	:List[str]):



          """simple docstring"""

          return 12





       @property
       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          return 12





       @property
       def 				UpperCamelCase__						(				self	:List[Any]):



          """simple docstring"""

          return 32





       @property
       def 				UpperCamelCase__						(				self	:Any):



          """simple docstring"""

          torch.manual_seed(0)
          _lowercase              =VQModel(
              block_out_channels=[32, 64],      in_channels=3,      out_channels=3,      down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D'],      up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D'],      latent_channels=3,      num_vq_embeddings=self.num_embed,      vq_embed_dim=3,      )
          return model





       @property
       def 				UpperCamelCase__						(				self	:List[Any]):



          """simple docstring"""

          _lowercase              =CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')
          return tokenizer





       @property
       def 				UpperCamelCase__						(				self	:Tuple):



          """simple docstring"""

          torch.manual_seed(0)
          _lowercase              =CLIPTextConfig(
              bos_token_id=0,      eos_token_id=2,      hidden_size=self.text_embedder_hidden_size,      intermediate_size=37,      layer_norm_eps=1e-0_5,      num_attention_heads=4,      num_hidden_layers=5,      pad_token_id=1,      vocab_size=1000,      )
          return CLIPTextModel(snake_case)





       @property
       def 				UpperCamelCase__						(				self	:List[Any]):



          """simple docstring"""

          torch.manual_seed(0)

          _lowercase              =12
          _lowercase              =12

          _lowercase              ={
              'attention_bias': True,
              'cross_attention_dim': 32,
              'attention_head_dim': height * width,
              'num_attention_heads': 1,
              'num_vector_embeds': self.num_embed,
              'num_embeds_ada_norm': self.num_embeds_ada_norm,
              'norm_num_groups': 32,
              'sample_size': width,
              'activation_fn': 'geglu-approximate',
          }

          _lowercase              =TransformeraDModel(**snake_case)
          return model





       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

          _lowercase              ='cpu'

          _lowercase              =self.dummy_vqvae
          _lowercase              =self.dummy_text_encoder
          _lowercase              =self.dummy_tokenizer
          _lowercase              =self.dummy_transformer
          _lowercase              =VQDiffusionScheduler(self.num_embed)
          _lowercase              =LearnedClassifierFreeSamplingEmbeddings(learnable=snake_case)

          _lowercase              =VQDiffusionPipeline(
              vqvae=snake_case,      text_encoder=snake_case,      tokenizer=snake_case,      transformer=snake_case,      scheduler=snake_case,      learned_classifier_free_sampling_embeddings=snake_case,      )
          _lowercase              =pipe.to(snake_case)
          pipe.set_progress_bar_config(disable=snake_case)

          _lowercase              ='teddy bear playing in the pool'

          _lowercase              =torch.Generator(device=snake_case).manual_seed(0)
          _lowercase              =pipe([prompt],      generator=snake_case,      num_inference_steps=2,      output_type='np')
          _lowercase              =output.images

          _lowercase              =torch.Generator(device=snake_case).manual_seed(0)
          _lowercase              =pipe(
              [prompt],      generator=snake_case,      output_type='np',      return_dict=snake_case,      num_inference_steps=2)[0]

          _lowercase              =image[0, -3:, -3:, -1]
          _lowercase              =image_from_tuple[0, -3:, -3:, -1]

          assert image.shape == (1, 24, 24, 3)

          _lowercase              =np.array([0.6_5_5_1, 0.6_1_6_8, 0.5_0_0_8, 0.5_6_7_6, 0.5_6_5_9, 0.4_2_9_5, 0.6_0_7_3, 0.5_5_9_9, 0.4_9_9_2])

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





       def 				UpperCamelCase__						(				self	:List[str]):



          """simple docstring"""

          _lowercase              ='cpu'

          _lowercase              =self.dummy_vqvae
          _lowercase              =self.dummy_text_encoder
          _lowercase              =self.dummy_tokenizer
          _lowercase              =self.dummy_transformer
          _lowercase              =VQDiffusionScheduler(self.num_embed)
          _lowercase              =LearnedClassifierFreeSamplingEmbeddings(
              learnable=snake_case,      hidden_size=self.text_embedder_hidden_size,      length=tokenizer.model_max_length)

          _lowercase              =VQDiffusionPipeline(
              vqvae=snake_case,      text_encoder=snake_case,      tokenizer=snake_case,      transformer=snake_case,      scheduler=snake_case,      learned_classifier_free_sampling_embeddings=snake_case,      )
          _lowercase              =pipe.to(snake_case)
          pipe.set_progress_bar_config(disable=snake_case)

          _lowercase              ='teddy bear playing in the pool'

          _lowercase              =torch.Generator(device=snake_case).manual_seed(0)
          _lowercase              =pipe([prompt],      generator=snake_case,      num_inference_steps=2,      output_type='np')
          _lowercase              =output.images

          _lowercase              =torch.Generator(device=snake_case).manual_seed(0)
          _lowercase              =pipe(
              [prompt],      generator=snake_case,      output_type='np',      return_dict=snake_case,      num_inference_steps=2)[0]

          _lowercase              =image[0, -3:, -3:, -1]
          _lowercase              =image_from_tuple[0, -3:, -3:, -1]

          assert image.shape == (1, 24, 24, 3)

          _lowercase              =np.array([0.6_6_9_3, 0.6_0_7_5, 0.4_9_5_9, 0.5_7_0_1, 0.5_5_8_3, 0.4_3_3_3, 0.6_1_7_1, 0.5_6_8_4, 0.4_9_8_8])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 2.0
          assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2






@slow
@require_torch_gpu
class 					SCREAMING_SNAKE_CASE_  (       unittest.TestCase ):
       """simple docstring"""


       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

          _lowercase              =load_numpy(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
              '/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy')

          _lowercase              =VQDiffusionPipeline.from_pretrained('microsoft/vq-diffusion-ithq')
          _lowercase              =pipeline.to(snake_case)
          pipeline.set_progress_bar_config(disable=snake_case)

          # requires GPU generator for gumbel softmax
          # don't use GPU generator in tests though
          _lowercase              =torch.Generator(device=snake_case).manual_seed(0)
          _lowercase              =pipeline(
              'teddy bear playing in the pool',      num_images_per_prompt=1,      generator=snake_case,      output_type='np',      )

          _lowercase              =output.images[0]

          assert image.shape == (256, 256, 3)
          assert np.abs(expected_image - image).max() < 2.0










import unittest

from transformers import MraConfig, 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, floats_tensor, ids_tensor, random_attention_mask


if is_torch_available():
    import torch

    from transformers import (
        MraForMaskedLM,
        MraForMultipleChoice,
        MraForQuestionAnswering,
        MraForSequenceClassification,
        MraForTokenClassification,
        MraModel,
    )
    from transformers.models.mra.modeling_mra import MRA_PRETRAINED_MODEL_ARCHIVE_LIST






class 					SCREAMING_SNAKE_CASE_  :
       """simple docstring"""


       def __init__(				self	:str,      snake_case	:Any,      snake_case	:str=2,      snake_case	:int=8,      snake_case	:List[Any]=True,      snake_case	:str=True,      snake_case	:Optional[int]=True,      snake_case	:str=True,      snake_case	:Optional[Any]=99,      snake_case	:Tuple=16,      snake_case	:Optional[Any]=5,      snake_case	:Optional[Any]=2,      snake_case	:Optional[Any]=36,      snake_case	:List[str]="gelu",      snake_case	:str=0.0,      snake_case	:List[Any]=0.0,      snake_case	:List[str]=512,      snake_case	:Dict=16,      snake_case	:int=2,      snake_case	:Tuple=0.0_2,      snake_case	:List[str]=3,      snake_case	:Any=4,      snake_case	:Union[str, Any]=None,      ):



          """simple docstring"""

          _lowercase              =parent
          _lowercase              =batch_size
          _lowercase              =seq_length
          _lowercase              =is_training
          _lowercase              =use_input_mask
          _lowercase              =use_token_type_ids
          _lowercase              =use_labels
          _lowercase              =vocab_size
          _lowercase              =hidden_size
          _lowercase              =num_hidden_layers
          _lowercase              =num_attention_heads
          _lowercase              =intermediate_size
          _lowercase              =hidden_act
          _lowercase              =hidden_dropout_prob
          _lowercase              =attention_probs_dropout_prob
          _lowercase              =max_position_embeddings
          _lowercase              =type_vocab_size
          _lowercase              =type_sequence_label_size
          _lowercase              =initializer_range
          _lowercase              =num_labels
          _lowercase              =num_choices
          _lowercase              =scope





       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

          _lowercase              =ids_tensor([self.batch_size, self.seq_length],      self.vocab_size)

          _lowercase              =None
          if self.use_input_mask:
             _lowercase              =random_attention_mask([self.batch_size, self.seq_length])

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

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

          _lowercase              =self.get_config()

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





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



          """simple docstring"""

          return MraConfig(
              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,      is_decoder=snake_case,      initializer_range=self.initializer_range,      )





       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          _lowercase              =self.get_config()
          _lowercase              =300
          return config





       def 				UpperCamelCase__						(				self	:Any):



          """simple docstring"""

          (
              (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       
          )              =self.prepare_config_and_inputs()

          _lowercase              =True
          _lowercase              =floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
          _lowercase              =ids_tensor([self.batch_size, self.seq_length],      vocab_size=2)

          return (
              config,
              input_ids,
              token_type_ids,
              input_mask,
              sequence_labels,
              token_labels,
              choice_labels,
              encoder_hidden_states,
              encoder_attention_mask,
          )





       def 				UpperCamelCase__						(				self	:Optional[int],      snake_case	:Optional[int],      snake_case	:List[str],      snake_case	:int,      snake_case	:Tuple,      snake_case	:Optional[int],      snake_case	:Dict,      snake_case	:str):



          """simple docstring"""

          _lowercase              =MraModel(config=snake_case)
          model.to(snake_case)
          model.eval()
          _lowercase              =model(snake_case,      attention_mask=snake_case,      token_type_ids=snake_case)
          _lowercase              =model(snake_case,      token_type_ids=snake_case)
          _lowercase              =model(snake_case)
          self.parent.assertEqual(result.last_hidden_state.shape,      (self.batch_size, self.seq_length, self.hidden_size))





       def 				UpperCamelCase__						(				self	:Any,      snake_case	:Any,      snake_case	:str,      snake_case	:Optional[int],      snake_case	:int,      snake_case	:int,      snake_case	:int,      snake_case	:Optional[Any],      snake_case	:Tuple,      snake_case	:Tuple,      ):



          """simple docstring"""

          _lowercase              =True
          _lowercase              =MraModel(snake_case)
          model.to(snake_case)
          model.eval()
          _lowercase              =model(
              snake_case,      attention_mask=snake_case,      token_type_ids=snake_case,      encoder_hidden_states=snake_case,      encoder_attention_mask=snake_case,      )
          _lowercase              =model(
              snake_case,      attention_mask=snake_case,      token_type_ids=snake_case,      encoder_hidden_states=snake_case,      )
          _lowercase              =model(snake_case,      attention_mask=snake_case,      token_type_ids=snake_case)
          self.parent.assertEqual(result.last_hidden_state.shape,      (self.batch_size, self.seq_length, self.hidden_size))





       def 				UpperCamelCase__						(				self	:str,      snake_case	:Optional[int],      snake_case	:Optional[int],      snake_case	:int,      snake_case	:Optional[int],      snake_case	:List[str],      snake_case	:Optional[int],      snake_case	:Optional[Any]):



          """simple docstring"""

          _lowercase              =MraForMaskedLM(config=snake_case)
          model.to(snake_case)
          model.eval()
          _lowercase              =model(snake_case,      attention_mask=snake_case,      token_type_ids=snake_case,      labels=snake_case)
          self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.seq_length, self.vocab_size))





       def 				UpperCamelCase__						(				self	:Tuple,      snake_case	:Union[str, Any],      snake_case	:Dict,      snake_case	:Any,      snake_case	:Union[str, Any],      snake_case	:List[Any],      snake_case	:Tuple,      snake_case	:Any):



          """simple docstring"""

          _lowercase              =MraForQuestionAnswering(config=snake_case)
          model.to(snake_case)
          model.eval()
          _lowercase              =model(
              snake_case,      attention_mask=snake_case,      token_type_ids=snake_case,      start_positions=snake_case,      end_positions=snake_case,      )
          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	:int,      snake_case	:List[str],      snake_case	:int,      snake_case	:Optional[Any],      snake_case	:Union[str, Any],      snake_case	:str,      snake_case	:List[Any],      snake_case	:Optional[Any]):



          """simple docstring"""

          _lowercase              =self.num_labels
          _lowercase              =MraForSequenceClassification(snake_case)
          model.to(snake_case)
          model.eval()
          _lowercase              =model(snake_case,      attention_mask=snake_case,      token_type_ids=snake_case,      labels=snake_case)
          self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.num_labels))





       def 				UpperCamelCase__						(				self	:Any,      snake_case	:Optional[int],      snake_case	:List[Any],      snake_case	:Optional[Any],      snake_case	:int,      snake_case	:int,      snake_case	:Dict,      snake_case	:List[Any]):



          """simple docstring"""

          _lowercase              =self.num_labels
          _lowercase              =MraForTokenClassification(config=snake_case)
          model.to(snake_case)
          model.eval()
          _lowercase              =model(snake_case,      attention_mask=snake_case,      token_type_ids=snake_case,      labels=snake_case)
          self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.seq_length, self.num_labels))





       def 				UpperCamelCase__						(				self	:Optional[int],      snake_case	:List[str],      snake_case	:Optional[int],      snake_case	:Dict,      snake_case	:str,      snake_case	:Optional[Any],      snake_case	:Tuple,      snake_case	:Dict):



          """simple docstring"""

          _lowercase              =self.num_choices
          _lowercase              =MraForMultipleChoice(config=snake_case)
          model.to(snake_case)
          model.eval()
          _lowercase              =input_ids.unsqueeze(1).expand(-1,      self.num_choices,      -1).contiguous()
          _lowercase              =token_type_ids.unsqueeze(1).expand(-1,      self.num_choices,      -1).contiguous()
          _lowercase              =input_mask.unsqueeze(1).expand(-1,      self.num_choices,      -1).contiguous()
          _lowercase              =model(
              snake_case,      attention_mask=snake_case,      token_type_ids=snake_case,      labels=snake_case,      )
          self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.num_choices))





       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          _lowercase              =self.prepare_config_and_inputs()
          (
              (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       
          )              =config_and_inputs
          _lowercase              ={'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
          return config, inputs_dict






@require_torch
class 					SCREAMING_SNAKE_CASE_  (       _a  ,						unittest.TestCase ):
       """simple docstring"""


       __lowerCAmelCase    :		Optional[int]					=(
           (
               MraModel,
               MraForMaskedLM,
               MraForMultipleChoice,
               MraForQuestionAnswering,
               MraForSequenceClassification,
               MraForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       __lowerCAmelCase    :		int					=False
       __lowerCAmelCase    :		Union[str, Any]					=False
       __lowerCAmelCase    :		str					=False
       __lowerCAmelCase    :		str					=False

       __lowerCAmelCase    :		Dict					=()





       def 				UpperCamelCase__						(				self	:Any):



          """simple docstring"""

          _lowercase              =MraModelTester(self)
          _lowercase              =ConfigTester(self,      config_class=snake_case,      hidden_size=37)





       def 				UpperCamelCase__						(				self	:str):



          """simple docstring"""

          self.config_tester.run_common_tests()





       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*snake_case)





       def 				UpperCamelCase__						(				self	:str):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          for type in ["absolute", "relative_key", "relative_key_query"]:
             _lowercase              =type
             self.model_tester.create_and_check_model(*snake_case)





       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*snake_case)





       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_multiple_choice(*snake_case)





       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_question_answering(*snake_case)





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



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_sequence_classification(*snake_case)





       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*snake_case)





       @slow
       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

          for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _lowercase              =MraModel.from_pretrained(snake_case)
             self.assertIsNotNone(snake_case)





       @unittest.skip(reason='MRA does not output attentions')
       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          return






@require_torch
class 					SCREAMING_SNAKE_CASE_  (       unittest.TestCase ):
       """simple docstring"""


       @slow
       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

          _lowercase              =MraModel.from_pretrained('uw-madison/mra-base-512-4')
          _lowercase              =torch.arange(256).unsqueeze(0)

          with torch.no_grad():
             _lowercase              =model(snake_case)[0]

          _lowercase              =torch.Size((1, 256, 768))
          self.assertEqual(output.shape,      snake_case)

          _lowercase              =torch.tensor(
              [[[-0.0_1_4_0, 0.0_8_3_0, -0.0_3_8_1], [0.1_5_4_6, 0.1_4_0_2, 0.0_2_2_0], [0.1_1_6_2, 0.0_8_5_1, 0.0_1_6_5]]])

          self.assertTrue(torch.allclose(output[:, :3, :3],      snake_case,      atol=1e-4))





       @slow
       def 				UpperCamelCase__						(				self	:List[Any]):



          """simple docstring"""

          _lowercase              =MraForMaskedLM.from_pretrained('uw-madison/mra-base-512-4')
          _lowercase              =torch.arange(256).unsqueeze(0)

          with torch.no_grad():
             _lowercase              =model(snake_case)[0]

          _lowercase              =5_0265

          _lowercase              =torch.Size((1, 256, vocab_size))
          self.assertEqual(output.shape,      snake_case)

          _lowercase              =torch.tensor(
              [[[9.2_5_9_5, -3.6_0_3_8, 1_1.8_8_1_9], [9.3_8_6_9, -3.2_6_9_3, 1_1.0_9_5_6], [1_1.8_5_2_4, -3.4_9_3_8, 1_3.1_2_1_0]]])

          self.assertTrue(torch.allclose(output[:, :3, :3],      snake_case,      atol=1e-4))





       @slow
       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

          _lowercase              =MraForMaskedLM.from_pretrained('uw-madison/mra-base-4096-8-d3')
          _lowercase              =torch.arange(4096).unsqueeze(0)

          with torch.no_grad():
             _lowercase              =model(snake_case)[0]

          _lowercase              =5_0265

          _lowercase              =torch.Size((1, 4096, vocab_size))
          self.assertEqual(output.shape,      snake_case)

          _lowercase              =torch.tensor(
              [[[5.4_7_8_9, -2.3_5_6_4, 7.5_0_6_4], [7.9_0_6_7, -1.3_3_6_9, 9.9_6_6_8], [9.0_7_1_2, -1.8_1_0_6, 7.0_3_8_0]]])

          self.assertTrue(torch.allclose(output[:, :3, :3],      snake_case,      atol=1e-4))


