





'''simple docstring'''






# Copyright (c) 2021-, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

####################################################################################################

#
# Note: If when running this conversion script you're getting an exception:
#     ModuleNotFoundError: No module named 'megatron.model.enums'
# you need to tell python where to find the clone of Megatron-LM, e.g.:
#
# cd /tmp
# git clone https://github.com/NVIDIA/Megatron-LM
# PYTHONPATH=/tmp/Megatron-LM python src/transformers/models/megatron_bert/convert_megatron_bert_checkpoint.py ...
#
# if you already have it cloned elsewhere, simply adjust the path to the existing path
#
# If the training was done using a Megatron-LM fork, e.g.,
# https://github.com/microsoft/Megatron-DeepSpeed/ then chances are that you need to have that one
# in your path, i.e., /path/to/Megatron-DeepSpeed/
#

import argparse
import os
import re
import zipfile

import torch

from transformers import MegatronBertConfig
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase=0						):





      """simple docstring"""


      if name is None:
            _lowerCAmelCase							  =							None
      else:
            _lowerCAmelCase							  =							""".""" * max(0							,			spaces - 2						) + """# {:""" + str(50 - spaces						) + """s}"""
            _lowerCAmelCase							  =							fmt.format(lowerCAmelCase						)

      # Print and recurse (if needed).
      if isinstance(lowerCAmelCase							,			lowerCAmelCase						):
            if msg is not None:
                  print(lowerCAmelCase						)
            for k in val.keys():
                  recursive_print(lowerCAmelCase							,			val[k]							,			spaces + 2						)
      elif isinstance(lowerCAmelCase							,			torch.Tensor						):
            print(lowerCAmelCase							,			""":"""							,			val.size()						)
      else:
            print(lowerCAmelCase							,			""":"""							,			lowerCAmelCase						)
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase						):





      """simple docstring"""


      _lowerCAmelCase							  =							param.size()
      if checkpoint_version == 1.0:
            # version 1.0 stores [num_heads * hidden_size * num_splits, :]
            _lowerCAmelCase							  =							(num_heads, hidden_size, num_splits) + input_shape[1:]
            _lowerCAmelCase							  =							param.view(*lowerCAmelCase						)
            _lowerCAmelCase							  =							param.transpose(0							,			2						)
            _lowerCAmelCase							  =							param.transpose(1							,			2						).contiguous()
      elif checkpoint_version >= 2.0:
            # other versions store [num_heads * num_splits * hidden_size, :]
            _lowerCAmelCase							  =							(num_heads, num_splits, hidden_size) + input_shape[1:]
            _lowerCAmelCase							  =							param.view(*lowerCAmelCase						)
            _lowerCAmelCase							  =							param.transpose(0							,			1						).contiguous()
      _lowerCAmelCase							  =							param.view(*lowerCAmelCase						)
      return param
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase						):





      """simple docstring"""


      _lowerCAmelCase							  =							{}

      # old versions did not store training args
      _lowerCAmelCase							  =							input_state_dict.get("""args"""							,			lowerCAmelCase						)
      if ds_args is not None:
            # do not make the user write a config file when the exact dimensions/sizes are already in the checkpoint
            # from pprint import pprint
            # pprint(vars(ds_args))

            _lowerCAmelCase							  =							ds_args.tokenizer_type
            _lowerCAmelCase							  =							ds_args.padded_vocab_size
            _lowerCAmelCase							  =							ds_args.max_position_embeddings
            _lowerCAmelCase							  =							ds_args.hidden_size
            _lowerCAmelCase							  =							ds_args.num_layers
            _lowerCAmelCase							  =							ds_args.num_attention_heads
            _lowerCAmelCase							  =							ds_args.ffn_hidden_size if """ffn_hidden_size""" in ds_args else 4 * ds_args.hidden_size
            # pprint(config)

      # The number of heads.
      _lowerCAmelCase							  =							config.num_attention_heads
      # The hidden_size per head.
      _lowerCAmelCase							  =							config.hidden_size // heads
      # Megatron-LM checkpoint version
      if "checkpoint_version" in input_state_dict.keys():
            _lowerCAmelCase							  =							input_state_dict["""checkpoint_version"""]
      else:
            _lowerCAmelCase							  =							0.0

      # The model.
      _lowerCAmelCase							  =							input_state_dict["""model"""]
      # The language model.
      _lowerCAmelCase							  =							model["""language_model"""]
      # The embeddings.
      _lowerCAmelCase							  =							lm["""embedding"""]

      # The word embeddings.
      _lowerCAmelCase							  =							embeddings["""word_embeddings"""]["""weight"""]
      # Truncate the embedding table to vocab_size rows.
      _lowerCAmelCase							  =							word_embeddings[: config.vocab_size, :]
      # Store the word embeddings.
      _lowerCAmelCase							  =							word_embeddings

      # The position embeddings.
      _lowerCAmelCase							  =							embeddings["""position_embeddings"""]["""weight"""]
      assert pos_embeddings.size(0						) == config.max_position_embeddings and pos_embeddings.size(1						) == config.hidden_size
      # Store the position embeddings.
      _lowerCAmelCase							  =							pos_embeddings

      # The token-type embeddings.
      _lowerCAmelCase							  =							embeddings["""tokentype_embeddings"""]["""weight"""]
      # Store the position embeddings.
      _lowerCAmelCase							  =							tokentype_embeddings

      # The transformer.
      _lowerCAmelCase							  =							lm["""transformer"""] if """transformer""" in lm.keys() else lm["""encoder"""]

      # The regex to extract layer names.
      _lowerCAmelCase							  =							re.compile(r"""layers\.(\d+)\.([a-z0-9_.]+)\.([a-z]+)"""						)

      # The simple map of names for "automated" rules.
      _lowerCAmelCase							  =							{
          """attention.dense""": """.attention.output.dense.""",
          """self_attention.dense""": """.attention.output.dense.""",
          """mlp.dense_h_to_4h""": """.intermediate.dense.""",
          """mlp.dense_4h_to_h""": """.output.dense.""",
      }

      # Keep track of the attention/query/value tensor.
      _lowerCAmelCase							  =							None

      # Extract the layers.
      for key, val in transformer.items():
            # Match the name.
            _lowerCAmelCase							  =							layer_re.match(lowerCAmelCase						)

            # Stop if that's not a layer
            if m is None:
                  break

            # The index of the layer.
            _lowerCAmelCase							  =							int(m.group(1						)						)
            # The name of the operation.
            _lowerCAmelCase							  =							m.group(2						)
            # Is it a weight or a bias?
            _lowerCAmelCase							  =							m.group(3						)

            # The name of the layer.
            _lowerCAmelCase							  =							f"bert.encoder.layer.{layer_idx}"

            # For layernorm(s), simply store the layer norm.
            if op_name.endswith("""layernorm"""						):
                  _lowerCAmelCase							  =							"""attention.ln""" if op_name.startswith("""input"""						) else """ln"""
                  _lowerCAmelCase							  =							val

            # Transpose the QKV matrix.
            elif (
                op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
            ) and weight_or_bias == "weight":
                  # Make sure the QKV pointer is nil.
                  assert attention_qkv_weight is None, ""

                  _lowerCAmelCase							  =							fix_query_key_value_ordering(lowerCAmelCase							,			lowerCAmelCase							,			3							,			lowerCAmelCase							,			lowerCAmelCase						)
                  # Store the tensor as we need the bias as well to interleave QKV and biases.
                  _lowerCAmelCase							  =							out_val

            # Transpose the bias.
            elif (
                op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
            ) and weight_or_bias == "bias":
                  # Make sure we read the weight tensor.
                  assert attention_qkv_weight is not None, ""

                  # Split the QKV matrix into Q, K and V. Megatron stores Q,K,V interleaved.
                  _lowerCAmelCase							  =							attention_qkv_weight[0 * config.hidden_size : 1 * config.hidden_size, :]
                  _lowerCAmelCase							  =							attention_qkv_weight[1 * config.hidden_size : 2 * config.hidden_size, :]
                  _lowerCAmelCase							  =							attention_qkv_weight[2 * config.hidden_size : 3 * config.hidden_size, :]

                  _lowerCAmelCase							  =							fix_query_key_value_ordering(lowerCAmelCase							,			lowerCAmelCase							,			3							,			lowerCAmelCase							,			lowerCAmelCase						)
                  # Split the bias.
                  _lowerCAmelCase							  =							out_val[0 * config.hidden_size : 1 * config.hidden_size]
                  _lowerCAmelCase							  =							out_val[1 * config.hidden_size : 2 * config.hidden_size]
                  _lowerCAmelCase							  =							out_val[2 * config.hidden_size : 3 * config.hidden_size]

                  # Store.
                  _lowerCAmelCase							  =							q
                  _lowerCAmelCase							  =							q_bias
                  _lowerCAmelCase							  =							k
                  _lowerCAmelCase							  =							k_bias
                  _lowerCAmelCase							  =							v
                  _lowerCAmelCase							  =							v_bias

                  # Clear the stored tensor.
                  _lowerCAmelCase							  =							None

            # Copy weights and biases as is.
            elif weight_or_bias in ["weight", "bias"]:
                  _lowerCAmelCase							  =							megatron_to_transformers[op_name]
                  _lowerCAmelCase							  =							val

    # The final layernorm.
      _lowerCAmelCase							  =							transformer["""final_layernorm.weight"""]
      _lowerCAmelCase							  =							transformer["""final_layernorm.bias"""]

      # The pooler.
      _lowerCAmelCase							  =							lm["""pooler"""]

      # Store the matrix and the bias.
      _lowerCAmelCase							  =							pooler["""dense.weight"""]
      _lowerCAmelCase							  =							pooler["""dense.bias"""]

      # The LM head from Megatron (for RACE).
      _lowerCAmelCase							  =							model["""lm_head"""]

      # The transform matrix.
      _lowerCAmelCase							  =							lm_head["""dense.weight"""]
      _lowerCAmelCase							  =							lm_head["""dense.bias"""]

      # The transform LN.
      _lowerCAmelCase							  =							lm_head["""layernorm.weight"""]
      _lowerCAmelCase							  =							lm_head["""layernorm.bias"""]

      # For the decoder, we replicate the weights.
      _lowerCAmelCase							  =							word_embeddings
      _lowerCAmelCase							  =							lm_head["""bias"""]

      # The classifier from Megatron (for MLNI).
      _lowerCAmelCase							  =							model["""binary_head"""]

      # Store the classifier.
      _lowerCAmelCase							  =							binary_head["""weight"""]
      _lowerCAmelCase							  =							binary_head["""bias"""]

      # It should be done!
      return output_state_dict
