





'''simple docstring'''






import argparse
import json

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import TimesformerConfig, TimesformerForVideoClassification, VideoMAEImageProcessor







def      UpperCamelCase__   (	lowerCAmelCase						):





      """simple docstring"""


      _lowerCAmelCase							  =							TimesformerConfig()

      if "large" in model_name:
            _lowerCAmelCase							  =							96

      if "hr" in model_name:
            _lowerCAmelCase							  =							16
            _lowerCAmelCase							  =							4_48

      _lowerCAmelCase							  =							"""huggingface/label-files"""
      if "k400" in model_name:
            _lowerCAmelCase							  =							4_00
            _lowerCAmelCase							  =							"""kinetics400-id2label.json"""
      elif "k600" in model_name:
            _lowerCAmelCase							  =							6_00
            _lowerCAmelCase							  =							"""kinetics600-id2label.json"""
      elif "ssv2" in model_name:
            _lowerCAmelCase							  =							1_74
            _lowerCAmelCase							  =							"""something-something-v2-id2label.json"""
      else:
            raise ValueError("""Model name should either contain 'k400', 'k600' or 'ssv2'."""						)
      _lowerCAmelCase							  =							json.load(open(hf_hub_download(lowerCAmelCase							,			lowerCAmelCase							,			repo_type="""dataset"""						)							,			"""r"""						)						)
      _lowerCAmelCase							  =							{int(lowerCAmelCase						): v for k, v in idalabel.items()}
      _lowerCAmelCase							  =							idalabel
      _lowerCAmelCase							  =							{v: k for k, v in idalabel.items()}

      return config







def      UpperCamelCase__   (	lowerCAmelCase						):





      """simple docstring"""


      if "encoder." in name:
            _lowerCAmelCase							  =							name.replace("""encoder."""							,			""""""						)
      if "cls_token" in name:
            _lowerCAmelCase							  =							name.replace("""cls_token"""							,			"""timesformer.embeddings.cls_token"""						)
      if "pos_embed" in name:
            _lowerCAmelCase							  =							name.replace("""pos_embed"""							,			"""timesformer.embeddings.position_embeddings"""						)
      if "time_embed" in name:
            _lowerCAmelCase							  =							name.replace("""time_embed"""							,			"""timesformer.embeddings.time_embeddings"""						)
      if "patch_embed.proj" in name:
            _lowerCAmelCase							  =							name.replace("""patch_embed.proj"""							,			"""timesformer.embeddings.patch_embeddings.projection"""						)
      if "patch_embed.norm" in name:
            _lowerCAmelCase							  =							name.replace("""patch_embed.norm"""							,			"""timesformer.embeddings.norm"""						)
      if "blocks" in name:
            _lowerCAmelCase							  =							name.replace("""blocks"""							,			"""timesformer.encoder.layer"""						)
      if "attn.proj" in name:
            _lowerCAmelCase							  =							name.replace("""attn.proj"""							,			"""attention.output.dense"""						)
      if "attn" in name and "bias" not in name and "temporal" not in name:
            _lowerCAmelCase							  =							name.replace("""attn"""							,			"""attention.self"""						)
      if "attn" in name and "temporal" not in name:
            _lowerCAmelCase							  =							name.replace("""attn"""							,			"""attention.attention"""						)
      if "temporal_norm1" in name:
            _lowerCAmelCase							  =							name.replace("""temporal_norm1"""							,			"""temporal_layernorm"""						)
      if "temporal_attn.proj" in name:
            _lowerCAmelCase							  =							name.replace("""temporal_attn"""							,			"""temporal_attention.output.dense"""						)
      if "temporal_fc" in name:
            _lowerCAmelCase							  =							name.replace("""temporal_fc"""							,			"""temporal_dense"""						)
      if "norm1" in name and "temporal" not in name:
            _lowerCAmelCase							  =							name.replace("""norm1"""							,			"""layernorm_before"""						)
      if "norm2" in name:
            _lowerCAmelCase							  =							name.replace("""norm2"""							,			"""layernorm_after"""						)
      if "mlp.fc1" in name:
            _lowerCAmelCase							  =							name.replace("""mlp.fc1"""							,			"""intermediate.dense"""						)
      if "mlp.fc2" in name:
            _lowerCAmelCase							  =							name.replace("""mlp.fc2"""							,			"""output.dense"""						)
      if "norm.weight" in name and "fc" not in name and "temporal" not in name:
            _lowerCAmelCase							  =							name.replace("""norm.weight"""							,			"""timesformer.layernorm.weight"""						)
      if "norm.bias" in name and "fc" not in name and "temporal" not in name:
            _lowerCAmelCase							  =							name.replace("""norm.bias"""							,			"""timesformer.layernorm.bias"""						)
      if "head" in name:
            _lowerCAmelCase							  =							name.replace("""head"""							,			"""classifier"""						)

      return name







