





import inspect
import unittest

import numpy as np

from transformers import BeitConfig
from transformers.testing_utils import require_flax, require_vision, slow
from transformers.utils import cached_property, is_flax_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor


if is_flax_available():
     import jax

     from transformers import FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling, FlaxBeitModel

if is_vision_available():
     from PIL import Image

     from transformers import BeitImageProcessor



class      _SCREAMING_SNAKE_CASE		(			unittest.TestCase					):







 '''simple docstring'''

 def __init__(self       :     Union[str, Any] ,				UpperCAmelCase_       :     List[str] ,				UpperCAmelCase_       :     List[Any]=100 ,				UpperCAmelCase_       :     Any=13 ,				UpperCAmelCase_       :     Dict=30 ,				UpperCAmelCase_       :     Optional[int]=2 ,				UpperCAmelCase_       :     str=3 ,				UpperCAmelCase_       :     Tuple=True ,				UpperCAmelCase_       :     str=True ,				UpperCAmelCase_       :     Union[str, Any]=32 ,				UpperCAmelCase_       :     int=5 ,				UpperCAmelCase_       :     List[Any]=4 ,				UpperCAmelCase_       :     Optional[Any]=37 ,				UpperCAmelCase_       :     Any="gelu" ,				UpperCAmelCase_       :     List[str]=0.1 ,				UpperCAmelCase_       :     int=0.1 ,				UpperCAmelCase_       :     Tuple=10 ,				UpperCAmelCase_       :     int=0.02 ,				UpperCAmelCase_       :     Tuple=3 ,				) ->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:   Dict    						=parent
   lowerCamelCase__:   List[str]    						=vocab_size
   lowerCamelCase__:   Tuple    						=batch_size
   lowerCamelCase__:   Optional[int]    						=image_size
   lowerCamelCase__:   str    						=patch_size
   lowerCamelCase__:   int    						=num_channels
   lowerCamelCase__:   Union[str, Any]    						=is_training
   lowerCamelCase__:   Optional[int]    						=use_labels
   lowerCamelCase__:   Any    						=hidden_size
   lowerCamelCase__:   str    						=num_hidden_layers
   lowerCamelCase__:   Union[str, Any]    						=num_attention_heads
   lowerCamelCase__:   str    						=intermediate_size
   lowerCamelCase__:   Dict    						=hidden_act
   lowerCamelCase__:   Optional[Any]    						=hidden_dropout_prob
   lowerCamelCase__:   List[str]    						=attention_probs_dropout_prob
   lowerCamelCase__:   str    						=type_sequence_label_size
   lowerCamelCase__:   Optional[Any]    						=initializer_range

   # in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
   lowerCamelCase__:   Union[str, Any]    						=(image_size // patch_size) ** 2
   lowerCamelCase__:   List[str]    						=num_patches + 1

 def 							SCREAMING_SNAKE_CASE_      (self       :     Optional[Any]) ->List[str]:



   '''simple docstring'''



   lowerCamelCase__:   Tuple    						=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

   lowerCamelCase__:   Any    						=None
   if self.use_labels:
     lowerCamelCase__:   Dict    						=ids_tensor([self.batch_size] ,				self.type_sequence_label_size)

   lowerCamelCase__:   List[str]    						=BeitConfig(
       vocab_size=self.vocab_size ,				image_size=self.image_size ,				patch_size=self.patch_size ,				num_channels=self.num_channels ,				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 ,				is_decoder=A_ ,				initializer_range=self.initializer_range ,				)

   return config, pixel_values, labels

 def 							SCREAMING_SNAKE_CASE_      (self       :     Optional[int] ,				UpperCAmelCase_       :     str ,				UpperCAmelCase_       :     str ,				UpperCAmelCase_       :     List[str]) ->List[Any]:



   '''simple docstring'''



   lowerCamelCase__:   Optional[int]    						=FlaxBeitModel(config=A_)
   lowerCamelCase__:   Optional[Any]    						=model(A_)
   self.parent.assertEqual(result.last_hidden_state.shape ,				(self.batch_size, self.seq_length, self.hidden_size))

 def 							SCREAMING_SNAKE_CASE_      (self       :     List[str] ,				UpperCAmelCase_       :     Any ,				UpperCAmelCase_       :     Any ,				UpperCAmelCase_       :     str) ->str:



   '''simple docstring'''



   lowerCamelCase__:   Union[str, Any]    						=FlaxBeitForMaskedImageModeling(config=A_)
   lowerCamelCase__:   int    						=model(A_)
   self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.seq_length - 1, self.vocab_size))

 def 							SCREAMING_SNAKE_CASE_      (self       :     str ,				UpperCAmelCase_       :     List[Any] ,				UpperCAmelCase_       :     Optional[int] ,				UpperCAmelCase_       :     List[str]) ->str:



   '''simple docstring'''



   lowerCamelCase__:   Any    						=self.type_sequence_label_size
   lowerCamelCase__:   Any    						=FlaxBeitForImageClassification(config=A_)
   lowerCamelCase__:   Tuple    						=model(A_)
   self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.type_sequence_label_size))

   # test greyscale images
   lowerCamelCase__:   Optional[Any]    						=1
   lowerCamelCase__:   Any    						=FlaxBeitForImageClassification(A_)

   lowerCamelCase__:   str    						=floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
   lowerCamelCase__:   int    						=model(A_)







 def 							SCREAMING_SNAKE_CASE_      (self       :     Dict) ->List[str]:



   '''simple docstring'''



   lowerCamelCase__:   int    						=self.prepare_config_and_inputs()
   (
       (
       lowerCamelCase__	
   )						,      (
       lowerCamelCase__	
   )						,      (
       lowerCamelCase__	
   )						,      	
   ):   int    						=config_and_inputs
   lowerCamelCase__:   Union[str, Any]    						={"pixel_values": pixel_values}
   return config, inputs_dict