def      UpperCamelCase__   (	):





      """simple docstring"""


      _lowerCAmelCase							  =							argparse.ArgumentParser()
      parser.add_argument("""--print-checkpoint-structure"""							,			action="""store_true"""						)
      parser.add_argument("""path_to_checkpoint"""							,			type=lowerCAmelCase							,			help="""Path to the ZIP file containing the checkpoint"""						)
      parser.add_argument(
          """--config_file"""							,			default=""""""							,			type=lowerCAmelCase							,			help="""An optional config json file describing the pre-trained model."""							,			)
      _lowerCAmelCase							  =							parser.parse_args()

      # Extract the basename.
      _lowerCAmelCase							  =							os.path.dirname(args.path_to_checkpoint						)

      # Load the model.
      # the .zip is very optional, let's keep it for backward compatibility
      print(f"Extracting PyTorch state dictionary from \"{args.path_to_checkpoint}\""						)
      if args.path_to_checkpoint.endswith(""".zip"""						):
            with zipfile.ZipFile(args.path_to_checkpoint							,			"""r"""						) as checkpoint:
                  with checkpoint.open("""release/mp_rank_00/model_optim_rng.pt"""						) as pytorch_dict:
                        _lowerCAmelCase							  =							torch.load(lowerCAmelCase							,			map_location="""cpu"""						)
      else:
            _lowerCAmelCase							  =							torch.load(args.path_to_checkpoint							,			map_location="""cpu"""						)

      if args.config_file == "":
            # Default config of megatron-bert 345m
            _lowerCAmelCase							  =							MegatronBertConfig()

            # different megatron-bert-*-345m models have different vocab sizes, so override the default
            # config (which is for megatron-bert-cased-345m) with the actual vocab dimension
            _lowerCAmelCase							  =							input_state_dict["""model"""]["""lm_head"""]["""bias"""].numel()
      else:
            _lowerCAmelCase							  =							MegatronBertConfig.from_json_file(args.config_file						)

      # Convert.
      print("""Converting"""						)
      _lowerCAmelCase							  =							convert_megatron_checkpoint(lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase						)

      # Print the structure of converted state dict.
      if args.print_checkpoint_structure:
            recursive_print(lowerCAmelCase							,			lowerCAmelCase						)

      # Store the config to file.
      print("""Saving config"""						)
      config.save_pretrained(lowerCAmelCase						)

      # Store the state_dict to file.
      _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			"""pytorch_model.bin"""						)
      print(f"Saving checkpoint to \"{output_checkpoint_file}\""						)
      torch.save(lowerCAmelCase							,			lowerCAmelCase						)


