



import os
import tempfile
import unittest

import numpy as np
import torch

from diffusers.models import ModelMixin, UNetaDConditionModel
from diffusers.models.attention_processor import AttnProcessor, LoRAAttnProcessor
from diffusers.utils import (
    floats_tensor,
    logging,
    skip_mps,
    torch_device,
)
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()

__lowerCamelCase								=      logging.get_logger(__name__)







def  _snake_case			(							__snake_case       ,	__snake_case = True     )  ->  Any:






      '''simple docstring'''







      UpperCAmelCase_	:						Dict					   =		{}
      for name in model.attn_processors.keys():
            UpperCAmelCase_	:						Tuple					   =		name.endswith("attn2.processor"     ) and not (
                name.startswith("transformer_in"     ) or "temp_attentions" in name.split("."     )
            )
            UpperCAmelCase_	:						List[str]					   =		model.config.cross_attention_dim if has_cross_attention else None
            if name.startswith("mid_block"     ):
                  UpperCAmelCase_	:						str					   =		model.config.block_out_channels[-1]
            elif name.startswith("up_blocks"     ):
                  UpperCAmelCase_	:						Optional[Any]					   =		int(name[len("up_blocks."     )]     )
                  UpperCAmelCase_	:						int					   =		list(reversed(model.config.block_out_channels     )     )[block_id]
            elif name.startswith("down_blocks"     ):
                  UpperCAmelCase_	:						List[str]					   =		int(name[len("down_blocks."     )]     )
                  UpperCAmelCase_	:						Tuple					   =		model.config.block_out_channels[block_id]
            elif name.startswith("transformer_in"     ):
                  # Note that the `8 * ...` comes from: https://github.com/huggingface/diffusers/blob/7139f0e874f10b2463caa8cbd585762a309d12d6/src/diffusers/models/unet_3d_condition.py#L148
                  UpperCAmelCase_	:						Optional[int]					   =		8 * model.config.attention_head_dim

            UpperCAmelCase_	:						Union[str, Any]					   =		LoRAAttnProcessor(hidden_size=__snake_case       ,	cross_attention_dim=__snake_case     )
            UpperCAmelCase_	:						Tuple					   =		lora_attn_procs[name].to(model.device     )

            if mock_weights:
                  # add 1 to weights to mock trained weights
                  with torch.no_grad():
                        lora_attn_procs[name].to_q_lora.up.weight += 1
                        lora_attn_procs[name].to_k_lora.up.weight += 1
                        lora_attn_procs[name].to_v_lora.up.weight += 1
                        lora_attn_procs[name].to_out_lora.up.weight += 1

      return lora_attn_procs





