





import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import VideoMAEConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
				import torch
				from torch import nn

				from transformers import (
				    MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING,
				    VideoMAEForPreTraining,
				    VideoMAEForVideoClassification,
				    VideoMAEModel,
				)
				from transformers.models.videomae.modeling_videomae import VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
				from transformers import VideoMAEImageProcessor

class 		SCREAMING_SNAKE_CASE__    :


						def __init__(self	:				Any    ,      a__	:				str    ,      a__	:				Tuple=13    ,      a__	:				int=10    ,      a__	:				Tuple=3    ,      a__	:				Tuple=2    ,      a__	:				List[str]=2    ,      a__	:				Union[str, Any]=2    ,      a__	:				Dict=True    ,      a__	:				Union[str, Any]=True    ,      a__	:				List[str]=32    ,      a__	:				Dict=5    ,      a__	:				Tuple=4    ,      a__	:				List[str]=37    ,      a__	:				Dict="gelu"    ,      a__	:				Dict=0.1    ,      a__	:				int=0.1    ,      a__	:				Optional[Any]=10    ,      a__	:				List[str]=0.0_2    ,      a__	:				str=0.9    ,      a__	:				Optional[Any]=None    ,      ):





									"""simple docstring"""



									__snake_case								=       parent
									__snake_case								=       batch_size
									__snake_case								=       image_size
									__snake_case								=       num_channels
									__snake_case								=       patch_size
									__snake_case								=       tubelet_size
									__snake_case								=       num_frames
									__snake_case								=       is_training
									__snake_case								=       use_labels
									__snake_case								=       hidden_size
									__snake_case								=       num_hidden_layers
									__snake_case								=       num_attention_heads
									__snake_case								=       intermediate_size
									__snake_case								=       hidden_act
									__snake_case								=       hidden_dropout_prob
									__snake_case								=       attention_probs_dropout_prob
									__snake_case								=       type_sequence_label_size
									__snake_case								=       initializer_range
									__snake_case								=       mask_ratio
									__snake_case								=       scope

									# in VideoMAE, the number of tokens equals num_frames/tubelet_size * num_patches per frame
									__snake_case								=       (image_size // patch_size) ** 2
									__snake_case								=       (num_frames // tubelet_size) * self.num_patches_per_frame

									# use this variable to define bool_masked_pos
									__snake_case								=       int(mask_ratio * self.seq_length				)


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





									"""simple docstring"""



									__snake_case								=       floats_tensor(
									    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]				)

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

									__snake_case								=       self.get_config()

									return config, pixel_values, labels


						def        a			(self	:				Optional[Any]				):





									"""simple docstring"""



									return VideoMAEConfig(
									    image_size=self.image_size    ,      patch_size=self.patch_size    ,      num_channels=self.num_channels    ,      num_frames=self.num_frames    ,      tubelet_size=self.tubelet_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    ,      is_decoder=a__    ,      initializer_range=self.initializer_range    ,      )


						def        a			(self	:				List[Any]    ,      a__	:				int    ,      a__	:				Tuple    ,      a__	:				List[Any]				):





									"""simple docstring"""



									__snake_case								=       VideoMAEModel(config=a__				)
									model.to(a__				)
									model.eval()
									__snake_case								=       model(a__				)
									self.parent.assertEqual(result.last_hidden_state.shape    ,      (self.batch_size, self.seq_length, self.hidden_size)				)


						def        a			(self	:				str    ,      a__	:				Union[str, Any]    ,      a__	:				Optional[Any]    ,      a__	:				Tuple				):





									"""simple docstring"""



									__snake_case								=       VideoMAEForPreTraining(a__				)
									model.to(a__				)
									model.eval()
									# important: each video needs to have the same number of masked patches
									# hence we define a single mask, which we then repeat for each example in the batch
									__snake_case								=       torch.ones((self.num_masks,)				)
									__snake_case								=       torch.cat([mask, torch.zeros(self.seq_length - mask.size(0				)				)]				)
									__snake_case								=       mask.expand(self.batch_size    ,      -1				).bool()

									__snake_case								=       model(a__    ,      a__				)
									# model only returns predictions for masked patches
									__snake_case								=       mask.sum().item()
									__snake_case								=       3 * self.tubelet_size * self.patch_size**2
									self.parent.assertEqual(result.logits.shape    ,      (self.batch_size, num_masked_patches, decoder_num_labels)				)


						def        a			(self	:				Any				):





									"""simple docstring"""



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






@require_torch
class 		SCREAMING_SNAKE_CASE__    (	_UpperCAmelCase						,				_UpperCAmelCase						,				unittest.TestCase  ):
						A_						:    Dict							=   (
						    (VideoMAEModel, VideoMAEForPreTraining, VideoMAEForVideoClassification) if is_torch_available() else ()
						)
						A_						:    Any							=   (
						    {'feature-extraction': VideoMAEModel, 'video-classification': VideoMAEForVideoClassification}
						    if is_torch_available()
						    else {}
						)

						A_						:    str							=   False
						A_						:    Tuple							=   False
						A_						:    List[str]							=   False
						A_						:    Tuple							=   False


						def        a			(self	:				List[str]				):





									"""simple docstring"""



									__snake_case								=       VideoMAEModelTester(self				)
									__snake_case								=       ConfigTester(self    ,      config_class=a__    ,      has_text_modality=a__    ,      hidden_size=37				)


						def        a			(self	:				Tuple    ,      a__	:				Any    ,      a__	:				Tuple    ,      a__	:				Optional[Any]=False				):





									"""simple docstring"""



									__snake_case								=       copy.deepcopy(a__				)

									if model_class == VideoMAEForPreTraining:
												# important: each video needs to have the same number of masked patches
												# hence we define a single mask, which we then repeat for each example in the batch
												__snake_case								=       torch.ones((self.model_tester.num_masks,)				)
												__snake_case								=       torch.cat([mask, torch.zeros(self.model_tester.seq_length - mask.size(0				)				)]				)
												__snake_case								=       mask.expand(self.model_tester.batch_size    ,      -1				).bool()
												__snake_case								=       bool_masked_pos.to(a__				)

									if return_labels:
												if model_class in [
												    *get_values(a__				),
												]:
															__snake_case								=       torch.zeros(
															    self.model_tester.batch_size    ,      dtype=torch.long    ,      device=a__				)

									return inputs_dict


						def        a			(self	:				Dict				):





									"""simple docstring"""



									self.config_tester.run_common_tests()


						@unittest.skip(reason='''VideoMAE does not use inputs_embeds'''				)
						def        a			(self	:				Tuple				):





									"""simple docstring"""



									pass


						def        a			(self	:				int				):





									"""simple docstring"""



									__snake_case	,     __snake_case								=       self.model_tester.prepare_config_and_inputs_for_common()

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


						def        a			(self	:				List[Any]				):





									"""simple docstring"""



									__snake_case	,     __snake_case								=       self.model_tester.prepare_config_and_inputs_for_common()

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

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


						def        a			(self	:				Dict				):





									"""simple docstring"""



									__snake_case								=       self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*a__				)


						def        a			(self	:				Any				):





									"""simple docstring"""



									__snake_case								=       self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_pretraining(*a__				)


						@slow
						def        a			(self	:				Optional[int]				):





									"""simple docstring"""



									for model_name in VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												__snake_case								=       VideoMAEModel.from_pretrained(a__				)
												self.assertIsNotNone(a__				)


						def        a			(self	:				List[str]				):





									"""simple docstring"""



									if not self.has_attentions:
												pass

									else:
												__snake_case	,     __snake_case								=       self.model_tester.prepare_config_and_inputs_for_common()
												__snake_case								=       True

												for model_class in self.all_model_classes:
															__snake_case								=       self.model_tester.seq_length - self.model_tester.num_masks
															__snake_case								=       (
															    num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length
															)

															__snake_case								=       True
															__snake_case								=       False
															__snake_case								=       True
															__snake_case								=       model_class(a__				)
															model.to(a__				)
															model.eval()
															with torch.no_grad():
																		__snake_case								=       model(**self._prepare_for_class(a__    ,      a__				)				)
															__snake_case								=       outputs.attentions
															self.assertEqual(len(a__				)    ,      self.model_tester.num_hidden_layers				)

															# check that output_attentions also work using config
															del inputs_dict["output_attentions"]
															__snake_case								=       True
															__snake_case								=       model_class(a__				)
															model.to(a__				)
															model.eval()
															with torch.no_grad():
																		__snake_case								=       model(**self._prepare_for_class(a__    ,      a__				)				)
															__snake_case								=       outputs.attentions
															self.assertEqual(len(a__				)    ,      self.model_tester.num_hidden_layers				)

															self.assertListEqual(
															    list(attentions[0].shape[-3:]				)    ,      [self.model_tester.num_attention_heads, seq_len, seq_len]    ,      )
															__snake_case								=       len(a__				)

															# Check attention is always last and order is fine
															__snake_case								=       True
															__snake_case								=       True
															__snake_case								=       model_class(a__				)
															model.to(a__				)
															model.eval()
															with torch.no_grad():
																		__snake_case								=       model(**self._prepare_for_class(a__    ,      a__				)				)

															self.assertEqual(out_len + 1    ,      len(a__				)				)

															__snake_case								=       outputs.attentions

															self.assertEqual(len(a__				)    ,      self.model_tester.num_hidden_layers				)
															self.assertListEqual(
															    list(self_attentions[0].shape[-3:]				)    ,      [self.model_tester.num_attention_heads, seq_len, seq_len]    ,      )


						def        a			(self	:				Any				):





									"""simple docstring"""



									def check_hidden_states_output(a__	:				Optional[int]    ,      a__	:				str    ,      a__	:				int				):
												__snake_case								=       model_class(a__				)
												model.to(a__				)
												model.eval()

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

												__snake_case								=       outputs.hidden_states
												__snake_case								=       self.model_tester.num_hidden_layers + 1
												self.assertEqual(len(a__				)    ,      a__				)

												__snake_case								=       self.model_tester.seq_length - self.model_tester.num_masks
												__snake_case								=       num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length

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

									__snake_case	,     __snake_case								=       self.model_tester.prepare_config_and_inputs_for_common()

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

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

												check_hidden_states_output(a__    ,      a__    ,      a__				)


						@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''				)
						def        a			(self	:				int				):





									"""simple docstring"""



									pass