@require_flax
class      _SCREAMING_SNAKE_CASE		(			_UpperCAmelCase					,							unittest.TestCase					):







 '''simple docstring'''
 lowercase_	=    (
     (FlaxBeitModel, FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling) if is_flax_available() else ()
 )

 def 							SCREAMING_SNAKE_CASE_      (self       :     List[Any]) ->None:



   '''simple docstring'''



   lowerCamelCase__:   Dict    						=FlaxBeitModelTester(self)
   lowerCamelCase__:   Dict    						=ConfigTester(self ,				config_class=A_ ,				has_text_modality=A_ ,				hidden_size=37)

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



   '''simple docstring'''



   self.config_tester.run_common_tests()

 def 							SCREAMING_SNAKE_CASE_      (self       :     int) ->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__						,      lowerCamelCase__:   Tuple    						=self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     lowerCamelCase__:   Tuple    						=model_class(A_)
     lowerCamelCase__:   List[str]    						=inspect.signature(model.__call__)
     # signature.parameters is an OrderedDict => so arg_names order is deterministic
     lowerCamelCase__:   str    						=[*signature.parameters.keys()]

     lowerCamelCase__:   Tuple    						=["pixel_values"]
     self.assertListEqual(arg_names[:1] ,				A_)

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



   '''simple docstring'''



   lowerCamelCase__						,      lowerCamelCase__:   List[Any]    						=self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     with self.subTest(model_class.__name__):
       lowerCamelCase__:   Optional[int]    						=self._prepare_for_class(A_ ,				A_)
       lowerCamelCase__:   Tuple    						=model_class(A_)

       @jax.jit
       def model_jitted(UpperCAmelCase_       :     Union[str, Any] ,				**UpperCAmelCase_       :     Union[str, Any]):
         return model(pixel_values=A_ ,				**A_)

       with self.subTest("JIT Enabled"):
         lowerCamelCase__:   Tuple    						=model_jitted(**A_).to_tuple()

       with self.subTest("JIT Disabled"):
         with jax.disable_jit():
           lowerCamelCase__:   List[Any]    						=model_jitted(**A_).to_tuple()

       self.assertEqual(len(A_) ,				len(A_))
       for jitted_output, output in zip(A_ ,				A_):
         self.assertEqual(jitted_output.shape ,				output.shape)

 def 							SCREAMING_SNAKE_CASE_      (self       :     str) ->List[Any]:



   '''simple docstring'''



   lowerCamelCase__:   List[str]    						=self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*A_)

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



   '''simple docstring'''



   lowerCamelCase__:   Optional[int]    						=self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_masked_lm(*A_)

 def 							SCREAMING_SNAKE_CASE_      (self       :     str) ->int:



   '''simple docstring'''



   lowerCamelCase__:   List[str]    						=self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_classification(*A_)







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



   '''simple docstring'''



   for model_class_name in self.all_model_classes:
     lowerCamelCase__:   str    						=model_class_name.from_pretrained("microsoft/beit-base-patch16-224")
     lowerCamelCase__:   Union[str, Any]    						=model(np.ones((1, 3, 224, 224)))
     self.assertIsNotNone(A_)





