
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()

_A							=       logging.get_logger(__name__)




def 					lowerCamelCase__					(    a__	:	Dict						,   a__	:	bool = True  )		->						List[str]:
       UpperCamelCase_											=							{}
       for name in model.attn_processors.keys():
              UpperCamelCase_											=							name.endswith("""attn2.processor"""  ) and not (
                  name.startswith("""transformer_in"""  ) or """temp_attentions""" in name.split("""."""  )
              )
              UpperCamelCase_											=							model.config.cross_attention_dim if has_cross_attention else None
              if name.startswith("""mid_block"""  ):
                     UpperCamelCase_											=							model.config.block_out_channels[-1]
              elif name.startswith("""up_blocks"""  ):
                     UpperCamelCase_											=							int(name[len("""up_blocks."""  )]  )
                     UpperCamelCase_											=							list(reversed(model.config.block_out_channels  )  )[block_id]
              elif name.startswith("""down_blocks"""  ):
                     UpperCamelCase_											=							int(name[len("""down_blocks."""  )]  )
                     UpperCamelCase_											=							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_											=							8 * model.config.attention_head_dim

              UpperCamelCase_											=							LoRAAttnProcessor(hidden_size=a_						,   cross_attention_dim=a_  )
              UpperCamelCase_											=							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 		lowercase_							(							__SCREAMING_SNAKE_CASE				,						__SCREAMING_SNAKE_CASE				,						unittest.TestCase					):
 A__    :					List[Any]    =     UNetaDConditionModel
 A__    :					Optional[int]    =     "sample"

 @property
 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_											=							4
        UpperCamelCase_											=							4
        UpperCamelCase_											=							4
        UpperCamelCase_											=							(3_2, 3_2)

        UpperCamelCase_											=							floats_tensor((batch_size, num_channels, num_frames) + sizes  ).to(__UpperCamelCase  )
        UpperCamelCase_											=							torch.tensor([1_0]  ).to(__UpperCamelCase  )
        UpperCamelCase_											=							floats_tensor((batch_size, 4, 3_2)  ).to(__UpperCamelCase  )

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

 @property
 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        return (4, 4, 3_2, 3_2)

 @property
 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        return (4, 4, 3_2, 3_2)

 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_											=							{
            """block_out_channels""": (3_2, 6_4),
            """down_block_types""": (
                """CrossAttnDownBlock3D""",
                """DownBlock3D""",
            ),
            """up_block_types""": ("""UpBlock3D""", """CrossAttnUpBlock3D"""),
            """cross_attention_dim""": 3_2,
            """attention_head_dim""": 8,
            """out_channels""": 4,
            """in_channels""": 4,
            """layers_per_block""": 1,
            """sample_size""": 3_2,
        }
        UpperCamelCase_											=							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 lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_							,					UpperCamelCase_											=							self.prepare_init_args_and_inputs_for_common()
        UpperCamelCase_											=							self.model_class(**__UpperCamelCase  )

        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 lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_							,					UpperCamelCase_											=							self.prepare_init_args_and_inputs_for_common()

        UpperCamelCase_											=							3_2

        UpperCamelCase_											=							self.model_class(**__UpperCamelCase  )
        model.to(__UpperCamelCase  )
        model.eval()

        with torch.no_grad():
               UpperCamelCase_											=							model(**__UpperCamelCase  )

               if isinstance(__UpperCamelCase  ,						__UpperCamelCase  ):
                      UpperCamelCase_											=							output.sample

        self.assertIsNotNone(__UpperCamelCase  )
        UpperCamelCase_											=							inputs_dict["""sample"""].shape
        self.assertEqual(output.shape  ,						__UpperCamelCase  ,						"""Input and output shapes do not match"""  )

 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_							,					UpperCamelCase_											=							self.prepare_init_args_and_inputs_for_common()
        UpperCamelCase_											=							self.model_class(**__UpperCamelCase  )
        model.to(__UpperCamelCase  )
        model.eval()

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

               UpperCamelCase_											=							model(**__UpperCamelCase  )
               if isinstance(__UpperCamelCase  ,						__UpperCamelCase  ):
                      UpperCamelCase_											=							first.sample

               UpperCamelCase_											=							model(**__UpperCamelCase  )
               if isinstance(__UpperCamelCase  ,						__UpperCamelCase  ):
                      UpperCamelCase_											=							second.sample

        UpperCamelCase_											=							first.cpu().numpy()
        UpperCamelCase_											=							second.cpu().numpy()
        UpperCamelCase_											=							out_a[~np.isnan(__UpperCamelCase  )]
        UpperCamelCase_											=							out_a[~np.isnan(__UpperCamelCase  )]
        UpperCamelCase_											=							np.amax(np.abs(out_a - out_a  )  )
        self.assertLessEqual(__UpperCamelCase  ,						1e-5  )

 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_							,					UpperCamelCase_											=							self.prepare_init_args_and_inputs_for_common()

        UpperCamelCase_											=							8

        UpperCamelCase_											=							self.model_class(**__UpperCamelCase  )
        model.to(__UpperCamelCase  )
        model.eval()

        model.set_attention_slice("""auto"""  )
        with torch.no_grad():
               UpperCamelCase_											=							model(**__UpperCamelCase  )
        assert output is not None

        model.set_attention_slice("""max"""  )
        with torch.no_grad():
               UpperCamelCase_											=							model(**__UpperCamelCase  )
        assert output is not None

        model.set_attention_slice(2  )
        with torch.no_grad():
               UpperCamelCase_											=							model(**__UpperCamelCase  )
        assert output is not None

 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_							,					UpperCamelCase_											=							self.prepare_init_args_and_inputs_for_common()

        UpperCamelCase_											=							8

        UpperCamelCase_											=							self.model_class(**__UpperCamelCase  )
        model.to(__UpperCamelCase  )

        with torch.no_grad():
               UpperCamelCase_											=							model(**__UpperCamelCase  ).sample

        UpperCamelCase_											=							create_lora_layers(__UpperCamelCase  )

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

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

        with torch.no_grad():
               UpperCamelCase_											=							model(**__UpperCamelCase  ,						cross_attention_kwargs={"""scale""": 0.0}  ).sample
               UpperCamelCase_											=							model(**__UpperCamelCase  ,						cross_attention_kwargs={"""scale""": 0.5}  ).sample
               UpperCamelCase_											=							model(**__UpperCamelCase  ,						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 lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_							,					UpperCamelCase_											=							self.prepare_init_args_and_inputs_for_common()

        UpperCamelCase_											=							8

        torch.manual_seed(0  )
        UpperCamelCase_											=							self.model_class(**__UpperCamelCase  )
        model.to(__UpperCamelCase  )

        with torch.no_grad():
               UpperCamelCase_											=							model(**__UpperCamelCase  ).sample

        UpperCamelCase_											=							create_lora_layers(__UpperCamelCase  )
        model.set_attn_processor(__UpperCamelCase  )

        with torch.no_grad():
               UpperCamelCase_											=							model(**__UpperCamelCase  ,						cross_attention_kwargs={"""scale""": 0.5}  ).sample

        with tempfile.TemporaryDirectory() as tmpdirname:
               model.save_attn_procs(__UpperCamelCase  )
               self.assertTrue(os.path.isfile(os.path.join(__UpperCamelCase  ,						"""pytorch_lora_weights.bin"""  )  )  )
               torch.manual_seed(0  )
               UpperCamelCase_											=							self.model_class(**__UpperCamelCase  )
               new_model.to(__UpperCamelCase  )
               new_model.load_attn_procs(__UpperCamelCase  )

        with torch.no_grad():
               UpperCamelCase_											=							new_model(**__UpperCamelCase  ,						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 lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_							,					UpperCamelCase_											=							self.prepare_init_args_and_inputs_for_common()

        UpperCamelCase_											=							8

        torch.manual_seed(0  )
        UpperCamelCase_											=							self.model_class(**__UpperCamelCase  )
        model.to(__UpperCamelCase  )

        with torch.no_grad():
               UpperCamelCase_											=							model(**__UpperCamelCase  ).sample

        UpperCamelCase_											=							create_lora_layers(__UpperCamelCase  )
        model.set_attn_processor(__UpperCamelCase  )

        with torch.no_grad():
               UpperCamelCase_											=							model(**__UpperCamelCase  ,						cross_attention_kwargs={"""scale""": 0.5}  ).sample

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

        with torch.no_grad():
               UpperCamelCase_											=							new_model(**__UpperCamelCase  ,						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 lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_							,					UpperCamelCase_											=							self.prepare_init_args_and_inputs_for_common()

        UpperCamelCase_											=							8

        torch.manual_seed(0  )
        UpperCamelCase_											=							self.model_class(**__UpperCamelCase  )
        model.to(__UpperCamelCase  )

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

 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_							,					UpperCamelCase_											=							self.prepare_init_args_and_inputs_for_common()

        UpperCamelCase_											=							8

        torch.manual_seed(0  )
        UpperCamelCase_											=							self.model_class(**__UpperCamelCase  )
        model.to(__UpperCamelCase  )

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

 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_							,					UpperCamelCase_											=							self.prepare_init_args_and_inputs_for_common()

        UpperCamelCase_											=							8

        torch.manual_seed(0  )
        UpperCamelCase_											=							self.model_class(**__UpperCamelCase  )
        model.to(__UpperCamelCase  )

        with torch.no_grad():
               UpperCamelCase_											=							model(**__UpperCamelCase  ).sample

        UpperCamelCase_											=							create_lora_layers(__UpperCamelCase  )
        model.set_attn_processor(__UpperCamelCase  )

        with torch.no_grad():
               UpperCamelCase_											=							model(**__UpperCamelCase  ,						cross_attention_kwargs={"""scale""": 0.0}  ).sample

        model.set_attn_processor(AttnProcessor()  )

        with torch.no_grad():
               UpperCamelCase_											=							model(**__UpperCamelCase  ).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 lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_							,					UpperCamelCase_											=							self.prepare_init_args_and_inputs_for_common()

        UpperCamelCase_											=							4

        torch.manual_seed(0  )
        UpperCamelCase_											=							self.model_class(**__UpperCamelCase  )
        model.to(__UpperCamelCase  )
        UpperCamelCase_											=							create_lora_layers(__UpperCamelCase  )
        model.set_attn_processor(__UpperCamelCase  )

        # default
        with torch.no_grad():
               UpperCamelCase_											=							model(**__UpperCamelCase  ).sample

               model.enable_xformers_memory_efficient_attention()
               UpperCamelCase_											=							model(**__UpperCamelCase  ).sample

               model.disable_xformers_memory_efficient_attention()
               UpperCamelCase_											=							model(**__UpperCamelCase  ).sample

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

 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_							,					UpperCamelCase_											=							self.prepare_init_args_and_inputs_for_common()
        UpperCamelCase_											=							3_2

        UpperCamelCase_											=							self.model_class(**__UpperCamelCase  )
        model.to(__UpperCamelCase  )
        model.eval()

        with torch.no_grad():
               UpperCamelCase_											=							model(**__UpperCamelCase  )[0]

        model.enable_forward_chunking()
        with torch.no_grad():
               UpperCamelCase_											=							model(**__UpperCamelCase  )[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 collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


A__							=				logging.get_logger(__name__)

A__							=				{
    "nghuyong/ernie-1.0-base-zh": "https://huggingface.co/nghuyong/ernie-1.0-base-zh/resolve/main/config.json",
    "nghuyong/ernie-2.0-base-en": "https://huggingface.co/nghuyong/ernie-2.0-base-en/resolve/main/config.json",
    "nghuyong/ernie-2.0-large-en": "https://huggingface.co/nghuyong/ernie-2.0-large-en/resolve/main/config.json",
    "nghuyong/ernie-3.0-base-zh": "https://huggingface.co/nghuyong/ernie-3.0-base-zh/resolve/main/config.json",
    "nghuyong/ernie-3.0-medium-zh": "https://huggingface.co/nghuyong/ernie-3.0-medium-zh/resolve/main/config.json",
    "nghuyong/ernie-3.0-mini-zh": "https://huggingface.co/nghuyong/ernie-3.0-mini-zh/resolve/main/config.json",
    "nghuyong/ernie-3.0-micro-zh": "https://huggingface.co/nghuyong/ernie-3.0-micro-zh/resolve/main/config.json",
    "nghuyong/ernie-3.0-nano-zh": "https://huggingface.co/nghuyong/ernie-3.0-nano-zh/resolve/main/config.json",
    "nghuyong/ernie-gram-zh": "https://huggingface.co/nghuyong/ernie-gram-zh/resolve/main/config.json",
    "nghuyong/ernie-health-zh": "https://huggingface.co/nghuyong/ernie-health-zh/resolve/main/config.json",
}

class 					__UpperCamelCase					(       SCREAMING_SNAKE_CASE      ):
	_lowercase	: Dict												=	"ernie"







	def __init__( self:			List[Any]	,			__UpperCamelCase:			Dict=3_05_22	,			__UpperCamelCase:			str=7_68	,			__UpperCamelCase:			Tuple=12	,			__UpperCamelCase:			int=12	,			__UpperCamelCase:			Optional[Any]=30_72	,			__UpperCamelCase:			str="gelu"	,			__UpperCamelCase:			Any=0.1	,			__UpperCamelCase:			int=0.1	,			__UpperCamelCase:			int=5_12	,			__UpperCamelCase:			Any=2	,			__UpperCamelCase:			Tuple=3	,			__UpperCamelCase:			Any=False	,			__UpperCamelCase:			Optional[int]=0.02	,			__UpperCamelCase:			List[Any]=1E-12	,			__UpperCamelCase:			Union[str, Any]=0	,			__UpperCamelCase:			Tuple="absolute"	,			__UpperCamelCase:			Union[str, Any]=True	,			__UpperCamelCase:			Optional[Any]=None	,			**__UpperCamelCase:			List[Any]	,			):

						'''simple docstring'''





						super().__init__(pad_token_id=__UpperCamelCase	,			**__UpperCamelCase   )

						__magic_name__				=     vocab_size
						__magic_name__				=     hidden_size
						__magic_name__				=     num_hidden_layers
						__magic_name__				=     num_attention_heads
						__magic_name__				=     hidden_act
						__magic_name__				=     intermediate_size
						__magic_name__				=     hidden_dropout_prob
						__magic_name__				=     attention_probs_dropout_prob
						__magic_name__				=     max_position_embeddings
						__magic_name__				=     type_vocab_size
						__magic_name__				=     task_type_vocab_size
						__magic_name__				=     use_task_id
						__magic_name__				=     initializer_range
						__magic_name__				=     layer_norm_eps
						__magic_name__				=     position_embedding_type
						__magic_name__				=     use_cache
						__magic_name__				=     classifier_dropout




class 					__UpperCamelCase					(       SCREAMING_SNAKE_CASE      ):







	@property
	def 						_SCREAMING_SNAKE_CASE				( self:			Tuple   ):

						'''simple docstring'''





						if self.task == "multiple-choice":
											__magic_name__				=     {0: 'batch', 1: 'choice', 2: 'sequence'}
						else:
											__magic_name__				=     {0: 'batch', 1: 'sequence'}
						return OrderedDict(
						    [
						        ('input_ids', dynamic_axis),
						        ('attention_mask', dynamic_axis),
						        ('token_type_ids', dynamic_axis),
						        ('task_type_ids', dynamic_axis),
						    ]   )