@skip_mps
class 	snake_case_			(lowercase__			,			lowercase__			,			unittest.TestCase	):




       """simple docstring"""






       _lowerCamelCase           =					UNetaDConditionModel
       _lowerCamelCase           =					"""sample"""



       @property
       def 			A_ (  self):

             """simple docstring"""





             UpperCAmelCase_	:						Union[str, Any]					   =		4
             UpperCAmelCase_	:						Tuple					   =		4
             UpperCAmelCase_	:						Optional[Any]					   =		4
             UpperCAmelCase_	:						Optional[int]					   =		(32, 32)

             UpperCAmelCase_	:						Any					   =		floats_tensor((batch_size, num_channels, num_frames) + sizes).to(lowercase)
             UpperCAmelCase_	:						Optional[Any]					   =		torch.tensor([10]).to(lowercase)
             UpperCAmelCase_	:						str					   =		floats_tensor((batch_size, 4, 32)).to(lowercase)

             return {"sample": noise, "timestep": time_step, "encoder_hidden_states": encoder_hidden_states}



       @property
       def 			A_ (  self):

             """simple docstring"""





             return (4, 4, 32, 32)



       @property
       def 			A_ (  self):

             """simple docstring"""





             return (4, 4, 32, 32)



       def 			A_ (  self):

             """simple docstring"""





             UpperCAmelCase_	:						Union[str, Any]					   =		{
                 "block_out_channels": (32, 64),
                 "down_block_types": (
                     "CrossAttnDownBlock3D",
                     "DownBlock3D",
                 ),
                 "up_block_types": ("UpBlock3D", "CrossAttnUpBlock3D"),
                 "cross_attention_dim": 32,
                 "attention_head_dim": 8,
                 "out_channels": 4,
                 "in_channels": 4,
                 "layers_per_block": 1,
                 "sample_size": 32,
             }
             UpperCAmelCase_	:						Any					   =		self.dummy_input
             return init_dict, inputs_dict



       @unittest.skipIf(
           torch_device != "cuda" or not is_xformers_available()  ,reason="XFormers attention is only available with CUDA and `xformers` installed"  ,)
       def 			A_ (  self):

             """simple docstring"""





             UpperCAmelCase_   ,   UpperCAmelCase_	:						List[Any]					   =		self.prepare_init_args_and_inputs_for_common()
             UpperCAmelCase_	:						Optional[Any]					   =		self.model_class(**lowercase)

             model.enable_xformers_memory_efficient_attention()

             assert (
                 model.mid_block.attentions[0].transformer_blocks[0].attna.processor.__class__.__name__
                 == "XFormersAttnProcessor"
             ), "xformers is not enabled"



       def 			A_ (  self):

             """simple docstring"""





             UpperCAmelCase_   ,   UpperCAmelCase_	:						Union[str, Any]					   =		self.prepare_init_args_and_inputs_for_common()

             UpperCAmelCase_	:						Dict					   =		32

             UpperCAmelCase_	:						Optional[int]					   =		self.model_class(**lowercase)
             model.to(lowercase)
             model.eval()

             with torch.no_grad():
                   UpperCAmelCase_	:						int					   =		model(**lowercase)

                   if isinstance(lowercase  ,lowercase):
                         UpperCAmelCase_	:						List[Any]					   =		output.sample

             self.assertIsNotNone(lowercase)
             UpperCAmelCase_	:						Any					   =		inputs_dict["sample"].shape
             self.assertEqual(output.shape  ,lowercase  ,"Input and output shapes do not match")



       def 			A_ (  self):

             """simple docstring"""





             UpperCAmelCase_   ,   UpperCAmelCase_	:						Optional[int]					   =		self.prepare_init_args_and_inputs_for_common()
             UpperCAmelCase_	:						List[Any]					   =		self.model_class(**lowercase)
             model.to(lowercase)
             model.eval()

             with torch.no_grad():
                   # Warmup pass when using mps (see #372)
                   if torch_device == "mps" and isinstance(lowercase  ,lowercase):
                         model(**self.dummy_input)

                   UpperCAmelCase_	:						Optional[int]					   =		model(**lowercase)
                   if isinstance(lowercase  ,lowercase):
                         UpperCAmelCase_	:						Optional[Any]					   =		first.sample

                   UpperCAmelCase_	:						List[Any]					   =		model(**lowercase)
                   if isinstance(lowercase  ,lowercase):
                         UpperCAmelCase_	:						List[str]					   =		second.sample

             UpperCAmelCase_	:						Optional[int]					   =		first.cpu().numpy()
             UpperCAmelCase_	:						Tuple					   =		second.cpu().numpy()
             UpperCAmelCase_	:						List[Any]					   =		out_a[~np.isnan(lowercase)]
             UpperCAmelCase_	:						str					   =		out_a[~np.isnan(lowercase)]
             UpperCAmelCase_	:						List[str]					   =		np.amax(np.abs(out_a - out_a))
             self.assertLessEqual(lowercase  ,1E-5)



       def 			A_ (  self):

             """simple docstring"""





             UpperCAmelCase_   ,   UpperCAmelCase_	:						Tuple					   =		self.prepare_init_args_and_inputs_for_common()

             UpperCAmelCase_	:						Dict					   =		8

             UpperCAmelCase_	:						Optional[Any]					   =		self.model_class(**lowercase)
             model.to(lowercase)
             model.eval()

             model.set_attention_slice("auto")
             with torch.no_grad():
                   UpperCAmelCase_	:						Union[str, Any]					   =		model(**lowercase)
             assert output is not None

             model.set_attention_slice("max")
             with torch.no_grad():
                   UpperCAmelCase_	:						List[str]					   =		model(**lowercase)
             assert output is not None

             model.set_attention_slice(2)
             with torch.no_grad():
                   UpperCAmelCase_	:						List[Any]					   =		model(**lowercase)
             assert output is not None



       def 			A_ (  self):

             """simple docstring"""





             UpperCAmelCase_   ,   UpperCAmelCase_	:						int					   =		self.prepare_init_args_and_inputs_for_common()

             UpperCAmelCase_	:						List[str]					   =		8

             UpperCAmelCase_	:						Dict					   =		self.model_class(**lowercase)
             model.to(lowercase)

             with torch.no_grad():
                   UpperCAmelCase_	:						int					   =		model(**lowercase).sample

             UpperCAmelCase_	:						Any					   =		create_lora_layers(lowercase)

             # make sure we can set a list of attention processors
             model.set_attn_processor(lowercase)
             model.to(lowercase)

             # test that attn processors can be set to itself
             model.set_attn_processor(model.attn_processors)

             with torch.no_grad():
                   UpperCAmelCase_	:						Tuple					   =		model(**lowercase  ,cross_attention_kwargs={"scale": 0.0}).sample
                   UpperCAmelCase_	:						Dict					   =		model(**lowercase  ,cross_attention_kwargs={"scale": 0.5}).sample
                   UpperCAmelCase_	:						int					   =		model(**lowercase  ,cross_attention_kwargs={"scale": 0.5}).sample

             assert (samplea - samplea).abs().max() < 3E-3
             assert (samplea - samplea).abs().max() < 3E-3

             # sample 2 and sample 3 should be different
             assert (samplea - samplea).abs().max() > 3E-3



       def 			A_ (  self):

             """simple docstring"""





             UpperCAmelCase_   ,   UpperCAmelCase_	:						Optional[int]					   =		self.prepare_init_args_and_inputs_for_common()

             UpperCAmelCase_	:						List[Any]					   =		8

             torch.manual_seed(0)
             UpperCAmelCase_	:						List[Any]					   =		self.model_class(**lowercase)
             model.to(lowercase)

             with torch.no_grad():
                   UpperCAmelCase_	:						str					   =		model(**lowercase).sample

             UpperCAmelCase_	:						Tuple					   =		create_lora_layers(lowercase)
             model.set_attn_processor(lowercase)

             with torch.no_grad():
                   UpperCAmelCase_	:						Any					   =		model(**lowercase  ,cross_attention_kwargs={"scale": 0.5}).sample

             with tempfile.TemporaryDirectory() as tmpdirname:
                   model.save_attn_procs(lowercase)
                   self.assertTrue(os.path.isfile(os.path.join(lowercase  ,"pytorch_lora_weights.bin")))
                   torch.manual_seed(0)
                   UpperCAmelCase_	:						Optional[Any]					   =		self.model_class(**lowercase)
                   new_model.to(lowercase)
                   new_model.load_attn_procs(lowercase)

             with torch.no_grad():
                   UpperCAmelCase_	:						List[Any]					   =		new_model(**lowercase  ,cross_attention_kwargs={"scale": 0.5}).sample

             assert (sample - new_sample).abs().max() < 1E-3

             # LoRA and no LoRA should NOT be the same
             assert (sample - old_sample).abs().max() > 1E-4



       def 			A_ (  self):

             """simple docstring"""





             UpperCAmelCase_   ,   UpperCAmelCase_	:						Dict					   =		self.prepare_init_args_and_inputs_for_common()

             UpperCAmelCase_	:						str					   =		8

             torch.manual_seed(0)
             UpperCAmelCase_	:						Union[str, Any]					   =		self.model_class(**lowercase)
             model.to(lowercase)

             with torch.no_grad():
                   UpperCAmelCase_	:						Tuple					   =		model(**lowercase).sample

             UpperCAmelCase_	:						Dict					   =		create_lora_layers(lowercase)
             model.set_attn_processor(lowercase)

             with torch.no_grad():
                   UpperCAmelCase_	:						Any					   =		model(**lowercase  ,cross_attention_kwargs={"scale": 0.5}).sample

             with tempfile.TemporaryDirectory() as tmpdirname:
                   model.save_attn_procs(lowercase  ,safe_serialization=lowercase)
                   self.assertTrue(os.path.isfile(os.path.join(lowercase  ,"pytorch_lora_weights.safetensors")))
                   torch.manual_seed(0)
                   UpperCAmelCase_	:						Any					   =		self.model_class(**lowercase)
                   new_model.to(lowercase)
                   new_model.load_attn_procs(lowercase)

             with torch.no_grad():
                   UpperCAmelCase_	:						Optional[int]					   =		new_model(**lowercase  ,cross_attention_kwargs={"scale": 0.5}).sample

             assert (sample - new_sample).abs().max() < 3E-3

             # LoRA and no LoRA should NOT be the same
             assert (sample - old_sample).abs().max() > 1E-4



       def 			A_ (  self):

             """simple docstring"""





             UpperCAmelCase_   ,   UpperCAmelCase_	:						str					   =		self.prepare_init_args_and_inputs_for_common()

             UpperCAmelCase_	:						Dict					   =		8

             torch.manual_seed(0)
             UpperCAmelCase_	:						Optional[int]					   =		self.model_class(**lowercase)
             model.to(lowercase)

             UpperCAmelCase_	:						int					   =		create_lora_layers(lowercase  ,mock_weights=lowercase)
             model.set_attn_processor(lowercase)
             # Saving as torch, properly reloads with directly filename
             with tempfile.TemporaryDirectory() as tmpdirname:
                   model.save_attn_procs(lowercase)
                   self.assertTrue(os.path.isfile(os.path.join(lowercase  ,"pytorch_lora_weights.bin")))
                   torch.manual_seed(0)
                   UpperCAmelCase_	:						Optional[Any]					   =		self.model_class(**lowercase)
                   new_model.to(lowercase)
                   new_model.load_attn_procs(lowercase  ,weight_name="pytorch_lora_weights.bin")



       def 			A_ (  self):

             """simple docstring"""





             UpperCAmelCase_   ,   UpperCAmelCase_	:						Optional[Any]					   =		self.prepare_init_args_and_inputs_for_common()

             UpperCAmelCase_	:						str					   =		8

             torch.manual_seed(0)
             UpperCAmelCase_	:						List[Any]					   =		self.model_class(**lowercase)
             model.to(lowercase)

             UpperCAmelCase_	:						Tuple					   =		create_lora_layers(lowercase  ,mock_weights=lowercase)
             model.set_attn_processor(lowercase)
             # Saving as torch, properly reloads with directly filename
             with tempfile.TemporaryDirectory() as tmpdirname:
                   model.save_attn_procs(lowercase)
                   self.assertTrue(os.path.isfile(os.path.join(lowercase  ,"pytorch_lora_weights.bin")))
                   torch.manual_seed(0)
                   UpperCAmelCase_	:						Dict					   =		self.model_class(**lowercase)
                   new_model.to(lowercase)
                   with self.assertRaises(lowercase) as e:
                         new_model.load_attn_procs(lowercase  ,use_safetensors=lowercase)
                   self.assertIn("Error no file named pytorch_lora_weights.safetensors"  ,str(e.exception))



       def 			A_ (  self):

             """simple docstring"""





             UpperCAmelCase_   ,   UpperCAmelCase_	:						int					   =		self.prepare_init_args_and_inputs_for_common()

             UpperCAmelCase_	:						List[Any]					   =		8

             torch.manual_seed(0)
             UpperCAmelCase_	:						Union[str, Any]					   =		self.model_class(**lowercase)
             model.to(lowercase)

             with torch.no_grad():
                   UpperCAmelCase_	:						Dict					   =		model(**lowercase).sample

             UpperCAmelCase_	:						Optional[int]					   =		create_lora_layers(lowercase)
             model.set_attn_processor(lowercase)

             with torch.no_grad():
                   UpperCAmelCase_	:						List[Any]					   =		model(**lowercase  ,cross_attention_kwargs={"scale": 0.0}).sample

             model.set_attn_processor(AttnProcessor())

             with torch.no_grad():
                   UpperCAmelCase_	:						str					   =		model(**lowercase).sample

             assert (sample - new_sample).abs().max() < 1E-4
             assert (sample - old_sample).abs().max() < 3E-3



       @unittest.skipIf(
           torch_device != "cuda" or not is_xformers_available()  ,reason="XFormers attention is only available with CUDA and `xformers` installed"  ,)
       def 			A_ (  self):

             """simple docstring"""





             UpperCAmelCase_   ,   UpperCAmelCase_	:						Tuple					   =		self.prepare_init_args_and_inputs_for_common()

             UpperCAmelCase_	:						Optional[Any]					   =		4

             torch.manual_seed(0)
             UpperCAmelCase_	:						List[Any]					   =		self.model_class(**lowercase)
             model.to(lowercase)
             UpperCAmelCase_	:						str					   =		create_lora_layers(lowercase)
             model.set_attn_processor(lowercase)

             # default
             with torch.no_grad():
                   UpperCAmelCase_	:						int					   =		model(**lowercase).sample

                   model.enable_xformers_memory_efficient_attention()
                   UpperCAmelCase_	:						Any					   =		model(**lowercase).sample

                   model.disable_xformers_memory_efficient_attention()
                   UpperCAmelCase_	:						int					   =		model(**lowercase).sample

             assert (sample - on_sample).abs().max() < 1E-4
             assert (sample - off_sample).abs().max() < 1E-4



       def 			A_ (  self):

             """simple docstring"""





             UpperCAmelCase_   ,   UpperCAmelCase_	:						Optional[Any]					   =		self.prepare_init_args_and_inputs_for_common()
             UpperCAmelCase_	:						int					   =		32

             UpperCAmelCase_	:						Optional[Any]					   =		self.model_class(**lowercase)
             model.to(lowercase)
             model.eval()

             with torch.no_grad():
                   UpperCAmelCase_	:						Tuple					   =		model(**lowercase)[0]

             model.enable_forward_chunking()
             with torch.no_grad():
                   UpperCAmelCase_	:						str					   =		model(**lowercase)[0]

             self.assertEqual(output.shape  ,output_a.shape  ,"Shape doesn't match")
             assert np.abs(output.cpu() - output_a.cpu()).max() < 1E-2


