





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 	UpperCamelCase_							(  SCREAMING_SNAKE_CASE_					,	unittest.TestCase     ):
							lowerCamelCase_							     =       TextToVideoSDPipeline
							lowerCamelCase_							     =       TEXT_TO_IMAGE_PARAMS
							lowerCamelCase_							     =       TEXT_TO_IMAGE_BATCH_PARAMS
							# No `output_type`.
							lowerCamelCase_							     =       frozenset(
							    [
							        "num_inference_steps",
							        "generator",
							        "latents",
							        "return_dict",
							        "callback",
							        "callback_steps",
							    ]     )


							def   _snake_case	(				self     :Tuple		) -> Optional[int]:



									"""simple docstring"""



									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   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				,					)
									SCREAMING_SNAKE_CASE__								=   DDIMScheduler(
									    beta_start=0.0_0_0_8_5				,					beta_end=0.0_1_2				,					beta_schedule="""scaled_linear"""				,					clip_sample=__A				,					set_alpha_to_one=__A				,					)
									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   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		)
									SCREAMING_SNAKE_CASE__								=   CLIPTextConfig(
									    bos_token_id=0				,					eos_token_id=2				,					hidden_size=32				,					intermediate_size=37				,					layer_norm_eps=1E-05				,					num_attention_heads=4				,					num_hidden_layers=5				,					pad_token_id=1				,					vocab_size=1000				,					hidden_act="""gelu"""				,					projection_dim=512				,					)
									SCREAMING_SNAKE_CASE__								=   CLIPTextModel(__A		)
									SCREAMING_SNAKE_CASE__								=   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""		)

									SCREAMING_SNAKE_CASE__								=   {
									    'unet': unet,
									    'scheduler': scheduler,
									    'vae': vae,
									    'text_encoder': text_encoder,
									    'tokenizer': tokenizer,
									}
									return components


							def   _snake_case	(				self     :Optional[int]				,					__A     :Dict				,					__A     :Optional[Any]=0		) -> List[Any]:



									"""simple docstring"""



									if str(__A		).startswith("""mps"""		):
											SCREAMING_SNAKE_CASE__								=   torch.manual_seed(__A		)
									else:
											SCREAMING_SNAKE_CASE__								=   torch.Generator(device=__A		).manual_seed(__A		)
									SCREAMING_SNAKE_CASE__								=   {
									    '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   _snake_case	(				self     :Optional[int]		) -> Any:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   'cpu'  # ensure determinism for the device-dependent torch.Generator
									SCREAMING_SNAKE_CASE__								=   self.get_dummy_components()
									SCREAMING_SNAKE_CASE__								=   TextToVideoSDPipeline(**__A		)
									SCREAMING_SNAKE_CASE__								=   sd_pipe.to(__A		)
									sd_pipe.set_progress_bar_config(disable=__A		)

									SCREAMING_SNAKE_CASE__								=   self.get_dummy_inputs(__A		)
									SCREAMING_SNAKE_CASE__								=   'np'
									SCREAMING_SNAKE_CASE__								=   sd_pipe(**__A		).frames
									SCREAMING_SNAKE_CASE__								=   frames[0][-3:, -3:, -1]

									assert frames[0].shape == (64, 64, 3)
									SCREAMING_SNAKE_CASE__								=   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   _snake_case	(				self     :str		) -> Dict:



									"""simple docstring"""



									self._test_attention_slicing_forward_pass(test_mean_pixel_difference=__A				,					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   _snake_case	(				self     :Tuple		) -> Any:



									"""simple docstring"""



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


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



									"""simple docstring"""



									pass


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



									"""simple docstring"""



									pass


							@unittest.skip(reason="""`num_images_per_prompt` argument is not supported for this pipeline."""		)
							def   _snake_case	(				self     :Optional[Any]		) -> Union[str, Any]:



									"""simple docstring"""



									pass


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



									"""simple docstring"""



									return super().test_progress_bar()




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


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



									"""simple docstring"""



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

									SCREAMING_SNAKE_CASE__								=   TextToVideoSDPipeline.from_pretrained("""damo-vilab/text-to-video-ms-1.7b"""		)
									SCREAMING_SNAKE_CASE__								=   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config		)
									SCREAMING_SNAKE_CASE__								=   pipe.to("""cuda"""		)

									SCREAMING_SNAKE_CASE__								=   'Spiderman is surfing'
									SCREAMING_SNAKE_CASE__								=   torch.Generator(device="""cpu"""		).manual_seed(0		)

									SCREAMING_SNAKE_CASE__								=   pipe(__A				,					generator=__A				,					num_inference_steps=25				,					output_type="""pt"""		).frames
									SCREAMING_SNAKE_CASE__								=   video_frames.cpu().numpy()

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


							def   _snake_case	(				self     :List[Any]		) -> Tuple:



									"""simple docstring"""



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

									SCREAMING_SNAKE_CASE__								=   TextToVideoSDPipeline.from_pretrained("""damo-vilab/text-to-video-ms-1.7b"""		)
									SCREAMING_SNAKE_CASE__								=   pipe.to("""cuda"""		)

									SCREAMING_SNAKE_CASE__								=   'Spiderman is surfing'
									SCREAMING_SNAKE_CASE__								=   torch.Generator(device="""cpu"""		).manual_seed(0		)

									SCREAMING_SNAKE_CASE__								=   pipe(__A				,					generator=__A				,					num_inference_steps=2				,					output_type="""pt"""		).frames
									SCREAMING_SNAKE_CASE__								=   video_frames.cpu().numpy()

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




"""simple docstring"""







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 					A_(unittest.TestCase      ):







			"""simple docstring"""



			def __init__(	self				,   A				,   A=100				,   A=13				,   A=30				,   A=2				,   A=3				,   A=True				,   A=True				,   A=32				,   A=5				,   A=4				,   A=37				,   A="gelu"				,   A=0.1				,   A=0.1				,   A=10				,   A=0.0_2				,   A=3				,   ):
						_lowerCamelCase					:					str	       =    parent
						_lowerCamelCase					:					Optional[Any]	       =    vocab_size
						_lowerCamelCase					:					int	       =    batch_size
						_lowerCamelCase					:					Tuple	       =    image_size
						_lowerCamelCase					:					Optional[Any]	       =    patch_size
						_lowerCamelCase					:					Union[str, Any]	       =    num_channels
						_lowerCamelCase					:					Tuple	       =    is_training
						_lowerCamelCase					:					str	       =    use_labels
						_lowerCamelCase					:					Tuple	       =    hidden_size
						_lowerCamelCase					:					Tuple	       =    num_hidden_layers
						_lowerCamelCase					:					Any	       =    num_attention_heads
						_lowerCamelCase					:					int	       =    intermediate_size
						_lowerCamelCase					:					List[Any]	       =    hidden_act
						_lowerCamelCase					:					int	       =    hidden_dropout_prob
						_lowerCamelCase					:					Dict	       =    attention_probs_dropout_prob
						_lowerCamelCase					:					List[str]	       =    type_sequence_label_size
						_lowerCamelCase					:					int	       =    initializer_range

						# in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
						_lowerCamelCase					:					Optional[Any]	       =    (image_size // patch_size) ** 2
						_lowerCamelCase					:					Optional[int]	       =    num_patches + 1
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					Optional[Any]	       =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

						_lowerCamelCase					:					int	       =    None
						if self.use_labels:
									_lowerCamelCase					:					int	       =    ids_tensor([self.batch_size]				,   self.type_sequence_label_size							)

						_lowerCamelCase					:					List[Any]	       =    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    _lowerCAmelCase		(	self				,   A				,   A				,   A							):
						_lowerCamelCase					:					Tuple	       =    FlaxBeitModel(config=A							)
						_lowerCamelCase					:					List[str]	       =    model(A							)
						self.parent.assertEqual(result.last_hidden_state.shape				,   (self.batch_size, self.seq_length, self.hidden_size)							)
			def    _lowerCAmelCase		(	self				,   A				,   A				,   A							):
						_lowerCamelCase					:					List[str]	       =    FlaxBeitForMaskedImageModeling(config=A							)
						_lowerCamelCase					:					Any	       =    model(A							)
						self.parent.assertEqual(result.logits.shape				,   (self.batch_size, self.seq_length - 1, self.vocab_size)							)
			def    _lowerCAmelCase		(	self				,   A				,   A				,   A							):
						_lowerCamelCase					:					Optional[Any]	       =    self.type_sequence_label_size
						_lowerCamelCase					:					Tuple	       =    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					:					List[Any]	       =    1
						_lowerCamelCase					:					str	       =    FlaxBeitForImageClassification(A							)

						_lowerCamelCase					:					Optional[int]	       =    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]							)
						_lowerCamelCase					:					List[Any]	       =    model(A							)

			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					str	       =    self.prepare_config_and_inputs()
						(
						    (
						    _lowerCamelCase
						)       ,							(
						    _lowerCamelCase
						)       ,							(
						    _lowerCamelCase
						)       ,							
						)					:					Optional[Any]	       =    config_and_inputs
						_lowerCamelCase					:					int	       =    {'pixel_values': pixel_values}
						return config, inputs_dict