def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase						):





      """simple docstring"""


      for key in orig_state_dict.copy().keys():
            _lowerCAmelCase							  =							orig_state_dict.pop(lowerCAmelCase						)

            if key.startswith("""model."""						):
                  _lowerCAmelCase							  =							key.replace("""model."""							,			""""""						)

            if "qkv" in key:
                  _lowerCAmelCase							  =							key.split("""."""						)
                  _lowerCAmelCase							  =							int(key_split[1]						)
                  _lowerCAmelCase							  =							"""timesformer.encoder.layer."""
                  if "temporal" in key:
                        _lowerCAmelCase							  =							""".temporal_attention.attention.qkv."""
                  else:
                        _lowerCAmelCase							  =							""".attention.attention.qkv."""
                  if "weight" in key:
                        _lowerCAmelCase							  =							val
                  else:
                        _lowerCAmelCase							  =							val
            else:
                  _lowerCAmelCase							  =							val

      return orig_state_dict







def      UpperCamelCase__   (	):





      """simple docstring"""


      _lowerCAmelCase							  =							hf_hub_download(
          repo_id="""hf-internal-testing/spaghetti-video"""							,			filename="""eating_spaghetti.npy"""							,			repo_type="""dataset"""						)
      _lowerCAmelCase							  =							np.load(lowerCAmelCase						)
      return list(lowerCAmelCase						)







