




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

_lowercase		 =      logging.get_logger(__name__)



def 						lowerCAmelCase							(	lowerCamelCase ,     lowerCamelCase = True						):
						'''simple docstring'''


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

												snake_case__    :   Optional[int]							     =   LoRAAttnProcessor(hidden_size=lowerCamelCase ,     cross_attention_dim=lowerCamelCase						)
												snake_case__    :   Optional[Any]							     =   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    (     snake_case_  ,				snake_case_  ,				unittest.TestCase		):


				'''simple docstring'''


				_lowerCAmelCase		     =UNetaDConditionModel
				_lowerCAmelCase		     ="sample"

				@property
				def  _lowerCamelCase   (				self				:       Union[str, Any] ):
										snake_case__    :   Optional[Any]							     =   4
										snake_case__    :   Union[str, Any]							     =   4
										snake_case__    :   Any							     =   4
										snake_case__    :   int							     =   (3_2, 3_2)

										snake_case__    :   Tuple							     =   floats_tensor((batch_size, num_channels, num_frames) + sizes ).to(SCREAMING_SNAKE_CASE_ )
										snake_case__    :   Union[str, Any]							     =   torch.tensor([1_0] ).to(SCREAMING_SNAKE_CASE_ )
										snake_case__    :   List[str]							     =   floats_tensor((batch_size, 4, 3_2) ).to(SCREAMING_SNAKE_CASE_ )

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

				@property
				def  _lowerCamelCase   (				self				:       List[Any] ):
										return (4, 4, 3_2, 3_2)

				@property
				def  _lowerCamelCase   (				self				:       Optional[Any] ):
										return (4, 4, 3_2, 3_2)

				def  _lowerCamelCase   (				self				:       Any ):
										snake_case__    :   List[str]							     =   {
										    '''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,
										}
										snake_case__    :   str							     =   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  _lowerCamelCase   (				self				:       Optional[int] ):
										snake_case__    :   List[Any]							     =   self.prepare_init_args_and_inputs_for_common()
										snake_case__    :   List[str]							     =   self.model_class(**SCREAMING_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  _lowerCamelCase   (				self				:       Tuple ):
										snake_case__    :   int							     =   self.prepare_init_args_and_inputs_for_common()

										snake_case__    :   List[str]							     =   3_2

										snake_case__    :   Tuple							     =   self.model_class(**SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )
										model.eval()

										with torch.no_grad():
																snake_case__    :   List[str]							     =   model(**SCREAMING_SNAKE_CASE_ )

																if isinstance(SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ ):
																						snake_case__    :   Union[str, Any]							     =   output.sample

										self.assertIsNotNone(SCREAMING_SNAKE_CASE_ )
										snake_case__    :   Union[str, Any]							     =   inputs_dict['''sample'''].shape
										self.assertEqual(output.shape ,SCREAMING_SNAKE_CASE_ ,'''Input and output shapes do not match''' )

				def  _lowerCamelCase   (				self				:       str ):
										snake_case__    :   Optional[Any]							     =   self.prepare_init_args_and_inputs_for_common()
										snake_case__    :   List[str]							     =   self.model_class(**SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )
										model.eval()

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

																snake_case__    :   Optional[Any]							     =   model(**SCREAMING_SNAKE_CASE_ )
																if isinstance(SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ ):
																						snake_case__    :   Optional[Any]							     =   first.sample

																snake_case__    :   Any							     =   model(**SCREAMING_SNAKE_CASE_ )
																if isinstance(SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ ):
																						snake_case__    :   Union[str, Any]							     =   second.sample

										snake_case__    :   int							     =   first.cpu().numpy()
										snake_case__    :   Tuple							     =   second.cpu().numpy()
										snake_case__    :   Union[str, Any]							     =   out_a[~np.isnan(SCREAMING_SNAKE_CASE_ )]
										snake_case__    :   Optional[int]							     =   out_a[~np.isnan(SCREAMING_SNAKE_CASE_ )]
										snake_case__    :   Optional[int]							     =   np.amax(np.abs(out_a - out_a ) )
										self.assertLessEqual(SCREAMING_SNAKE_CASE_ ,1E-5 )

				def  _lowerCamelCase   (				self				:       int ):
										snake_case__    :   Optional[int]							     =   self.prepare_init_args_and_inputs_for_common()

										snake_case__    :   Tuple							     =   8

										snake_case__    :   Optional[Any]							     =   self.model_class(**SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )
										model.eval()

										model.set_attention_slice('''auto''' )
										with torch.no_grad():
																snake_case__    :   Optional[int]							     =   model(**SCREAMING_SNAKE_CASE_ )
										assert output is not None

										model.set_attention_slice('''max''' )
										with torch.no_grad():
																snake_case__    :   Optional[int]							     =   model(**SCREAMING_SNAKE_CASE_ )
										assert output is not None

										model.set_attention_slice(2 )
										with torch.no_grad():
																snake_case__    :   Any							     =   model(**SCREAMING_SNAKE_CASE_ )
										assert output is not None

				def  _lowerCamelCase   (				self				:       List[str] ):
										snake_case__    :   Tuple							     =   self.prepare_init_args_and_inputs_for_common()

										snake_case__    :   Tuple							     =   8

										snake_case__    :   str							     =   self.model_class(**SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )

										with torch.no_grad():
																snake_case__    :   Any							     =   model(**SCREAMING_SNAKE_CASE_ ).sample

										snake_case__    :   Optional[int]							     =   create_lora_layers(SCREAMING_SNAKE_CASE_ )

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

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

										with torch.no_grad():
																snake_case__    :   Tuple							     =   model(**SCREAMING_SNAKE_CASE_ ,cross_attention_kwargs={'''scale''': 0.0} ).sample
																snake_case__    :   Dict							     =   model(**SCREAMING_SNAKE_CASE_ ,cross_attention_kwargs={'''scale''': 0.5} ).sample
																snake_case__    :   Optional[Any]							     =   model(**SCREAMING_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  _lowerCamelCase   (				self				:       Tuple ):
										snake_case__    :   int							     =   self.prepare_init_args_and_inputs_for_common()

										snake_case__    :   List[str]							     =   8

										torch.manual_seed(0 )
										snake_case__    :   Optional[Any]							     =   self.model_class(**SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )

										with torch.no_grad():
																snake_case__    :   Optional[Any]							     =   model(**SCREAMING_SNAKE_CASE_ ).sample

										snake_case__    :   Tuple							     =   create_lora_layers(SCREAMING_SNAKE_CASE_ )
										model.set_attn_processor(SCREAMING_SNAKE_CASE_ )

										with torch.no_grad():
																snake_case__    :   Optional[Any]							     =   model(**SCREAMING_SNAKE_CASE_ ,cross_attention_kwargs={'''scale''': 0.5} ).sample

										with tempfile.TemporaryDirectory() as tmpdirname:
																model.save_attn_procs(SCREAMING_SNAKE_CASE_ )
																self.assertTrue(os.path.isfile(os.path.join(SCREAMING_SNAKE_CASE_ ,'''pytorch_lora_weights.bin''' ) ) )
																torch.manual_seed(0 )
																snake_case__    :   List[Any]							     =   self.model_class(**SCREAMING_SNAKE_CASE_ )
																new_model.to(SCREAMING_SNAKE_CASE_ )
																new_model.load_attn_procs(SCREAMING_SNAKE_CASE_ )

										with torch.no_grad():
																snake_case__    :   List[Any]							     =   new_model(**SCREAMING_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  _lowerCamelCase   (				self				:       Tuple ):
										snake_case__    :   List[str]							     =   self.prepare_init_args_and_inputs_for_common()

										snake_case__    :   str							     =   8

										torch.manual_seed(0 )
										snake_case__    :   Union[str, Any]							     =   self.model_class(**SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )

										with torch.no_grad():
																snake_case__    :   List[str]							     =   model(**SCREAMING_SNAKE_CASE_ ).sample

										snake_case__    :   Optional[Any]							     =   create_lora_layers(SCREAMING_SNAKE_CASE_ )
										model.set_attn_processor(SCREAMING_SNAKE_CASE_ )

										with torch.no_grad():
																snake_case__    :   Optional[Any]							     =   model(**SCREAMING_SNAKE_CASE_ ,cross_attention_kwargs={'''scale''': 0.5} ).sample

										with tempfile.TemporaryDirectory() as tmpdirname:
																model.save_attn_procs(SCREAMING_SNAKE_CASE_ ,safe_serialization=SCREAMING_SNAKE_CASE_ )
																self.assertTrue(os.path.isfile(os.path.join(SCREAMING_SNAKE_CASE_ ,'''pytorch_lora_weights.safetensors''' ) ) )
																torch.manual_seed(0 )
																snake_case__    :   Optional[int]							     =   self.model_class(**SCREAMING_SNAKE_CASE_ )
																new_model.to(SCREAMING_SNAKE_CASE_ )
																new_model.load_attn_procs(SCREAMING_SNAKE_CASE_ )

										with torch.no_grad():
																snake_case__    :   Tuple							     =   new_model(**SCREAMING_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  _lowerCamelCase   (				self				:       Optional[Any] ):
										# enable deterministic behavior for gradient checkpointing
										snake_case__    :   str							     =   self.prepare_init_args_and_inputs_for_common()

										snake_case__    :   List[Any]							     =   8

										torch.manual_seed(0 )
										snake_case__    :   List[str]							     =   self.model_class(**SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )

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

				def  _lowerCamelCase   (				self				:       Any ):
										snake_case__    :   int							     =   self.prepare_init_args_and_inputs_for_common()

										snake_case__    :   str							     =   8

										torch.manual_seed(0 )
										snake_case__    :   Any							     =   self.model_class(**SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )

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

				def  _lowerCamelCase   (				self				:       Tuple ):
										snake_case__    :   Tuple							     =   self.prepare_init_args_and_inputs_for_common()

										snake_case__    :   List[Any]							     =   8

										torch.manual_seed(0 )
										snake_case__    :   Optional[int]							     =   self.model_class(**SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )

										with torch.no_grad():
																snake_case__    :   int							     =   model(**SCREAMING_SNAKE_CASE_ ).sample

										snake_case__    :   List[Any]							     =   create_lora_layers(SCREAMING_SNAKE_CASE_ )
										model.set_attn_processor(SCREAMING_SNAKE_CASE_ )

										with torch.no_grad():
																snake_case__    :   List[Any]							     =   model(**SCREAMING_SNAKE_CASE_ ,cross_attention_kwargs={'''scale''': 0.0} ).sample

										model.set_attn_processor(AttnProcessor() )

										with torch.no_grad():
																snake_case__    :   Optional[Any]							     =   model(**SCREAMING_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  _lowerCamelCase   (				self				:       Optional[Any] ):
										# enable deterministic behavior for gradient checkpointing
										snake_case__    :   Tuple							     =   self.prepare_init_args_and_inputs_for_common()

										snake_case__    :   int							     =   4

										torch.manual_seed(0 )
										snake_case__    :   List[Any]							     =   self.model_class(**SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )
										snake_case__    :   Any							     =   create_lora_layers(SCREAMING_SNAKE_CASE_ )
										model.set_attn_processor(SCREAMING_SNAKE_CASE_ )

										# default
										with torch.no_grad():
																snake_case__    :   Union[str, Any]							     =   model(**SCREAMING_SNAKE_CASE_ ).sample

																model.enable_xformers_memory_efficient_attention()
																snake_case__    :   Dict							     =   model(**SCREAMING_SNAKE_CASE_ ).sample

																model.disable_xformers_memory_efficient_attention()
																snake_case__    :   str							     =   model(**SCREAMING_SNAKE_CASE_ ).sample

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





				def  _lowerCamelCase   (				self				:       Tuple ):
										snake_case__    :   int							     =   self.prepare_init_args_and_inputs_for_common()
										snake_case__    :   str							     =   3_2

										snake_case__    :   Optional[int]							     =   self.model_class(**SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )
										model.eval()

										with torch.no_grad():
																snake_case__    :   Dict							     =   model(**SCREAMING_SNAKE_CASE_ )[0]

										model.enable_forward_chunking()
										with torch.no_grad():
																snake_case__    :   Optional[int]							     =   model(**SCREAMING_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.








"""simple docstring"""





from __future__ import annotations

import math
import random

def 						lowerCAmelCase							(	lowerCamelCase ,     lowerCamelCase ,     lowerCamelCase						):
						'''simple docstring'''


						snake_case__    :   Tuple							     =   d * e - 1
						snake_case__    :   Optional[int]							     =   0
						snake_case__    :   Tuple							     =   0
						while p == 0:
												snake_case__    :   Optional[int]							     =   random.randint(2 ,     n - 1						)
												snake_case__    :   Tuple							     =   k
												while True:
																		if t % 2 == 0:
																								snake_case__    :   List[Any]							     =   t // 2
																								snake_case__    :   Union[str, Any]							     =   (g**t) % n
																								snake_case__    :   Dict							     =   math.gcd(x - 1 ,     lowerCamelCase						)
																								if x > 1 and y > 1:
																														snake_case__    :   Union[str, Any]							     =   y
																														snake_case__    :   str							     =   n // y
																														break  # find the correct factors
																		else:
																								break  # t is not divisible by 2, break and choose another g
						return sorted([p, q]						)


if __name__ == "__main__":
							import doctest

							doctest.testmod()


