


"""simple docstring"""
import tempfile
import unittest

import torch

from diffusers import UNetaDConditionModel
from diffusers.training_utils import EMAModel
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device


enable_full_determinism()



class      _UpperCAmelCase					(				unittest.TestCase							):






		'''simple docstring'''

		lowerCamelCase__           ='hf-internal-testing/tiny-stable-diffusion-pipe'
		lowerCamelCase__           =1
		lowerCamelCase__           =77
		lowerCamelCase__           =32
		lowerCamelCase__           =4
		lowerCamelCase__           =lowerCamelCase__           =64
		lowerCamelCase__           =torch.manual_seed(0							)



		def 							SCREAMING_SNAKE_CASE    (self							,			a_=0.9999   ):





			'''simple docstring'''

			__snake_case      :		Optional[Any]         =							UNetaDConditionModel.from_pretrained(self.model_id							,			subfolder='''unet'''   )
			__snake_case      :		Union[str, Any]         =							unet.to(a_   )
			__snake_case      :		Tuple         =							EMAModel(unet.parameters()							,			decay=a_							,			model_cls=a_							,			model_config=unet.config   )
			return unet, ema_unet



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      :		Union[str, Any]         =							torch.randn(
			    self.batch_size							,			self.num_in_channels							,			self.latent_height							,			self.latent_width							,			generator=self.generator   ).to(a_   )
			__snake_case      :		int         =							torch.randint(0							,			10_00							,			size=(self.batch_size,)							,			generator=self.generator   ).to(a_   )
			__snake_case      :		int         =							torch.randn(
			    self.batch_size							,			self.prompt_length							,			self.text_encoder_hidden_dim							,			generator=self.generator   ).to(a_   )
			return noisy_latents, timesteps, encoder_hidden_states



		def 							SCREAMING_SNAKE_CASE    (self							,			a_   ):





			'''simple docstring'''

			__snake_case      :		List[Any]         =							{}
			for k, param in unet.state_dict().items():
				__snake_case      :		int         =							torch.randn_like(a_   ) + (param * torch.randn_like(a_   ))
				updated_state_dict.update({k: updated_param}   )
			unet.load_state_dict(a_   )
			return unet



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      ,				__snake_case      :		Optional[Any]         =							self.get_models()
			# Take the first (hypothetical) EMA step.
			ema_unet.step(unet.parameters()   )
			assert ema_unet.optimization_step == 1

			# Take two more.
			for _ in range(2   ):
				ema_unet.step(unet.parameters()   )
			assert ema_unet.optimization_step == 3



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      ,				__snake_case      :		Optional[Any]         =							self.get_models()
			# Since the `unet` is not being updated (i.e., backprop'd)
			# there won't be any difference between the `params` of `unet`
			# and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
			ema_unet.step(unet.parameters()   )
			__snake_case      :		Optional[Any]         =							list(unet.parameters()   )
			for s_param, param in zip(ema_unet.shadow_params							,			a_   ):
				assert torch.allclose(a_							,			a_   )

			# The above holds true even if we call `ema.step()` multiple times since
			# `unet` params are still not being updated.
			for _ in range(4   ):
				ema_unet.step(unet.parameters()   )
			for s_param, param in zip(ema_unet.shadow_params							,			a_   ):
				assert torch.allclose(a_							,			a_   )



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      ,				__snake_case      :		List[str]         =							self.get_models()
			# Here we simulate the parameter updates for `unet`. Since there might
			# be some parameters which are initialized to zero we take extra care to
			# initialize their values to something non-zero before the multiplication.
			__snake_case      :		Union[str, Any]         =							self.simulate_backprop(a_   )

			# Take the EMA step.
			ema_unet.step(unet_pseudo_updated_step_one.parameters()   )

			# Now the EMA'd parameters won't be equal to the original model parameters.
			__snake_case      :		Any         =							list(unet_pseudo_updated_step_one.parameters()   )
			for s_param, param in zip(ema_unet.shadow_params							,			a_   ):
				assert ~torch.allclose(a_							,			a_   )

			# Ensure this is the case when we take multiple EMA steps.
			for _ in range(4   ):
				ema_unet.step(unet.parameters()   )
			for s_param, param in zip(ema_unet.shadow_params							,			a_   ):
				assert ~torch.allclose(a_							,			a_   )



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      ,				__snake_case      :		Dict         =							self.get_models()

			# First backprop + EMA
			__snake_case      :		Optional[int]         =							self.simulate_backprop(a_   )
			ema_unet.step(unet_step_one.parameters()   )
			__snake_case      :		Union[str, Any]         =							ema_unet.shadow_params

			# Second backprop + EMA
			__snake_case      :		List[str]         =							self.simulate_backprop(a_   )
			ema_unet.step(unet_step_two.parameters()   )
			__snake_case      :		Tuple         =							ema_unet.shadow_params

			for step_one, step_two in zip(a_							,			a_   ):
				assert ~torch.allclose(a_							,			a_   )



		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      ,				__snake_case      :		Optional[int]         =							self.get_models(decay=0.0   )
			__snake_case      :		Any         =							self.simulate_backprop(a_   )
			ema_unet.step(unet_step_one.parameters()   )
			__snake_case      :		Any         =							ema_unet.shadow_params

			__snake_case      :		Optional[Any]         =							self.simulate_backprop(a_   )
			ema_unet.step(unet_step_two.parameters()   )
			__snake_case      :		List[str]         =							ema_unet.shadow_params

			for step_one, step_two in zip(a_							,			a_   ):
				assert torch.allclose(a_							,			a_   )



		@skip_mps
		def 							SCREAMING_SNAKE_CASE    (self   ):





			'''simple docstring'''

			__snake_case      ,				__snake_case      :		Optional[int]         =							self.get_models()
			__snake_case      ,				__snake_case      ,				__snake_case      :		Any         =							self.get_dummy_inputs()

			with tempfile.TemporaryDirectory() as tmpdir:
				ema_unet.save_pretrained(a_   )
				__snake_case      :		int         =							UNetaDConditionModel.from_pretrained(a_							,			model_cls=a_   )
				__snake_case      :		Union[str, Any]         =							loaded_unet.to(unet.device   )

			# Since no EMA step has been performed the outputs should match.
			__snake_case      :		List[str]         =							unet(a_							,			a_							,			a_   ).sample
			__snake_case      :		Tuple         =							loaded_unet(a_							,			a_							,			a_   ).sample

			assert torch.allclose(a_							,			a_							,			atol=1E-4   )










