






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

__UpperCAmelCase      =				logging.get_logger(__name__)





def 						A_					(						lowercase_	,  lowercase_ = True			)				->Optional[Any]:







 """simple docstring"""

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

  SCREAMING_SNAKE_CASE				    =		LoRAAttnProcessor(hidden_size=lowercase_	,  cross_attention_dim=lowercase_			)
  SCREAMING_SNAKE_CASE				    =		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 	a_(		lowercase__			,							lowercase__			,							unittest.TestCase	):



     """simple docstring"""



     __snake_case		:				List[str]  					=UNetaDConditionModel
     __snake_case		:				Dict  					='sample'





     @property
     def     __UpperCamelCase						(		self					:	List[Any])							->	List[Any]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		4
      SCREAMING_SNAKE_CASE				    =		4
      SCREAMING_SNAKE_CASE				    =		4
      SCREAMING_SNAKE_CASE				    =		(3_2, 3_2)

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

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





     @property
     def     __UpperCamelCase						(		self					:	int)							->	List[Any]:

      """simple docstring"""
      return (4, 4, 3_2, 3_2)





     @property
     def     __UpperCamelCase						(		self					:	List[str])							->	Any:

      """simple docstring"""
      return (4, 4, 3_2, 3_2)





     def     __UpperCamelCase						(		self					:	Tuple)							->	Any:

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

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.prepare_init_args_and_inputs_for_common()
      SCREAMING_SNAKE_CASE				    =		self.model_class(**_A)

      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     __UpperCamelCase						(		self					:	List[Any])							->	Any:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.prepare_init_args_and_inputs_for_common()

      SCREAMING_SNAKE_CASE				    =		3_2

      SCREAMING_SNAKE_CASE				    =		self.model_class(**_A)
      model.to(_A)
      model.eval()

      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		model(**_A)

       if isinstance(_A  ,  _A):
        SCREAMING_SNAKE_CASE				    =		output.sample

      self.assertIsNotNone(_A)
      SCREAMING_SNAKE_CASE				    =		inputs_dict['sample'].shape
      self.assertEqual(output.shape  ,  _A  ,  'Input and output shapes do not match')





     def     __UpperCamelCase						(		self					:	Optional[int])							->	Dict:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.prepare_init_args_and_inputs_for_common()
      SCREAMING_SNAKE_CASE				    =		self.model_class(**_A)
      model.to(_A)
      model.eval()

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

       SCREAMING_SNAKE_CASE				    =		model(**_A)
       if isinstance(_A  ,  _A):
        SCREAMING_SNAKE_CASE				    =		first.sample

       SCREAMING_SNAKE_CASE				    =		model(**_A)
       if isinstance(_A  ,  _A):
        SCREAMING_SNAKE_CASE				    =		second.sample

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





     def     __UpperCamelCase						(		self					:	Optional[Any])							->	Dict:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.prepare_init_args_and_inputs_for_common()

      SCREAMING_SNAKE_CASE				    =		8

      SCREAMING_SNAKE_CASE				    =		self.model_class(**_A)
      model.to(_A)
      model.eval()

      model.set_attention_slice('auto')
      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		model(**_A)
      assert output is not None

      model.set_attention_slice('max')
      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		model(**_A)
      assert output is not None

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





     def     __UpperCamelCase						(		self					:	Dict)							->	Optional[Any]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.prepare_init_args_and_inputs_for_common()

      SCREAMING_SNAKE_CASE				    =		8

      SCREAMING_SNAKE_CASE				    =		self.model_class(**_A)
      model.to(_A)

      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		model(**_A).sample

      SCREAMING_SNAKE_CASE				    =		create_lora_layers(_A)

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

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

      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		model(**_A  ,  cross_attention_kwargs={'scale': 0.0}).sample
       SCREAMING_SNAKE_CASE				    =		model(**_A  ,  cross_attention_kwargs={'scale': 0.5}).sample
       SCREAMING_SNAKE_CASE				    =		model(**_A  ,  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     __UpperCamelCase						(		self					:	List[str])							->	Any:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.prepare_init_args_and_inputs_for_common()

      SCREAMING_SNAKE_CASE				    =		8

      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE				    =		self.model_class(**_A)
      model.to(_A)

      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		model(**_A).sample

      SCREAMING_SNAKE_CASE				    =		create_lora_layers(_A)
      model.set_attn_processor(_A)

      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		model(**_A  ,  cross_attention_kwargs={'scale': 0.5}).sample

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

      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		new_model(**_A  ,  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     __UpperCamelCase						(		self					:	Tuple)							->	str:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.prepare_init_args_and_inputs_for_common()

      SCREAMING_SNAKE_CASE				    =		8

      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE				    =		self.model_class(**_A)
      model.to(_A)

      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		model(**_A).sample

      SCREAMING_SNAKE_CASE				    =		create_lora_layers(_A)
      model.set_attn_processor(_A)

      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		model(**_A  ,  cross_attention_kwargs={'scale': 0.5}).sample

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

      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		new_model(**_A  ,  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     __UpperCamelCase						(		self					:	Any)							->	Tuple:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.prepare_init_args_and_inputs_for_common()

      SCREAMING_SNAKE_CASE				    =		8

      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE				    =		self.model_class(**_A)
      model.to(_A)

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





     def     __UpperCamelCase						(		self					:	Any)							->	Union[str, Any]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.prepare_init_args_and_inputs_for_common()

      SCREAMING_SNAKE_CASE				    =		8

      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE				    =		self.model_class(**_A)
      model.to(_A)

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





     def     __UpperCamelCase						(		self					:	List[Any])							->	str:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.prepare_init_args_and_inputs_for_common()

      SCREAMING_SNAKE_CASE				    =		8

      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE				    =		self.model_class(**_A)
      model.to(_A)

      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		model(**_A).sample

      SCREAMING_SNAKE_CASE				    =		create_lora_layers(_A)
      model.set_attn_processor(_A)

      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		model(**_A  ,  cross_attention_kwargs={'scale': 0.0}).sample

      model.set_attn_processor(AttnProcessor())

      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		model(**_A).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     __UpperCamelCase						(		self					:	Tuple)							->	Dict:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.prepare_init_args_and_inputs_for_common()

      SCREAMING_SNAKE_CASE				    =		4

      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE				    =		self.model_class(**_A)
      model.to(_A)
      SCREAMING_SNAKE_CASE				    =		create_lora_layers(_A)
      model.set_attn_processor(_A)

      # default
      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		model(**_A).sample

       model.enable_xformers_memory_efficient_attention()
       SCREAMING_SNAKE_CASE				    =		model(**_A).sample

       model.disable_xformers_memory_efficient_attention()
       SCREAMING_SNAKE_CASE				    =		model(**_A).sample

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





     def     __UpperCamelCase						(		self					:	str)							->	Union[str, Any]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.prepare_init_args_and_inputs_for_common()
      SCREAMING_SNAKE_CASE				    =		3_2

      SCREAMING_SNAKE_CASE				    =		self.model_class(**_A)
      model.to(_A)
      model.eval()

      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		model(**_A)[0]

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






'''simple docstring'''





def 							snake_case				(		UpperCAmelCase						)->      int:







      """simple docstring"""
      if not isinstance(UpperCAmelCase   ,	UpperCAmelCase						):
            raise ValueError('double_factorial() only accepts integral values'						)
      if num < 0:
            raise ValueError('double_factorial() not defined for negative values'						)
      __A													=      1
      for i in range(UpperCAmelCase   ,	0   ,	-2						):
            value *= i
      return value


if __name__ == "__main__":
 import doctest

 doctest.testmod()