# (todo: sayakpaul) implement SLOW tests.





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


__lowerCamelCase								=      logging.get_logger(__name__)

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


class 	snake_case_			(lowercase__	):




       """simple docstring"""






       _lowerCamelCase           =					"""videomae"""



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

             """simple docstring"""





             super().__init__(**lowercase)

             UpperCAmelCase_	:						Optional[int]					   =		image_size
             UpperCAmelCase_	:						Optional[Any]					   =		patch_size
             UpperCAmelCase_	:						Tuple					   =		num_channels
             UpperCAmelCase_	:						Optional[int]					   =		num_frames
             UpperCAmelCase_	:						str					   =		tubelet_size

             UpperCAmelCase_	:						Any					   =		hidden_size
             UpperCAmelCase_	:						List[str]					   =		num_hidden_layers
             UpperCAmelCase_	:						List[str]					   =		num_attention_heads
             UpperCAmelCase_	:						Union[str, Any]					   =		intermediate_size
             UpperCAmelCase_	:						List[str]					   =		hidden_act
             UpperCAmelCase_	:						Tuple					   =		hidden_dropout_prob
             UpperCAmelCase_	:						Optional[Any]					   =		attention_probs_dropout_prob
             UpperCAmelCase_	:						Optional[Any]					   =		initializer_range
             UpperCAmelCase_	:						Dict					   =		layer_norm_eps
             UpperCAmelCase_	:						Tuple					   =		qkv_bias
             UpperCAmelCase_	:						int					   =		use_mean_pooling

             UpperCAmelCase_	:						List[Any]					   =		decoder_num_attention_heads
             UpperCAmelCase_	:						Optional[int]					   =		decoder_hidden_size
             UpperCAmelCase_	:						Any					   =		decoder_num_hidden_layers
             UpperCAmelCase_	:						Dict					   =		decoder_intermediate_size
             UpperCAmelCase_	:						str					   =		norm_pix_loss