@require_flax
class 					A_(SCREAMING_SNAKE_CASE_   ,   unittest.TestCase      ):







			"""simple docstring"""



			a_       :			List[Any]												=			(
			    (FlaxBeitModel, FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling) if is_flax_available() else ()
			)
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					Tuple	       =    FlaxBeitModelTester(self							)
						_lowerCamelCase					:					List[str]	       =    ConfigTester(self				,   config_class=A				,   has_text_modality=A				,   hidden_size=37							)
			def    _lowerCAmelCase		(	self							):
						self.config_tester.run_common_tests()
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase       ,							_lowerCamelCase					:					Tuple	       =    self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									_lowerCamelCase					:					List[str]	       =    model_class(A							)
									_lowerCamelCase					:					int	       =    inspect.signature(model.__call__							)
									# signature.parameters is an OrderedDict => so arg_names order is deterministic
									_lowerCamelCase					:					Optional[Any]	       =    [*signature.parameters.keys()]

									_lowerCamelCase					:					List[Any]	       =    ['pixel_values']
									self.assertListEqual(arg_names[:1]				,   A							)
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase       ,							_lowerCamelCase					:					List[str]	       =    self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									with self.subTest(model_class.__name__							):
												_lowerCamelCase					:					Dict	       =    self._prepare_for_class(A				,   A							)
												_lowerCamelCase					:					str	       =    model_class(A							)

												@jax.jit
												def model_jitted(A				,   **A							):
															return model(pixel_values=A				,   **A							)

												with self.subTest('JIT Enabled'							):
															_lowerCamelCase					:					List[Any]	       =    model_jitted(**A							).to_tuple()

												with self.subTest('JIT Disabled'							):
															with jax.disable_jit():
																		_lowerCamelCase					:					Optional[int]	       =    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    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					Tuple	       =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*A							)
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					int	       =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*A							)
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					Optional[int]	       =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*A							)

			@slow
			def    _lowerCAmelCase		(	self							):
						for model_class_name in self.all_model_classes:
									_lowerCamelCase					:					str	       =    model_class_name.from_pretrained('microsoft/beit-base-patch16-224'							)
									_lowerCamelCase					:					Dict	       =    model(np.ones((1, 3, 224, 224)							)							)
									self.assertIsNotNone(A							)






