code: 
  
  
  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()
  
  SCREAMING_SNAKE_CASE : Any          =    logging.get_logger(__name__)
  
  
  
  
  
  def 						_lowerCamelCase		(							SCREAMING_SNAKE_CASE_				: int	,	SCREAMING_SNAKE_CASE_				: bool = True      ):
  
  
  
  
  
  
         """simple docstring"""
  
  
  
         a_       :   Dict			=	{}
         for name in model.attn_processors.keys():
                a_       :   Optional[Any]			=	name.endswith("""attn2.processor"""      ) and not (
                    name.startswith("""transformer_in"""      ) or """temp_attentions""" in name.split("""."""      )
                )
                a_       :   str			=	model.config.cross_attention_dim if has_cross_attention else None
                if name.startswith("""mid_block"""      ):
                       a_       :   Union[str, Any]			=	model.config.block_out_channels[-1]
                elif name.startswith("""up_blocks"""      ):
                       a_       :   List[Any]			=	int(name[len("""up_blocks."""      )]      )
                       a_       :   str			=	list(reversed(model.config.block_out_channels      )      )[block_id]
                elif name.startswith("""down_blocks"""      ):
                       a_       :   List[Any]			=	int(name[len("""down_blocks."""      )]      )
                       a_       :   Any			=	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_       :   Union[str, Any]			=	8 * model.config.attention_head_dim
  
                a_       :   List[str]			=	LoRAAttnProcessor(hidden_size=SCREAMING_SNAKE_CASE_	,	cross_attention_dim=SCREAMING_SNAKE_CASE_      )
                a_       :   Tuple			=	lora_attn_procs[name].to(model.device      )
  
                if mock_weights:
                       # add 1 to weights to mock trained weights
                       with torch.no_grad():
                              lora_attn_procs[name].to_q_lora.up.weight += 1
                              lora_attn_procs[name].to_k_lora.up.weight += 1
                              lora_attn_procs[name].to_v_lora.up.weight += 1
                              lora_attn_procs[name].to_out_lora.up.weight += 1
  
         return lora_attn_procs
  
  
  
  
  
  @skip_mps
  class 					snake_case__  ( __A      ,    __A      ,    unittest.TestCase       ):
   UpperCAmelCase   :							Dict         = UNetaDConditionModel
   UpperCAmelCase   :							Optional[Any]         = """sample"""
  
  
  
  
  
   @property
   def     A				(  self							)		->			List[Any]:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   Any			=	4
          a_       :   List[str]			=	4
          a_       :   int			=	4
          a_       :   List[str]			=	(32, 32)
  
          a_       :   List[Any]			=	floats_tensor((batch_size, num_channels, num_frames) + sizes							).to(UpperCamelCase_							)
          a_       :   Optional[Any]			=	torch.tensor([10]							).to(UpperCamelCase_							)
          a_       :   Tuple			=	floats_tensor((batch_size, 4, 32)							).to(UpperCamelCase_							)
  
          return {"sample": noise, "timestep": time_step, "encoder_hidden_states": encoder_hidden_states}
  
  
  
  
  
   @property
   def     A				(  self							)		->			List[Any]:
          """simple docstring"""
  
  
  
  
  
  
          return (4, 4, 32, 32)
  
  
  
  
  
   @property
   def     A				(  self							)		->			str:
          """simple docstring"""
  
  
  
  
  
  
          return (4, 4, 32, 32)
  
  
  
  
  
   def     A				(  self							)		->			List[str]:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   Any			=	{
              """block_out_channels""": (32, 64),
              """down_block_types""": (
                  """CrossAttnDownBlock3D""",
                  """DownBlock3D""",
              ),
              """up_block_types""": ("""UpBlock3D""", """CrossAttnUpBlock3D"""),
              """cross_attention_dim""": 32,
              """attention_head_dim""": 8,
              """out_channels""": 4,
              """in_channels""": 4,
              """layers_per_block""": 1,
              """sample_size""": 32,
          }
          a_       :   Union[str, Any]			=	self.dummy_input
          return init_dict, inputs_dict
  
  
  
  
  
   @unittest.skipIf(
       torch_device != """cuda""" or not is_xformers_available()				,    reason="""XFormers attention is only available with CUDA and `xformers` installed"""				,    )
   def     A				(  self							)		->			List[str]:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   Dict			=	self.prepare_init_args_and_inputs_for_common()
          a_       :   List[str]			=	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     A				(  self							)		->			List[Any]:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   Union[str, Any]			=	self.prepare_init_args_and_inputs_for_common()
  
          a_       :   int			=	32
  
          a_       :   Any			=	self.model_class(**UpperCamelCase_							)
          model.to(UpperCamelCase_							)
          model.eval()
  
          with torch.no_grad():
                 a_       :   int			=	model(**UpperCamelCase_							)
  
                 if isinstance(UpperCamelCase_				,    UpperCamelCase_							):
                        a_       :   Dict			=	output.sample
  
          self.assertIsNotNone(UpperCamelCase_							)
          a_       :   Optional[Any]			=	inputs_dict["""sample"""].shape
          self.assertEqual(output.shape				,    UpperCamelCase_				,    """Input and output shapes do not match"""							)
  
  
  
  
  
   def     A				(  self							)		->			Any:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   Any			=	self.prepare_init_args_and_inputs_for_common()
          a_       :   str			=	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							)
  
                 a_       :   Tuple			=	model(**UpperCamelCase_							)
                 if isinstance(UpperCamelCase_				,    UpperCamelCase_							):
                        a_       :   int			=	first.sample
  
                 a_       :   Optional[int]			=	model(**UpperCamelCase_							)
                 if isinstance(UpperCamelCase_				,    UpperCamelCase_							):
                        a_       :   Optional[Any]			=	second.sample
  
          a_       :   Union[str, Any]			=	first.cpu().numpy()
          a_       :   List[Any]			=	second.cpu().numpy()
          a_       :   Dict			=	out_a[~np.isnan(UpperCamelCase_							)]
          a_       :   List[str]			=	out_a[~np.isnan(UpperCamelCase_							)]
          a_       :   Dict			=	np.amax(np.abs(out_a - out_a							)							)
          self.assertLessEqual(UpperCamelCase_				,    1e-5							)
  
  
  
  
  
   def     A				(  self							)		->			Tuple:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   Dict			=	self.prepare_init_args_and_inputs_for_common()
  
          a_       :   List[Any]			=	8
  
          a_       :   str			=	self.model_class(**UpperCamelCase_							)
          model.to(UpperCamelCase_							)
          model.eval()
  
          model.set_attention_slice("""auto"""							)
          with torch.no_grad():
                 a_       :   str			=	model(**UpperCamelCase_							)
          assert output is not None
  
          model.set_attention_slice("""max"""							)
          with torch.no_grad():
                 a_       :   Any			=	model(**UpperCamelCase_							)
          assert output is not None
  
          model.set_attention_slice(2							)
          with torch.no_grad():
                 a_       :   int			=	model(**UpperCamelCase_							)
          assert output is not None
  
  
  
  
  
   def     A				(  self							)		->			Optional[Any]:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   Union[str, Any]			=	self.prepare_init_args_and_inputs_for_common()
  
          a_       :   Optional[int]			=	8
  
          a_       :   str			=	self.model_class(**UpperCamelCase_							)
          model.to(UpperCamelCase_							)
  
          with torch.no_grad():
                 a_       :   List[str]			=	model(**UpperCamelCase_							).sample
  
          a_       :   List[Any]			=	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():
                 a_       :   int			=	model(**UpperCamelCase_				,    cross_attention_kwargs={"""scale""": 0.0}							).sample
                 a_       :   List[str]			=	model(**UpperCamelCase_				,    cross_attention_kwargs={"""scale""": 0.5}							).sample
                 a_       :   Any			=	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     A				(  self							)		->			Optional[Any]:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   int			=	self.prepare_init_args_and_inputs_for_common()
  
          a_       :   int			=	8
  
          torch.manual_seed(0							)
          a_       :   Tuple			=	self.model_class(**UpperCamelCase_							)
          model.to(UpperCamelCase_							)
  
          with torch.no_grad():
                 a_       :   Any			=	model(**UpperCamelCase_							).sample
  
          a_       :   List[str]			=	create_lora_layers(UpperCamelCase_							)
          model.set_attn_processor(UpperCamelCase_							)
  
          with torch.no_grad():
                 a_       :   str			=	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							)
                 a_       :   Dict			=	self.model_class(**UpperCamelCase_							)
                 new_model.to(UpperCamelCase_							)
                 new_model.load_attn_procs(UpperCamelCase_							)
  
          with torch.no_grad():
                 a_       :   Optional[int]			=	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     A				(  self							)		->			str:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   Tuple			=	self.prepare_init_args_and_inputs_for_common()
  
          a_       :   Optional[int]			=	8
  
          torch.manual_seed(0							)
          a_       :   Dict			=	self.model_class(**UpperCamelCase_							)
          model.to(UpperCamelCase_							)
  
          with torch.no_grad():
                 a_       :   Optional[Any]			=	model(**UpperCamelCase_							).sample
  
          a_       :   Tuple			=	create_lora_layers(UpperCamelCase_							)
          model.set_attn_processor(UpperCamelCase_							)
  
          with torch.no_grad():
                 a_       :   Optional[int]			=	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							)
                 a_       :   Dict			=	self.model_class(**UpperCamelCase_							)
                 new_model.to(UpperCamelCase_							)
                 new_model.load_attn_procs(UpperCamelCase_							)
  
          with torch.no_grad():
                 a_       :   str			=	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     A				(  self							)		->			Optional[int]:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   Dict			=	self.prepare_init_args_and_inputs_for_common()
  
          a_       :   List[str]			=	8
  
          torch.manual_seed(0							)
          a_       :   Union[str, Any]			=	self.model_class(**UpperCamelCase_							)
          model.to(UpperCamelCase_							)
  
          a_       :   Dict			=	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							)
                 a_       :   Dict			=	self.model_class(**UpperCamelCase_							)
                 new_model.to(UpperCamelCase_							)
                 new_model.load_attn_procs(UpperCamelCase_				,    weight_name="""pytorch_lora_weights.bin"""							)
  
  
  
  
  
   def     A				(  self							)		->			str:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   Union[str, Any]			=	self.prepare_init_args_and_inputs_for_common()
  
          a_       :   Dict			=	8
  
          torch.manual_seed(0							)
          a_       :   Dict			=	self.model_class(**UpperCamelCase_							)
          model.to(UpperCamelCase_							)
  
          a_       :   List[Any]			=	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							)
                 a_       :   Tuple			=	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     A				(  self							)		->			Tuple:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   Tuple			=	self.prepare_init_args_and_inputs_for_common()
  
          a_       :   Optional[int]			=	8
  
          torch.manual_seed(0							)
          a_       :   Any			=	self.model_class(**UpperCamelCase_							)
          model.to(UpperCamelCase_							)
  
          with torch.no_grad():
                 a_       :   Dict			=	model(**UpperCamelCase_							).sample
  
          a_       :   Dict			=	create_lora_layers(UpperCamelCase_							)
          model.set_attn_processor(UpperCamelCase_							)
  
          with torch.no_grad():
                 a_       :   int			=	model(**UpperCamelCase_				,    cross_attention_kwargs={"""scale""": 0.0}							).sample
  
          model.set_attn_processor(AttnProcessor()							)
  
          with torch.no_grad():
                 a_       :   Any			=	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     A				(  self							)		->			Tuple:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   str			=	self.prepare_init_args_and_inputs_for_common()
  
          a_       :   List[Any]			=	4
  
          torch.manual_seed(0							)
          a_       :   Optional[Any]			=	self.model_class(**UpperCamelCase_							)
          model.to(UpperCamelCase_							)
          a_       :   Tuple			=	create_lora_layers(UpperCamelCase_							)
          model.set_attn_processor(UpperCamelCase_							)
  
          # default
          with torch.no_grad():
                 a_       :   Tuple			=	model(**UpperCamelCase_							).sample
  
                 model.enable_xformers_memory_efficient_attention()
                 a_       :   Tuple			=	model(**UpperCamelCase_							).sample
  
                 model.disable_xformers_memory_efficient_attention()
                 a_       :   Dict			=	model(**UpperCamelCase_							).sample
  
          assert (sample - on_sample).abs().max() < 1e-4
          assert (sample - off_sample).abs().max() < 1e-4
  
  
  
  
  
   def     A				(  self							)		->			List[str]:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   int			=	self.prepare_init_args_and_inputs_for_common()
          a_       :   List[Any]			=	32
  
          a_       :   str			=	self.model_class(**UpperCamelCase_							)
          model.to(UpperCamelCase_							)
          model.eval()
  
          with torch.no_grad():
                 a_       :   List[Any]			=	model(**UpperCamelCase_							)[0]
  
          model.enable_forward_chunking()
          with torch.no_grad():
                 a_       :   Tuple			=	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.
  
  
  
style_context: 
  
  
  from __future__ import annotations
  
  
  
  
  
  def 						_lowerCamelCase		(							SCREAMING_SNAKE_CASE_				: list      ):
  
  
  
  
  
  
  							"""simple docstring"""
  
  
  
  							a_       :   Tuple			=	sorted(SCREAMING_SNAKE_CASE_      )
  							a_       :   Tuple			=	len(SCREAMING_SNAKE_CASE_      )
  							a_       :   Tuple			=	length >> 1
  							return (
  							    (sorted_list[mid_index] + sorted_list[mid_index - 1]) / 2
  							    if length % 2 == 0
  							    else sorted_list[mid_index]
  							)
  
  
  
  
  
  def 						_lowerCamelCase		(							):
  
  
  
  
  
  
  							"""simple docstring"""
  
  
  
  							import doctest
  
  							doctest.testmod()
  
  
  if __name__ == "__main__":
  	main()
  
  
  