

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

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



 '''simple docstring'''


 __SCREAMING_SNAKE_CASE												=				UNetaDConditionModel
 __SCREAMING_SNAKE_CASE												=				"""sample"""




 @property
 def   UpperCamelCase__							(self ) ->					Any:






  """simple docstring"""





  UpperCAmelCase__							=				4
  UpperCAmelCase__							=				4
  UpperCAmelCase__							=				4
  UpperCAmelCase__							=				(32, 32)

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

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




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






  """simple docstring"""





  return (4, 4, 32, 32)




 @property
 def   UpperCamelCase__							(self ) ->					Dict:






  """simple docstring"""





  return (4, 4, 32, 32)




 def   UpperCamelCase__							(self ) ->					Optional[int]:






  """simple docstring"""





  UpperCAmelCase__							=				{
      '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,
  }
  UpperCAmelCase__							=				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 ) ->					int:






  """simple docstring"""





  UpperCAmelCase__				,    UpperCAmelCase__							=				self.prepare_init_args_and_inputs_for_common()
  UpperCAmelCase__							=				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 ) ->					Optional[int]:






  """simple docstring"""





  UpperCAmelCase__				,    UpperCAmelCase__							=				self.prepare_init_args_and_inputs_for_common()

  UpperCAmelCase__							=				32

  UpperCAmelCase__							=				self.model_class(**__a )
  model.to(__a )
  model.eval()

  with torch.no_grad():
   UpperCAmelCase__							=				model(**__a )

   if isinstance(__a       ,       __a ):
    UpperCAmelCase__							=				output.sample

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




 def   UpperCamelCase__							(self ) ->					List[Any]:






  """simple docstring"""





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

   UpperCAmelCase__							=				model(**__a )
   if isinstance(__a       ,       __a ):
    UpperCAmelCase__							=				first.sample

   UpperCAmelCase__							=				model(**__a )
   if isinstance(__a       ,       __a ):
    UpperCAmelCase__							=				second.sample

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




 def   UpperCamelCase__							(self ) ->					Any:






  """simple docstring"""





  UpperCAmelCase__				,    UpperCAmelCase__							=				self.prepare_init_args_and_inputs_for_common()

  UpperCAmelCase__							=				8

  UpperCAmelCase__							=				self.model_class(**__a )
  model.to(__a )
  model.eval()

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

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

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




 def   UpperCamelCase__							(self ) ->					Dict:






  """simple docstring"""





  UpperCAmelCase__				,    UpperCAmelCase__							=				self.prepare_init_args_and_inputs_for_common()

  UpperCAmelCase__							=				8

  UpperCAmelCase__							=				self.model_class(**__a )
  model.to(__a )

  with torch.no_grad():
   UpperCAmelCase__							=				model(**__a ).sample

  UpperCAmelCase__							=				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():
   UpperCAmelCase__							=				model(**__a       ,       cross_attention_kwargs={'scale': 0.0} ).sample
   UpperCAmelCase__							=				model(**__a       ,       cross_attention_kwargs={'scale': 0.5} ).sample
   UpperCAmelCase__							=				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[Any]:






  """simple docstring"""





  UpperCAmelCase__				,    UpperCAmelCase__							=				self.prepare_init_args_and_inputs_for_common()

  UpperCAmelCase__							=				8

  torch.manual_seed(0 )
  UpperCAmelCase__							=				self.model_class(**__a )
  model.to(__a )

  with torch.no_grad():
   UpperCAmelCase__							=				model(**__a ).sample

  UpperCAmelCase__							=				create_lora_layers(__a )
  model.set_attn_processor(__a )

  with torch.no_grad():
   UpperCAmelCase__							=				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 )
   UpperCAmelCase__							=				self.model_class(**__a )
   new_model.to(__a )
   new_model.load_attn_procs(__a )

  with torch.no_grad():
   UpperCAmelCase__							=				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 ) ->					int:






  """simple docstring"""





  UpperCAmelCase__				,    UpperCAmelCase__							=				self.prepare_init_args_and_inputs_for_common()

  UpperCAmelCase__							=				8

  torch.manual_seed(0 )
  UpperCAmelCase__							=				self.model_class(**__a )
  model.to(__a )

  with torch.no_grad():
   UpperCAmelCase__							=				model(**__a ).sample

  UpperCAmelCase__							=				create_lora_layers(__a )
  model.set_attn_processor(__a )

  with torch.no_grad():
   UpperCAmelCase__							=				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 )
   UpperCAmelCase__							=				self.model_class(**__a )
   new_model.to(__a )
   new_model.load_attn_procs(__a )

  with torch.no_grad():
   UpperCAmelCase__							=				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 ) ->					Union[str, Any]:






  """simple docstring"""





  UpperCAmelCase__				,    UpperCAmelCase__							=				self.prepare_init_args_and_inputs_for_common()

  UpperCAmelCase__							=				8

  torch.manual_seed(0 )
  UpperCAmelCase__							=				self.model_class(**__a )
  model.to(__a )

  UpperCAmelCase__							=				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 )
   UpperCAmelCase__							=				self.model_class(**__a )
   new_model.to(__a )
   new_model.load_attn_procs(__a       ,       weight_name='pytorch_lora_weights.bin' )




 def   UpperCamelCase__							(self ) ->					Optional[Any]:






  """simple docstring"""





  UpperCAmelCase__				,    UpperCAmelCase__							=				self.prepare_init_args_and_inputs_for_common()

  UpperCAmelCase__							=				8

  torch.manual_seed(0 )
  UpperCAmelCase__							=				self.model_class(**__a )
  model.to(__a )

  UpperCAmelCase__							=				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 )
   UpperCAmelCase__							=				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 ) ->					Any:






  """simple docstring"""





  UpperCAmelCase__				,    UpperCAmelCase__							=				self.prepare_init_args_and_inputs_for_common()

  UpperCAmelCase__							=				8

  torch.manual_seed(0 )
  UpperCAmelCase__							=				self.model_class(**__a )
  model.to(__a )

  with torch.no_grad():
   UpperCAmelCase__							=				model(**__a ).sample

  UpperCAmelCase__							=				create_lora_layers(__a )
  model.set_attn_processor(__a )

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

  model.set_attn_processor(AttnProcessor() )

  with torch.no_grad():
   UpperCAmelCase__							=				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 ) ->					int:






  """simple docstring"""





  UpperCAmelCase__				,    UpperCAmelCase__							=				self.prepare_init_args_and_inputs_for_common()

  UpperCAmelCase__							=				4

  torch.manual_seed(0 )
  UpperCAmelCase__							=				self.model_class(**__a )
  model.to(__a )
  UpperCAmelCase__							=				create_lora_layers(__a )
  model.set_attn_processor(__a )

  # default
  with torch.no_grad():
   UpperCAmelCase__							=				model(**__a ).sample

   model.enable_xformers_memory_efficient_attention()
   UpperCAmelCase__							=				model(**__a ).sample

   model.disable_xformers_memory_efficient_attention()
   UpperCAmelCase__							=				model(**__a ).sample

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




 def   UpperCamelCase__							(self ) ->					Optional[Any]:






  """simple docstring"""





  UpperCAmelCase__				,    UpperCAmelCase__							=				self.prepare_init_args_and_inputs_for_common()
  UpperCAmelCase__							=				32

  UpperCAmelCase__							=				self.model_class(**__a )
  model.to(__a )
  model.eval()

  with torch.no_grad():
   UpperCAmelCase__							=				model(**__a )[0]

  model.enable_forward_chunking()
  with torch.no_grad():
   UpperCAmelCase__							=				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		(					_SCREAMING_SNAKE_CASE     :						int = 4_000_000							)			->						int:







	"""simple docstring"""
	lowerCAmelCase    =	1
	lowerCAmelCase    =	2
	lowerCAmelCase    =	0
	while j <= n:
		if j % 2 == 0:
			total += j
		lowerCAmelCase,   lowerCAmelCase    =	j, i + j

	return total


if __name__ == "__main__":
				print(F'''{solution() = }''')