def lowerCamelCase__   (							)				->    Union[str, Any]:
			__snake_case								=       hf_hub_download(
			    repo_id='''hf-internal-testing/spaghetti-video'''	,				filename='''eating_spaghetti.npy'''	,				repo_type='''dataset'''	)
			__snake_case								=       np.load(snake_case_	)
			return list(snake_case_	)







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


						@cached_property
						def        a			(self	:				int				):





									"""simple docstring"""



									return (
									    VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]    ,      image_std=[0.5, 0.5, 0.5]				)
									    if is_vision_available()
									    else None
									)


						@slow
						def        a			(self	:				Union[str, Any]				):





									"""simple docstring"""



									__snake_case								=       VideoMAEForVideoClassification.from_pretrained('''MCG-NJU/videomae-base-finetuned-kinetics'''				).to(
									    a__				)

									__snake_case								=       self.default_image_processor
									__snake_case								=       prepare_video()
									__snake_case								=       image_processor(a__    ,      return_tensors='''pt'''				).to(a__				)

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

									# verify the logits
									__snake_case								=       torch.Size((1, 400)				)
									self.assertEqual(outputs.logits.shape    ,      a__				)

									__snake_case								=       torch.tensor([0.3_6_6_9, -0.0_6_8_8, -0.2_4_2_1]				).to(a__				)

									self.assertTrue(torch.allclose(outputs.logits[0, :3]    ,      a__    ,      atol=1E-4				)				)


						@slow
						def        a			(self	:				Tuple				):





									"""simple docstring"""



									__snake_case								=       VideoMAEForPreTraining.from_pretrained('''MCG-NJU/videomae-base-short'''				).to(a__				)

									__snake_case								=       self.default_image_processor
									__snake_case								=       prepare_video()
									__snake_case								=       image_processor(a__    ,      return_tensors='''pt'''				).to(a__				)

									# add boolean mask, indicating which patches to mask
									__snake_case								=       hf_hub_download(repo_id='''hf-internal-testing/bool-masked-pos'''    ,      filename='''bool_masked_pos.pt'''				)
									__snake_case								=       torch.load(a__				)

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

									# verify the logits
									__snake_case								=       torch.Size([1, 1408, 1536]				)
									__snake_case								=       torch.tensor(
									    [[0.7_9_9_4, 0.9_6_1_2, 0.8_5_0_8], [0.7_4_0_1, 0.8_9_5_8, 0.8_3_0_2], [0.5_8_6_2, 0.7_4_6_8, 0.7_3_2_5]]    ,      device=a__				)
									self.assertEqual(outputs.logits.shape    ,      a__				)
									self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]    ,      a__    ,      atol=1E-4				)				)

									# verify the loss (`config.norm_pix_loss` = `True`)
									__snake_case								=       torch.tensor([0.5_1_4_2]    ,      device=a__				)
									self.assertTrue(torch.allclose(outputs.loss    ,      a__    ,      atol=1E-4				)				)

									# verify the loss (`config.norm_pix_loss` = `False`)
									__snake_case								=       VideoMAEForPreTraining.from_pretrained('''MCG-NJU/videomae-base-short'''    ,      norm_pix_loss=a__				).to(
									    a__				)

									with torch.no_grad():
												__snake_case								=       model(**a__				)

									__snake_case								=       torch.tensor(torch.tensor([0.6_4_6_9]				)    ,      device=a__				)
									self.assertTrue(torch.allclose(outputs.loss    ,      a__    ,      atol=1E-4				)				)