def     UpperCAmelCase_					(				):




			'''simple docstring'''



			_lowerCamelCase					:					Union[str, Any]	       =    Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'						)
			return image







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







			"""simple docstring"""



			@cached_property
			def    _lowerCAmelCase		(	self							):
						return BeitImageProcessor.from_pretrained('microsoft/beit-base-patch16-224'							) if is_vision_available() else None
			@slow
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					List[str]	       =    FlaxBeitForMaskedImageModeling.from_pretrained('microsoft/beit-base-patch16-224-pt22k'							)

						_lowerCamelCase					:					int	       =    self.default_image_processor
						_lowerCamelCase					:					List[str]	       =    prepare_img()
						_lowerCamelCase					:					Any	       =    image_processor(images=A				,   return_tensors='np'							).pixel_values

						# prepare bool_masked_pos
						_lowerCamelCase					:					Optional[Any]	       =    np.ones((1, 196)				,   dtype=A							)

						# forward pass
						_lowerCamelCase					:					Optional[int]	       =    model(pixel_values=A				,   bool_masked_pos=A							)
						_lowerCamelCase					:					Optional[int]	       =    outputs.logits

						# verify the logits
						_lowerCamelCase					:					str	       =    (1, 196, 8192)
						self.assertEqual(logits.shape				,   A							)

						_lowerCamelCase					:					Dict	       =    np.array(
						    [[-3.2_4_3_7, 0.5_0_7_2, -1_3.9_1_7_4], [-3.2_4_5_6, 0.4_9_4_8, -1_3.9_4_0_1], [-3.2_0_3_3, 0.5_1_2_1, -1_3.8_5_5_0]]							)

						self.assertTrue(np.allclose(logits[bool_masked_pos][:3, :3]				,   A				,   atol=1E-2							)							)
			@slow
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					Any	       =    FlaxBeitForImageClassification.from_pretrained('microsoft/beit-base-patch16-224'							)

						_lowerCamelCase					:					Tuple	       =    self.default_image_processor
						_lowerCamelCase					:					Dict	       =    prepare_img()
						_lowerCamelCase					:					Optional[Any]	       =    image_processor(images=A				,   return_tensors='np'							)

						# forward pass
						_lowerCamelCase					:					Tuple	       =    model(**A							)
						_lowerCamelCase					:					Optional[int]	       =    outputs.logits

						# verify the logits
						_lowerCamelCase					:					int	       =    (1, 1000)
						self.assertEqual(logits.shape				,   A							)

						_lowerCamelCase					:					Tuple	       =    np.array([-1.2_3_8_5, -1.0_9_8_7, -1.0_1_0_8]							)

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

						_lowerCamelCase					:					str	       =    281
						self.assertEqual(logits.argmax(-1							).item()				,   A							)

			@slow
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					Optional[Any]	       =    FlaxBeitForImageClassification.from_pretrained('microsoft/beit-large-patch16-224-pt22k-ft22k'							)

						_lowerCamelCase					:					Any	       =    self.default_image_processor
						_lowerCamelCase					:					Union[str, Any]	       =    prepare_img()
						_lowerCamelCase					:					Tuple	       =    image_processor(images=A				,   return_tensors='np'							)

						# forward pass
						_lowerCamelCase					:					Union[str, Any]	       =    model(**A							)
						_lowerCamelCase					:					Optional[int]	       =    outputs.logits

						# verify the logits
						_lowerCamelCase					:					List[Any]	       =    (1, 2_1841)
						self.assertEqual(logits.shape				,   A							)

						_lowerCamelCase					:					Tuple	       =    np.array([1.6_8_8_1, -0.2_7_8_7, 0.5_9_0_1]							)

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

						_lowerCamelCase					:					List[Any]	       =    2396
						self.assertEqual(logits.argmax(-1							).item()				,   A							)