####################################################################################################

if __name__ == "__main__":
      main()

####################################################################################################













'''simple docstring'''






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__       :     str          =logging.get_logger(__name__)
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase = True						):





      """simple docstring"""


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

            _lowerCAmelCase							  =							LoRAAttnProcessor(hidden_size=lowerCAmelCase							,			cross_attention_dim=lowerCAmelCase						)
            _lowerCAmelCase							  =							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  (    snake_case_							,	snake_case_							,	unittest.TestCase     ):
      _lowercase:    int		  =     UNetaDConditionModel
      _lowercase:    Optional[int]		  =     '''sample'''




      @property
      def        lowercase__					(    self    :				Any						)       ->   int:
            _lowerCAmelCase							  =							4
            _lowerCAmelCase							  =							4
            _lowerCAmelCase							  =							4
            _lowerCAmelCase							  =							(32, 32)

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

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




      @property
      def        lowercase__					(    self    :				Dict						)       ->   Any:
            return (4, 4, 32, 32)




      @property
      def        lowercase__					(    self    :				List[str]						)       ->   Optional[int]:
            return (4, 4, 32, 32)




      def        lowercase__					(    self    :				Optional[Any]						)       ->   Optional[Any]:
            _lowerCAmelCase							  =							{
                """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,
            }
            _lowerCAmelCase							  =							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        lowercase__					(    self    :				List[str]						)       ->   int:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()
            _lowerCAmelCase							  =							self.model_class(**__snake_case						)

            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        lowercase__					(    self    :				str						)       ->   int:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							32

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

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

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

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




      def        lowercase__					(    self    :				Any						)       ->   Any:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()
            _lowerCAmelCase							  =							self.model_class(**__snake_case						)
            model.to(__snake_case						)
            model.eval()

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

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

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

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




      def        lowercase__					(    self    :				Any						)       ->   List[str]:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							8

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

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

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

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




      def        lowercase__					(    self    :				str						)       ->   Any:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							8

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

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

            _lowerCAmelCase							  =							create_lora_layers(__snake_case						)

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

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

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

            _lowerCAmelCase							  =							8

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

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

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

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

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

            with torch.no_grad():
                  _lowerCAmelCase							  =							new_model(**__snake_case    ,    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        lowercase__					(    self    :				Optional[int]						)       ->   Optional[int]:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							8

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

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

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

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

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

            with torch.no_grad():
                  _lowerCAmelCase							  =							new_model(**__snake_case    ,    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        lowercase__					(    self    :				Dict						)       ->   List[str]:
            # enable deterministic behavior for gradient checkpointing
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							8

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

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




      def        lowercase__					(    self    :				Optional[int]						)       ->   List[Any]:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							8

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

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




      def        lowercase__					(    self    :				Any						)       ->   Any:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							8

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

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

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

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

            model.set_attn_processor(AttnProcessor()						)

            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case						).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        lowercase__					(    self    :				Union[str, Any]						)       ->   Tuple:
            # enable deterministic behavior for gradient checkpointing
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							4

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

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

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

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

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




      def        lowercase__					(    self    :				Dict						)       ->   Optional[int]:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()
            _lowerCAmelCase							  =							32

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

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

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