import gc
import unittest

from diffusers import FlaxControlNetModel, FlaxStableDiffusionControlNetPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
    import jax
    import jax.numpy as jnp
    from flax.jax_utils import replicate
    from flax.training.common_utils import shard

@slow
@require_flax
class 		SCREAMING_SNAKE_CASE__    (	unittest.TestCase  ):


      def        a			(self	:				int				):





         """simple docstring"""



         super().tearDown()
         gc.collect()


      def        a			(self	:				Dict				):





         """simple docstring"""



         __snake_case	,     __snake_case								=       FlaxControlNetModel.from_pretrained(
             '''lllyasviel/sd-controlnet-canny'''    ,      from_pt=a__    ,      dtype=jnp.bfloataa				)
         __snake_case	,     __snake_case								=       FlaxStableDiffusionControlNetPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''    ,      controlnet=a__    ,      from_pt=a__    ,      dtype=jnp.bfloataa				)
         __snake_case								=       controlnet_params

         __snake_case								=       '''bird'''
         __snake_case								=       jax.device_count()
         __snake_case								=       pipe.prepare_text_inputs([prompts] * num_samples				)

         __snake_case								=       load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'''				)
         __snake_case								=       pipe.prepare_image_inputs([canny_image] * num_samples				)

         __snake_case								=       jax.random.PRNGKey(0				)
         __snake_case								=       jax.random.split(a__    ,      jax.device_count()				)

         __snake_case								=       replicate(a__				)
         __snake_case								=       shard(a__				)
         __snake_case								=       shard(a__				)

         __snake_case								=       pipe(
             prompt_ids=a__    ,      image=a__    ,      params=a__    ,      prng_seed=a__    ,      num_inference_steps=50    ,      jit=a__    ,      ).images
         assert images.shape == (jax.device_count(), 1, 768, 512, 3)

         __snake_case								=       images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]				)
         __snake_case								=       images[0, 253:256, 253:256, -1]

         __snake_case								=       jnp.asarray(jax.device_get(image_slice.flatten()				)				)
         __snake_case								=       jnp.array(
             [0.1_6_7_9_6_9, 0.1_1_6_6_9_9, 0.0_8_1_5_4_3, 0.1_5_4_2_9_7, 0.1_3_2_8_1_2, 0.1_0_8_8_8_7, 0.1_6_9_9_2_2, 0.1_6_9_9_2_2, 0.2_0_5_0_7_8]				)
         print(f"""output_slice: {output_slice}"""				)
         assert jnp.abs(output_slice - expected_slice				).max() < 1E-2


      def        a			(self	:				Dict				):





         """simple docstring"""



         __snake_case	,     __snake_case								=       FlaxControlNetModel.from_pretrained(
             '''lllyasviel/sd-controlnet-openpose'''    ,      from_pt=a__    ,      dtype=jnp.bfloataa				)
         __snake_case	,     __snake_case								=       FlaxStableDiffusionControlNetPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''    ,      controlnet=a__    ,      from_pt=a__    ,      dtype=jnp.bfloataa				)
         __snake_case								=       controlnet_params

         __snake_case								=       '''Chef in the kitchen'''
         __snake_case								=       jax.device_count()
         __snake_case								=       pipe.prepare_text_inputs([prompts] * num_samples				)

         __snake_case								=       load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/pose.png'''				)
         __snake_case								=       pipe.prepare_image_inputs([pose_image] * num_samples				)

         __snake_case								=       jax.random.PRNGKey(0				)
         __snake_case								=       jax.random.split(a__    ,      jax.device_count()				)

         __snake_case								=       replicate(a__				)
         __snake_case								=       shard(a__				)
         __snake_case								=       shard(a__				)

         __snake_case								=       pipe(
             prompt_ids=a__    ,      image=a__    ,      params=a__    ,      prng_seed=a__    ,      num_inference_steps=50    ,      jit=a__    ,      ).images
         assert images.shape == (jax.device_count(), 1, 768, 512, 3)

         __snake_case								=       images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]				)
         __snake_case								=       images[0, 253:256, 253:256, -1]

         __snake_case								=       jnp.asarray(jax.device_get(image_slice.flatten()				)				)
         __snake_case								=       jnp.array(
             [[0.2_7_1_4_8_4, 0.2_6_1_7_1_9, 0.2_7_5_3_9_1, 0.2_7_7_3_4_4, 0.2_7_9_2_9_7, 0.2_9_1_0_1_6, 0.2_9_4_9_2_2, 0.3_0_2_7_3_4, 0.3_0_2_7_3_4]]				)
         print(f"""output_slice: {output_slice}"""				)
         assert jnp.abs(output_slice - expected_slice				).max() < 1E-2

