






"""simple docstring"""



import unittest

from transformers import load_tool
from transformers.tools.agent_types import AGENT_TYPE_MAPPING

from .test_tools_common import ToolTesterMixin, output_types






class 				__SCREAMING_SNAKE_CASE    (					unittest.TestCase							,		UpperCamelCase		):


							def     _UpperCamelCase							(      self   :    Tuple				):
													'''simple docstring'''





													A__    :	str			=	load_tool("""translation"""				)
													self.tool.setup()
													A__    :	List[str]			=	load_tool("""translation"""      ,     remote=snake_case				)


							def     _UpperCamelCase							(      self   :    Optional[Any]				):
													'''simple docstring'''





													A__    :	Tuple			=	self.tool("""Hey, what's up?"""      ,     src_lang="""English"""      ,     tgt_lang="""French"""				)
													self.assertEqual(snake_case      ,     """- Hé, comment ça va?"""				)


							def     _UpperCamelCase							(      self   :    Any				):
													'''simple docstring'''





													A__    :	Any			=	self.remote_tool("""Hey, what's up?"""      ,     src_lang="""English"""      ,     tgt_lang="""French"""				)
													self.assertEqual(snake_case      ,     """- Hé, comment ça va?"""				)


							def     _UpperCamelCase							(      self   :    Union[str, Any]				):
													'''simple docstring'''





													A__    :	str			=	self.tool(text="""Hey, what's up?"""      ,     src_lang="""English"""      ,     tgt_lang="""French"""				)
													self.assertEqual(snake_case      ,     """- Hé, comment ça va?"""				)


							def     _UpperCamelCase							(      self   :    Tuple				):
													'''simple docstring'''





													A__    :	Optional[Any]			=	self.remote_tool(text="""Hey, what's up?"""      ,     src_lang="""English"""      ,     tgt_lang="""French"""				)
													self.assertEqual(snake_case      ,     """- Hé, comment ça va?"""				)


							def     _UpperCamelCase							(      self   :    Tuple				):
													'''simple docstring'''





													A__    :	Union[str, Any]			=	["""Hey, what's up?""", """English""", """Spanish"""]
													A__    :	Any			=	self.tool(*snake_case				)

													# There is a single output
													if len(self.tool.outputs				) == 1:
																			A__    :	int			=	[outputs]

													self.assertListEqual(output_types(snake_case				)      ,     self.tool.outputs				)


							def     _UpperCamelCase							(      self   :    Optional[Any]				):
													'''simple docstring'''





													A__    :	Union[str, Any]			=	["""Hey, what's up?""", """English""", """Spanish"""]
													A__    :	Tuple			=	self.tool(*snake_case				)

													if not isinstance(snake_case      ,     snake_case				):
																			A__    :	List[Any]			=	[outputs]

													self.assertEqual(len(snake_case				)      ,     len(self.tool.outputs				)				)

													for output, output_type in zip(snake_case      ,     self.tool.outputs				):
																			A__    :	List[Any]			=	AGENT_TYPE_MAPPING[output_type]
																			self.assertTrue(isinstance(snake_case      ,     snake_case				)				)


							def     _UpperCamelCase							(      self   :    Dict				):
													'''simple docstring'''





													A__    :	List[Any]			=	["""Hey, what's up?""", """English""", """Spanish"""]

													A__    :	str			=	[]

													for _input, input_type in zip(snake_case      ,     self.tool.inputs				):
																			if isinstance(snake_case      ,     snake_case				):
																									_inputs.append([AGENT_TYPE_MAPPING[_input_type](_input				) for _input_type in input_type]				)
																			else:
																									_inputs.append(AGENT_TYPE_MAPPING[input_type](_input				)				)

        # Should not raise an error
													A__    :	Dict			=	self.tool(*snake_case				)

													if not isinstance(snake_case      ,     snake_case				):
																			A__    :	int			=	[outputs]

													self.assertEqual(len(snake_case				)      ,     len(self.tool.outputs				)				)