def   lowerCAmelCase_				( )  ->	List[str]:







  """simple docstring"""





  lowerCamelCase__:   Tuple    						=Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"	)
  return image






@require_vision
@require_flax
class      _SCREAMING_SNAKE_CASE		(			unittest.TestCase					):







 '''simple docstring'''

 @cached_property
 def 							SCREAMING_SNAKE_CASE_      (self       :     str) ->str:



   '''simple docstring'''



   return BeitImageProcessor.from_pretrained("microsoft/beit-base-patch16-224") if is_vision_available() else None

 @slow
 def 							SCREAMING_SNAKE_CASE_      (self       :     Union[str, Any]) ->List[str]:



   '''simple docstring'''



   lowerCamelCase__:   Tuple    						=FlaxBeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k")

   lowerCamelCase__:   Tuple    						=self.default_image_processor
   lowerCamelCase__:   Tuple    						=prepare_img()
   lowerCamelCase__:   int    						=image_processor(images=A_ ,				return_tensors="np").pixel_values

   # prepare bool_masked_pos
   lowerCamelCase__:   List[str]    						=np.ones((1, 196) ,				dtype=A_)

   # forward pass
   lowerCamelCase__:   Any    						=model(pixel_values=A_ ,				bool_masked_pos=A_)
   lowerCamelCase__:   Tuple    						=outputs.logits

   # verify the logits
   lowerCamelCase__:   Optional[int]    						=(1, 196, 8_192)
   self.assertEqual(logits.shape ,				A_)

   lowerCamelCase__:   int    						=np.array(
       [[-3.2437, 0.5072, -13.9174], [-3.2456, 0.4948, -13.9401], [-3.2033, 0.5121, -13.8550]])

   self.assertTrue(np.allclose(logits[bool_masked_pos][:3, :3] ,				A_ ,				atol=1E-2))

 @slow
 def 							SCREAMING_SNAKE_CASE_      (self       :     List[Any]) ->List[str]:



   '''simple docstring'''



   lowerCamelCase__:   Any    						=FlaxBeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224")

   lowerCamelCase__:   Optional[Any]    						=self.default_image_processor
   lowerCamelCase__:   List[Any]    						=prepare_img()
   lowerCamelCase__:   Any    						=image_processor(images=A_ ,				return_tensors="np")

   # forward pass
   lowerCamelCase__:   Dict    						=model(**A_)
   lowerCamelCase__:   Dict    						=outputs.logits

   # verify the logits
   lowerCamelCase__:   int    						=(1, 1_000)
   self.assertEqual(logits.shape ,				A_)

   lowerCamelCase__:   str    						=np.array([-1.2385, -1.0987, -1.0108])

   self.assertTrue(np.allclose(logits[0, :3] ,				A_ ,				atol=1E-4))

   lowerCamelCase__:   int    						=281
   self.assertEqual(logits.argmax(-1).item() ,				A_)







 @slow
 def 							SCREAMING_SNAKE_CASE_      (self       :     int) ->str:



   '''simple docstring'''



   lowerCamelCase__:   Tuple    						=FlaxBeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k")

   lowerCamelCase__:   Union[str, Any]    						=self.default_image_processor
   lowerCamelCase__:   List[str]    						=prepare_img()
   lowerCamelCase__:   Optional[int]    						=image_processor(images=A_ ,				return_tensors="np")

   # forward pass
   lowerCamelCase__:   List[Any]    						=model(**A_)
   lowerCamelCase__:   List[Any]    						=outputs.logits

   # verify the logits
   lowerCamelCase__:   List[str]    						=(1, 21_841)
   self.assertEqual(logits.shape ,				A_)

   lowerCamelCase__:   List[str]    						=np.array([1.6881, -0.2787, 0.5901])

   self.assertTrue(np.allclose(logits[0, :3] ,				A_ ,				atol=1E-4))

   lowerCamelCase__:   int    						=2_396
   self.assertEqual(logits.argmax(-1).item() ,				A_)