from ...configuration_utils import PretrainedConfig
from ...utils import logging


A__       =			logging.get_logger(__name__)

A__       =			{
    """MCG-NJU/videomae-base""": """https://huggingface.co/MCG-NJU/videomae-base/resolve/main/config.json""",
}





class    __lowerCAmelCase					(	lowerCamelCase__			):
				__lowerCamelCase					=							'''videomae'''




				def __init__(				self		,							_snake_case=224		,							_snake_case=16		,							_snake_case=3		,							_snake_case=16		,							_snake_case=2		,							_snake_case=768		,							_snake_case=12		,							_snake_case=12		,							_snake_case=3072		,							_snake_case="gelu"		,							_snake_case=0.0		,							_snake_case=0.0		,							_snake_case=0.02		,							_snake_case=1e-12		,							_snake_case=True		,							_snake_case=True		,							_snake_case=6		,							_snake_case=384		,							_snake_case=4		,							_snake_case=1536		,							_snake_case=True		,							**_snake_case		,							):



									"""simple docstring"""





									super().__init__(**_snake_case						)

									_lowerCAmelCase         =       image_size
									_lowerCAmelCase         =       patch_size
									_lowerCAmelCase         =       num_channels
									_lowerCAmelCase         =       num_frames
									_lowerCAmelCase         =       tubelet_size

									_lowerCAmelCase         =       hidden_size
									_lowerCAmelCase         =       num_hidden_layers
									_lowerCAmelCase         =       num_attention_heads
									_lowerCAmelCase         =       intermediate_size
									_lowerCAmelCase         =       hidden_act
									_lowerCAmelCase         =       hidden_dropout_prob
									_lowerCAmelCase         =       attention_probs_dropout_prob
									_lowerCAmelCase         =       initializer_range
									_lowerCAmelCase         =       layer_norm_eps
									_lowerCAmelCase         =       qkv_bias
									_lowerCAmelCase         =       use_mean_pooling

									_lowerCAmelCase         =       decoder_num_attention_heads
									_lowerCAmelCase         =       decoder_hidden_size
									_lowerCAmelCase         =       decoder_num_hidden_layers
									_lowerCAmelCase         =       decoder_intermediate_size
									_lowerCAmelCase         =       norm_pix_loss