"""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_     =				logging.get_logger(__name__)




def   _lowerCAmelCase		(				UpperCAmelCase__							:       int,    UpperCAmelCase__							:       bool = True   )		->Tuple:
						A__    :	List[str]			=	{}
						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__    :	Tuple			=	model.config.cross_attention_dim if has_cross_attention else None
												if name.startswith("""mid_block"""   ):
																		A__    :	Tuple			=	model.config.block_out_channels[-1]
												elif name.startswith("""up_blocks"""   ):
																		A__    :	str			=	int(name[len("""up_blocks."""   )]   )
																		A__    :	Tuple			=	list(reversed(model.config.block_out_channels   )   )[block_id]
												elif name.startswith("""down_blocks"""   ):
																		A__    :	List[str]			=	int(name[len("""down_blocks."""   )]   )
																		A__    :	Tuple			=	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__    :	List[str]			=	8 * model.config.attention_head_dim

												A__    :	Optional[Any]			=	LoRAAttnProcessor(hidden_size=UpperCAmelCase__,    cross_attention_dim=UpperCAmelCase__   )
												A__    :	Optional[int]			=	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 				__SCREAMING_SNAKE_CASE    (					UpperCamelCase							,		UpperCamelCase							,		unittest.TestCase		):
							snake_case_  =				UNetaDConditionModel
							snake_case_  =				'sample'


							@property
							def     _UpperCamelCase							(      self   :    str				):
													'''simple docstring'''





													A__    :	Dict			=	4
													A__    :	str			=	4
													A__    :	str			=	4
													A__    :	Union[str, Any]			=	(32, 32)

													A__    :	List[Any]			=	floats_tensor((batch_size, num_channels, num_frames) + sizes				).to(snake_case				)
													A__    :	Any			=	torch.tensor([10]				).to(snake_case				)
													A__    :	Optional[Any]			=	floats_tensor((batch_size, 4, 32)				).to(snake_case				)

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


							@property
							def     _UpperCamelCase							(      self   :    Union[str, Any]				):
													'''simple docstring'''





													return (4, 4, 32, 32)


							@property
							def     _UpperCamelCase							(      self   :    Optional[int]				):
													'''simple docstring'''





													return (4, 4, 32, 32)


							def     _UpperCamelCase							(      self   :    int				):
													'''simple docstring'''





													A__    :	Tuple			=	{
													    """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__    :	Optional[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     _UpperCamelCase							(      self   :    int				):
													'''simple docstring'''





													A__	,	A__    :	Optional[Any]			=	self.prepare_init_args_and_inputs_for_common()
													A__    :	Dict			=	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     _UpperCamelCase							(      self   :    List[str]				):
													'''simple docstring'''





													A__	,	A__    :	Optional[int]			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	Tuple			=	32

													A__    :	List[Any]			=	self.model_class(**snake_case				)
													model.to(snake_case				)
													model.eval()

													with torch.no_grad():
																			A__    :	Dict			=	model(**snake_case				)

																			if isinstance(snake_case      ,     snake_case				):
																									A__    :	int			=	output.sample

													self.assertIsNotNone(snake_case				)
													A__    :	Any			=	inputs_dict["""sample"""].shape
													self.assertEqual(output.shape      ,     snake_case      ,     """Input and output shapes do not match"""				)


							def     _UpperCamelCase							(      self   :    int				):
													'''simple docstring'''





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

																			A__    :	List[str]			=	model(**snake_case				)
																			if isinstance(snake_case      ,     snake_case				):
																									A__    :	List[Any]			=	first.sample

																			A__    :	List[str]			=	model(**snake_case				)
																			if isinstance(snake_case      ,     snake_case				):
																									A__    :	str			=	second.sample

													A__    :	Optional[Any]			=	first.cpu().numpy()
													A__    :	Tuple			=	second.cpu().numpy()
													A__    :	List[Any]			=	out_a[~np.isnan(snake_case				)]
													A__    :	Dict			=	out_a[~np.isnan(snake_case				)]
													A__    :	str			=	np.amax(np.abs(out_a - out_a				)				)
													self.assertLessEqual(snake_case      ,     1e-5				)


							def     _UpperCamelCase							(      self   :    Any				):
													'''simple docstring'''





													A__	,	A__    :	List[Any]			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	List[str]			=	8

													A__    :	Optional[Any]			=	self.model_class(**snake_case				)
													model.to(snake_case				)
													model.eval()

													model.set_attention_slice("""auto"""				)
													with torch.no_grad():
																			A__    :	List[str]			=	model(**snake_case				)
													assert output is not None

													model.set_attention_slice("""max"""				)
													with torch.no_grad():
																			A__    :	Union[str, Any]			=	model(**snake_case				)
													assert output is not None

													model.set_attention_slice(2				)
													with torch.no_grad():
																			A__    :	Dict			=	model(**snake_case				)
													assert output is not None


							def     _UpperCamelCase							(      self   :    str				):
													'''simple docstring'''





													A__	,	A__    :	Dict			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	List[Any]			=	8

													A__    :	Union[str, Any]			=	self.model_class(**snake_case				)
													model.to(snake_case				)

													with torch.no_grad():
																			A__    :	Optional[int]			=	model(**snake_case				).sample

													A__    :	Dict			=	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():
																			A__    :	int			=	model(**snake_case      ,     cross_attention_kwargs={"""scale""": 0.0}				).sample
																			A__    :	Optional[Any]			=	model(**snake_case      ,     cross_attention_kwargs={"""scale""": 0.5}				).sample
																			A__    :	List[Any]			=	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     _UpperCamelCase							(      self   :    str				):
													'''simple docstring'''





													A__	,	A__    :	Dict			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	Optional[Any]			=	8

													torch.manual_seed(0				)
													A__    :	str			=	self.model_class(**snake_case				)
													model.to(snake_case				)

													with torch.no_grad():
																			A__    :	Any			=	model(**snake_case				).sample

													A__    :	str			=	create_lora_layers(snake_case				)
													model.set_attn_processor(snake_case				)

													with torch.no_grad():
																			A__    :	Union[str, Any]			=	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				)
																			A__    :	Union[str, Any]			=	self.model_class(**snake_case				)
																			new_model.to(snake_case				)
																			new_model.load_attn_procs(snake_case				)

													with torch.no_grad():
																			A__    :	List[Any]			=	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     _UpperCamelCase							(      self   :    List[Any]				):
													'''simple docstring'''





													A__	,	A__    :	Optional[Any]			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	Optional[Any]			=	8

													torch.manual_seed(0				)
													A__    :	Any			=	self.model_class(**snake_case				)
													model.to(snake_case				)

													with torch.no_grad():
																			A__    :	Tuple			=	model(**snake_case				).sample

													A__    :	Union[str, Any]			=	create_lora_layers(snake_case				)
													model.set_attn_processor(snake_case				)

													with torch.no_grad():
																			A__    :	Tuple			=	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				)
																			A__    :	Optional[Any]			=	self.model_class(**snake_case				)
																			new_model.to(snake_case				)
																			new_model.load_attn_procs(snake_case				)

													with torch.no_grad():
																			A__    :	Any			=	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     _UpperCamelCase							(      self   :    List[Any]				):
													'''simple docstring'''





													A__	,	A__    :	List[Any]			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	Any			=	8

													torch.manual_seed(0				)
													A__    :	str			=	self.model_class(**snake_case				)
													model.to(snake_case				)

													A__    :	int			=	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				)
																			A__    :	str			=	self.model_class(**snake_case				)
																			new_model.to(snake_case				)
																			new_model.load_attn_procs(snake_case      ,     weight_name="""pytorch_lora_weights.bin"""				)


							def     _UpperCamelCase							(      self   :    Union[str, Any]				):
													'''simple docstring'''





													A__	,	A__    :	Union[str, Any]			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	Optional[Any]			=	8

													torch.manual_seed(0				)
													A__    :	Any			=	self.model_class(**snake_case				)
													model.to(snake_case				)

													A__    :	Dict			=	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				)
																			A__    :	Optional[Any]			=	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     _UpperCamelCase							(      self   :    List[str]				):
													'''simple docstring'''





													A__	,	A__    :	Dict			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	Tuple			=	8

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

													with torch.no_grad():
																			A__    :	Dict			=	model(**snake_case				).sample

													A__    :	Union[str, Any]			=	create_lora_layers(snake_case				)
													model.set_attn_processor(snake_case				)

													with torch.no_grad():
																			A__    :	Any			=	model(**snake_case      ,     cross_attention_kwargs={"""scale""": 0.0}				).sample

													model.set_attn_processor(AttnProcessor()				)

													with torch.no_grad():
																			A__    :	str			=	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     _UpperCamelCase							(      self   :    List[str]				):
													'''simple docstring'''





													A__	,	A__    :	List[str]			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	List[Any]			=	4

													torch.manual_seed(0				)
													A__    :	Tuple			=	self.model_class(**snake_case				)
													model.to(snake_case				)
													A__    :	List[str]			=	create_lora_layers(snake_case				)
													model.set_attn_processor(snake_case				)

													# default
													with torch.no_grad():
																			A__    :	Optional[Any]			=	model(**snake_case				).sample

																			model.enable_xformers_memory_efficient_attention()
																			A__    :	List[str]			=	model(**snake_case				).sample

																			model.disable_xformers_memory_efficient_attention()
																			A__    :	Union[str, Any]			=	model(**snake_case				).sample

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


							def     _UpperCamelCase							(      self   :    Any				):
													'''simple docstring'''





													A__	,	A__    :	Optional[Any]			=	self.prepare_init_args_and_inputs_for_common()
													A__    :	List[Any]			=	32

													A__    :	Any			=	self.model_class(**snake_case				)
													model.to(snake_case				)
													model.eval()

													with torch.no_grad():
																			A__    :	List[Any]			=	model(**snake_case				)[0]

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