import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    TextToVideoSDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, load_numpy, skip_mps, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



@skip_mps
class      _SCREAMING_SNAKE_CASE		(			__SCREAMING_SNAKE_CASE					,							unittest.TestCase					):







 '''simple docstring'''
 lowercase_	=    TextToVideoSDPipeline
 lowercase_	=    TEXT_TO_IMAGE_PARAMS
 lowercase_	=    TEXT_TO_IMAGE_BATCH_PARAMS
 # No `output_type`.
 lowercase_	=    frozenset(
     [
         "num_inference_steps",
         "generator",
         "latents",
         "return_dict",
         "callback",
         "callback_steps",
     ]					)

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



   '''simple docstring'''



   torch.manual_seed(0)
   lowerCamelCase__:   Optional[int]    						=UNetaDConditionModel(
       block_out_channels=(32, 64, 64, 64) ,				layers_per_block=2 ,				sample_size=32 ,				in_channels=4 ,				out_channels=4 ,				down_block_types=("CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D") ,				up_block_types=("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D") ,				cross_attention_dim=32 ,				attention_head_dim=4 ,				)
   lowerCamelCase__:   Union[str, Any]    						=DDIMScheduler(
       beta_start=0.0_0085 ,				beta_end=0.012 ,				beta_schedule="scaled_linear" ,				clip_sample=UpperCAmelCase_ ,				set_alpha_to_one=UpperCAmelCase_ ,				)
   torch.manual_seed(0)
   lowerCamelCase__:   List[str]    						=AutoencoderKL(
       block_out_channels=[32, 64] ,				in_channels=3 ,				out_channels=3 ,				down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"] ,				up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"] ,				latent_channels=4 ,				sample_size=128 ,				)
   torch.manual_seed(0)
   lowerCamelCase__:   Optional[int]    						=CLIPTextConfig(
       bos_token_id=0 ,				eos_token_id=2 ,				hidden_size=32 ,				intermediate_size=37 ,				layer_norm_eps=1E-0_5 ,				num_attention_heads=4 ,				num_hidden_layers=5 ,				pad_token_id=1 ,				vocab_size=1_000 ,				hidden_act="gelu" ,				projection_dim=512 ,				)
   lowerCamelCase__:   Optional[Any]    						=CLIPTextModel(UpperCAmelCase_)
   lowerCamelCase__:   Dict    						=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

   lowerCamelCase__:   Tuple    						={
       "unet": unet,
       "scheduler": scheduler,
       "vae": vae,
       "text_encoder": text_encoder,
       "tokenizer": tokenizer,
   }
   return components

 def 							SCREAMING_SNAKE_CASE_      (self       :     str ,				UpperCAmelCase_       :     Dict ,				UpperCAmelCase_       :     List[Any]=0) ->Union[str, Any]:



   '''simple docstring'''



   if str(UpperCAmelCase_).startswith("mps"):
     lowerCamelCase__:   Optional[int]    						=torch.manual_seed(UpperCAmelCase_)
   else:
     lowerCamelCase__:   Any    						=torch.Generator(device=UpperCAmelCase_).manual_seed(UpperCAmelCase_)
   lowerCamelCase__:   List[Any]    						={
       "prompt": "A painting of a squirrel eating a burger",
       "generator": generator,
       "num_inference_steps": 2,
       "guidance_scale": 6.0,
       "output_type": "pt",
   }
   return inputs

 def 							SCREAMING_SNAKE_CASE_      (self       :     List[Any]) ->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:   List[str]    						="cpu"  # ensure determinism for the device-dependent torch.Generator
   lowerCamelCase__:   Optional[Any]    						=self.get_dummy_components()
   lowerCamelCase__:   List[Any]    						=TextToVideoSDPipeline(**UpperCAmelCase_)
   lowerCamelCase__:   int    						=sd_pipe.to(UpperCAmelCase_)
   sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_)

   lowerCamelCase__:   Any    						=self.get_dummy_inputs(UpperCAmelCase_)
   lowerCamelCase__:   List[Any]    						="np"
   lowerCamelCase__:   Optional[int]    						=sd_pipe(**UpperCAmelCase_).frames
   lowerCamelCase__:   Dict    						=frames[0][-3:, -3:, -1]

   assert frames[0].shape == (64, 64, 3)
   lowerCamelCase__:   Optional[int]    						=np.array([158.0, 160.0, 153.0, 125.0, 100.0, 121.0, 111.0, 93.0, 113.0])

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

 def 							SCREAMING_SNAKE_CASE_      (self       :     Tuple) ->str:



   '''simple docstring'''



   self._test_attention_slicing_forward_pass(test_mean_pixel_difference=UpperCAmelCase_ ,				expected_max_diff=3E-3)

 @unittest.skipIf(
     torch_device != "cuda" or not is_xformers_available() ,				reason="XFormers attention is only available with CUDA and `xformers` installed" ,				)
 def 							SCREAMING_SNAKE_CASE_      (self       :     Optional[Any]) ->List[Any]:



   '''simple docstring'''



   self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=UpperCAmelCase_ ,				expected_max_diff=1E-2)

 @unittest.skip(reason="Batching needs to be properly figured out first for this pipeline.")
 def 							SCREAMING_SNAKE_CASE_      (self       :     Optional[int]) ->str:



   '''simple docstring'''



   pass

 @unittest.skip(reason="Batching needs to be properly figured out first for this pipeline.")
 def 							SCREAMING_SNAKE_CASE_      (self       :     List[Any]) ->Optional[Any]:



   '''simple docstring'''



   pass

 @unittest.skip(reason="`num_images_per_prompt` argument is not supported for this pipeline.")
 def 							SCREAMING_SNAKE_CASE_      (self       :     List[Any]) ->Optional[int]:



   '''simple docstring'''



   pass







 def 							SCREAMING_SNAKE_CASE_      (self       :     Optional[int]) ->Any:



   '''simple docstring'''



   return super().test_progress_bar()



