




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 							lowerCamelCase    (							UpperCAmelCase_	:      Tuple   ,       UpperCAmelCase_	:      bool = True      )->			Union[str, Any]:

       """simple docstring"""




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

              a			   =LoRAAttnProcessor(hidden_size=UpperCAmelCase_   ,       cross_attention_dim=UpperCAmelCase_      )
              a			   =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        UpperCAmelCase__				(							UpperCAmelCase__      ,		UpperCAmelCase__      ,		unittest.TestCase				):



     '''simple docstring'''
     _SCREAMING_SNAKE_CASE      :					Union[str, Any]     					=    UNetaDConditionModel
     _SCREAMING_SNAKE_CASE      :					Any     					=    "sample"
     @property
     def        lowerCAmelCase__  (				self	):
            a			   =4
            a			   =4
            a			   =4
            a			   =(32, 32)

            a			   =floats_tensor((batch_size, num_channels, num_frames) + sizes	).to(_lowerCAmelCase	)
            a			   =torch.tensor([10]	).to(_lowerCAmelCase	)
            a			   =floats_tensor((batch_size, 4, 32)	).to(_lowerCAmelCase	)

            return {"sample": noise, "timestep": time_step, "encoder_hidden_states": encoder_hidden_states}
     @property
     def        lowerCAmelCase__  (				self	):
            return (4, 4, 32, 32)
     @property
     def        lowerCAmelCase__  (				self	):
            return (4, 4, 32, 32)
     def        lowerCAmelCase__  (				self	):
            a			   ={
                """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,
            }
            a			   =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	):
            a   ,      a			   =self.prepare_init_args_and_inputs_for_common()
            a			   =self.model_class(**_lowerCAmelCase	)

            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	):
            a   ,      a			   =self.prepare_init_args_and_inputs_for_common()

            a			   =32

            a			   =self.model_class(**_lowerCAmelCase	)
            model.to(_lowerCAmelCase	)
            model.eval()

            with torch.no_grad():
                   a			   =model(**_lowerCAmelCase	)

                   if isinstance(_lowerCAmelCase ,    _lowerCAmelCase	):
                          a			   =output.sample

            self.assertIsNotNone(_lowerCAmelCase	)
            a			   =inputs_dict["""sample"""].shape
            self.assertEqual(output.shape ,    _lowerCAmelCase ,    """Input and output shapes do not match"""	)
     def        lowerCAmelCase__  (				self	):
            a   ,      a			   =self.prepare_init_args_and_inputs_for_common()
            a			   =self.model_class(**_lowerCAmelCase	)
            model.to(_lowerCAmelCase	)
            model.eval()

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

                   a			   =model(**_lowerCAmelCase	)
                   if isinstance(_lowerCAmelCase ,    _lowerCAmelCase	):
                          a			   =first.sample

                   a			   =model(**_lowerCAmelCase	)
                   if isinstance(_lowerCAmelCase ,    _lowerCAmelCase	):
                          a			   =second.sample

            a			   =first.cpu().numpy()
            a			   =second.cpu().numpy()
            a			   =out_a[~np.isnan(_lowerCAmelCase	)]
            a			   =out_a[~np.isnan(_lowerCAmelCase	)]
            a			   =np.amax(np.abs(out_a - out_a	)	)
            self.assertLessEqual(_lowerCAmelCase ,    1E-5	)
     def        lowerCAmelCase__  (				self	):
            a   ,      a			   =self.prepare_init_args_and_inputs_for_common()

            a			   =8

            a			   =self.model_class(**_lowerCAmelCase	)
            model.to(_lowerCAmelCase	)
            model.eval()

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

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

            model.set_attention_slice(2	)
            with torch.no_grad():
                   a			   =model(**_lowerCAmelCase	)
            assert output is not None
     def        lowerCAmelCase__  (				self	):
            a   ,      a			   =self.prepare_init_args_and_inputs_for_common()

            a			   =8

            a			   =self.model_class(**_lowerCAmelCase	)
            model.to(_lowerCAmelCase	)

            with torch.no_grad():
                   a			   =model(**_lowerCAmelCase	).sample

            a			   =create_lora_layers(_lowerCAmelCase	)

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

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

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

            a			   =8

            torch.manual_seed(0	)
            a			   =self.model_class(**_lowerCAmelCase	)
            model.to(_lowerCAmelCase	)

            with torch.no_grad():
                   a			   =model(**_lowerCAmelCase	).sample

            a			   =create_lora_layers(_lowerCAmelCase	)
            model.set_attn_processor(_lowerCAmelCase	)

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

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

            with torch.no_grad():
                   a			   =new_model(**_lowerCAmelCase ,    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	):
            a   ,      a			   =self.prepare_init_args_and_inputs_for_common()

            a			   =8

            torch.manual_seed(0	)
            a			   =self.model_class(**_lowerCAmelCase	)
            model.to(_lowerCAmelCase	)

            with torch.no_grad():
                   a			   =model(**_lowerCAmelCase	).sample

            a			   =create_lora_layers(_lowerCAmelCase	)
            model.set_attn_processor(_lowerCAmelCase	)

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

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

            with torch.no_grad():
                   a			   =new_model(**_lowerCAmelCase ,    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	):
            # enable deterministic behavior for gradient checkpointing
            a   ,      a			   =self.prepare_init_args_and_inputs_for_common()

            a			   =8

            torch.manual_seed(0	)
            a			   =self.model_class(**_lowerCAmelCase	)
            model.to(_lowerCAmelCase	)

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

            a			   =8

            torch.manual_seed(0	)
            a			   =self.model_class(**_lowerCAmelCase	)
            model.to(_lowerCAmelCase	)

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

            a			   =8

            torch.manual_seed(0	)
            a			   =self.model_class(**_lowerCAmelCase	)
            model.to(_lowerCAmelCase	)

            with torch.no_grad():
                   a			   =model(**_lowerCAmelCase	).sample

            a			   =create_lora_layers(_lowerCAmelCase	)
            model.set_attn_processor(_lowerCAmelCase	)

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

            model.set_attn_processor(AttnProcessor()	)

            with torch.no_grad():
                   a			   =model(**_lowerCAmelCase	).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	):
            # enable deterministic behavior for gradient checkpointing
            a   ,      a			   =self.prepare_init_args_and_inputs_for_common()

            a			   =4

            torch.manual_seed(0	)
            a			   =self.model_class(**_lowerCAmelCase	)
            model.to(_lowerCAmelCase	)
            a			   =create_lora_layers(_lowerCAmelCase	)
            model.set_attn_processor(_lowerCAmelCase	)

            # default
            with torch.no_grad():
                   a			   =model(**_lowerCAmelCase	).sample

                   model.enable_xformers_memory_efficient_attention()
                   a			   =model(**_lowerCAmelCase	).sample

                   model.disable_xformers_memory_efficient_attention()
                   a			   =model(**_lowerCAmelCase	).sample

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






     def        lowerCAmelCase__  (				self	):
            a   ,      a			   =self.prepare_init_args_and_inputs_for_common()
            a			   =32

            a			   =self.model_class(**_lowerCAmelCase	)
            model.to(_lowerCAmelCase	)
            model.eval()

            with torch.no_grad():
                   a			   =model(**_lowerCAmelCase	)[0]

            model.enable_forward_chunking()
            with torch.no_grad():
                   a			   =model(**_lowerCAmelCase	)[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.






import functools
import operator

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


_lowerCamelCase							=      logging.get_logger(__name__)

_lowerCamelCase							=      {
    '''facebook/hubert-base-ls960''': '''https://huggingface.co/facebook/hubert-base-ls960/resolve/main/config.json''',
    # See all Hubert models at https://huggingface.co/models?filter=hubert
}


class        UpperCAmelCase__				(							UpperCAmelCase__				):



     '''simple docstring'''
     _SCREAMING_SNAKE_CASE      :					Dict     					=    "hubert"
     def __init__(				self ,    _lowerCAmelCase=32 ,    _lowerCAmelCase=768 ,    _lowerCAmelCase=12 ,    _lowerCAmelCase=12 ,    _lowerCAmelCase=3_072 ,    _lowerCAmelCase="gelu" ,    _lowerCAmelCase=0.1 ,    _lowerCAmelCase=0.1 ,    _lowerCAmelCase=0.1 ,    _lowerCAmelCase=True ,    _lowerCAmelCase=0.0 ,    _lowerCAmelCase=0.1 ,    _lowerCAmelCase=0.1 ,    _lowerCAmelCase=0.02 ,    _lowerCAmelCase=1E-5 ,    _lowerCAmelCase="group" ,    _lowerCAmelCase="gelu" ,    _lowerCAmelCase=(512, 512, 512, 512, 512, 512, 512) ,    _lowerCAmelCase=(5, 2, 2, 2, 2, 2, 2) ,    _lowerCAmelCase=(10, 3, 3, 3, 3, 2, 2) ,    _lowerCAmelCase=False ,    _lowerCAmelCase=128 ,    _lowerCAmelCase=16 ,    _lowerCAmelCase=False ,    _lowerCAmelCase=True ,    _lowerCAmelCase=0.05 ,    _lowerCAmelCase=10 ,    _lowerCAmelCase=2 ,    _lowerCAmelCase=0.0 ,    _lowerCAmelCase=10 ,    _lowerCAmelCase=0 ,    _lowerCAmelCase="sum" ,    _lowerCAmelCase=False ,    _lowerCAmelCase=False ,    _lowerCAmelCase=256 ,    _lowerCAmelCase=0 ,    _lowerCAmelCase=1 ,    _lowerCAmelCase=2 ,    **_lowerCAmelCase ,    ):
            super().__init__(**_lowerCAmelCase ,    pad_token_id=_lowerCAmelCase ,    bos_token_id=_lowerCAmelCase ,    eos_token_id=_lowerCAmelCase	)
            a			   =hidden_size
            a			   =feat_extract_norm
            a			   =feat_extract_activation
            a			   =list(_lowerCAmelCase	)
            a			   =list(_lowerCAmelCase	)
            a			   =list(_lowerCAmelCase	)
            a			   =conv_bias
            a			   =num_conv_pos_embeddings
            a			   =num_conv_pos_embedding_groups
            a			   =len(self.conv_dim	)
            a			   =num_hidden_layers
            a			   =intermediate_size
            a			   =hidden_act
            a			   =num_attention_heads
            a			   =hidden_dropout
            a			   =attention_dropout
            a			   =activation_dropout
            a			   =feat_proj_layer_norm
            a			   =feat_proj_dropout
            a			   =final_dropout
            a			   =layerdrop
            a			   =layer_norm_eps
            a			   =initializer_range
            a			   =vocab_size
            a			   =do_stable_layer_norm
            a			   =use_weighted_layer_sum
            a			   =classifier_proj_size

            if (
                (len(self.conv_stride	) != self.num_feat_extract_layers)
                or (len(self.conv_kernel	) != self.num_feat_extract_layers)
                or (len(self.conv_dim	) != self.num_feat_extract_layers)
            ):
                   raise ValueError(
                       """Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` =="""
                       """ `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) ="""
                       F''' {len(self.conv_dim	)}`, `len(config.conv_stride) = {len(self.conv_stride	)}`,'''
                       F''' `len(config.conv_kernel) = {len(self.conv_kernel	)}`.'''	)

            # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
            a			   =apply_spec_augment
            a			   =mask_time_prob
            a			   =mask_time_length
            a			   =mask_time_min_masks
            a			   =mask_feature_prob
            a			   =mask_feature_length
            a			   =mask_feature_min_masks

            # ctc loss
            a			   =ctc_loss_reduction
            a			   =ctc_zero_infinity






     @property
     def        lowerCAmelCase__  (				self	):
            return functools.reduce(operator.mul ,    self.conv_stride ,    1	)