def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase						):





      """simple docstring"""


      _lowerCAmelCase							  =							get_timesformer_config(lowerCAmelCase						)

      _lowerCAmelCase							  =							TimesformerForVideoClassification(lowerCAmelCase						)

      # download original checkpoint, hosted on Google Drive
      _lowerCAmelCase							  =							"""pytorch_model.bin"""
      gdown.cached_download(lowerCAmelCase							,			lowerCAmelCase							,			quiet=lowerCAmelCase						)
      _lowerCAmelCase							  =							torch.load(lowerCAmelCase							,			map_location="""cpu"""						)
      if "model" in files:
            _lowerCAmelCase							  =							files["""model"""]
      elif "module" in files:
            _lowerCAmelCase							  =							files["""module"""]
      else:
            _lowerCAmelCase							  =							files["""model_state"""]
      _lowerCAmelCase							  =							convert_state_dict(lowerCAmelCase							,			lowerCAmelCase						)

      model.load_state_dict(lowerCAmelCase						)
      model.eval()

      # verify model on basic input
      _lowerCAmelCase							  =							VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]							,			image_std=[0.5, 0.5, 0.5]						)
      _lowerCAmelCase							  =							prepare_video()
      _lowerCAmelCase							  =							image_processor(video[:8]							,			return_tensors="""pt"""						)

      _lowerCAmelCase							  =							model(**lowerCAmelCase						)
      _lowerCAmelCase							  =							outputs.logits

      _lowerCAmelCase							  =							[
          # Kinetics-400 checkpoints (hr = high resolution input of 448px instead of 224px)
          """timesformer-base-finetuned-k400""",
          """timesformer-large-finetuned-k400""",
          """timesformer-hr-finetuned-k400""",
          # Kinetics-600 checkpoints (hr = high resolution input of 448px instead of 224px)
          """timesformer-base-finetuned-k600""",
          """timesformer-large-finetuned-k600""",
          """timesformer-hr-finetuned-k600""",
          # Something-Something-v2 checkpoints (hr = high resolution input of 448px instead of 224px)
          """timesformer-base-finetuned-ssv2""",
          """timesformer-large-finetuned-ssv2""",
          """timesformer-hr-finetuned-ssv2""",
      ]

      # NOTE: logits were tested with image_mean and image_std equal to [0.5, 0.5, 0.5] and [0.5, 0.5, 0.5]
      if model_name == "timesformer-base-finetuned-k400":
            _lowerCAmelCase							  =							torch.Size([1, 4_00]						)
            _lowerCAmelCase							  =							torch.tensor([-0.3_016, -0.7_713, -0.4_205]						)
      elif model_name == "timesformer-base-finetuned-k600":
            _lowerCAmelCase							  =							torch.Size([1, 6_00]						)
            _lowerCAmelCase							  =							torch.tensor([-0.7_267, -0.7_466, 3.2_404]						)
      elif model_name == "timesformer-base-finetuned-ssv2":
            _lowerCAmelCase							  =							torch.Size([1, 1_74]						)
            _lowerCAmelCase							  =							torch.tensor([-0.9_059, 0.6_433, -3.1_457]						)
      elif model_name == "timesformer-large-finetuned-k400":
            _lowerCAmelCase							  =							torch.Size([1, 4_00]						)
            _lowerCAmelCase							  =							torch.tensor([0, 0, 0]						)
      elif model_name == "timesformer-large-finetuned-k600":
            _lowerCAmelCase							  =							torch.Size([1, 6_00]						)
            _lowerCAmelCase							  =							torch.tensor([0, 0, 0]						)
      elif model_name == "timesformer-large-finetuned-ssv2":
            _lowerCAmelCase							  =							torch.Size([1, 1_74]						)
            _lowerCAmelCase							  =							torch.tensor([0, 0, 0]						)
      elif model_name == "timesformer-hr-finetuned-k400":
            _lowerCAmelCase							  =							torch.Size([1, 4_00]						)
            _lowerCAmelCase							  =							torch.tensor([-0.9_617, -3.7_311, -3.7_708]						)
      elif model_name == "timesformer-hr-finetuned-k600":
            _lowerCAmelCase							  =							torch.Size([1, 6_00]						)
            _lowerCAmelCase							  =							torch.tensor([2.5_273, 0.7_127, 1.8_848]						)
      elif model_name == "timesformer-hr-finetuned-ssv2":
            _lowerCAmelCase							  =							torch.Size([1, 1_74]						)
            _lowerCAmelCase							  =							torch.tensor([-3.6_756, -0.7_513, 0.7_180]						)
      else:
            raise ValueError(f"Model name not supported. Should be one of {model_names}"						)

      # verify logits
      assert logits.shape == expected_shape
      assert torch.allclose(logits[0, :3]							,			lowerCAmelCase							,			atol=1e-4						)
      print("""Logits ok!"""						)

      if pytorch_dump_folder_path is not None:
            print(f"Saving model and image processor to {pytorch_dump_folder_path}"						)
            image_processor.save_pretrained(lowerCAmelCase						)
            model.save_pretrained(lowerCAmelCase						)

      if push_to_hub:
            print("""Pushing to the hub..."""						)
            model.push_to_hub(f"fcakyon/{model_name}"						)


if __name__ == "__main__":
      A__       :     Tuple          =argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          '''--checkpoint_url''',
          default='''https://drive.google.com/u/1/uc?id=17yvuYp9L4mn-HpIcK5Zo6K3UoOy1kA5l&export=download''',
          type=str,
          help=(
              '''URL of the original PyTorch checkpoint (on Google Drive) you\'d like to convert. Should be a direct'''
              ''' download link.'''
          ),
      )
      parser.add_argument(
          '''--pytorch_dump_folder_path''',
          default='''''',
          type=str,
          help='''Path to the output PyTorch model directory.''',
      )
      parser.add_argument('''--model_name''', default='''timesformer-base-finetuned-k400''', type=str, help='''Name of the model.''')
      parser.add_argument(
          '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
      )

      A__       :     List[Any]          =parser.parse_args()
      convert_timesformer_checkpoint(
          args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub
      )













'''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.