@slow
@skip_mps
class      _SCREAMING_SNAKE_CASE		(			unittest.TestCase					):







 '''simple docstring'''

 def 							SCREAMING_SNAKE_CASE_      (self       :     List[Any]) ->int:



   '''simple docstring'''



   lowerCamelCase__:   Dict    						=load_numpy(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video.npy")

   lowerCamelCase__:   Optional[int]    						=TextToVideoSDPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b")
   lowerCamelCase__:   str    						=DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
   lowerCamelCase__:   Tuple    						=pipe.to("cuda")

   lowerCamelCase__:   List[Any]    						="Spiderman is surfing"
   lowerCamelCase__:   Dict    						=torch.Generator(device="cpu").manual_seed(0)

   lowerCamelCase__:   Optional[int]    						=pipe(UpperCAmelCase_ ,				generator=UpperCAmelCase_ ,				num_inference_steps=25 ,				output_type="pt").frames
   lowerCamelCase__:   str    						=video_frames.cpu().numpy()

   assert np.abs(expected_video - video).mean() < 5E-2







 def 							SCREAMING_SNAKE_CASE_      (self       :     List[str]) ->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:   Optional[Any]    						=load_numpy(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video_2step.npy")

   lowerCamelCase__:   List[str]    						=TextToVideoSDPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b")
   lowerCamelCase__:   Any    						=pipe.to("cuda")

   lowerCamelCase__:   Dict    						="Spiderman is surfing"
   lowerCamelCase__:   Dict    						=torch.Generator(device="cpu").manual_seed(0)

   lowerCamelCase__:   Optional[int]    						=pipe(UpperCAmelCase_ ,				generator=UpperCAmelCase_ ,				num_inference_steps=2 ,				output_type="pt").frames
   lowerCamelCase__:   List[Any]    						=video_frames.cpu().numpy()

   assert np.abs(expected_video - video).mean() < 5E-2

