





'''simple docstring'''







import textwrap

import pyarrow as pa
import pytest

from datasets import Features, Image
from datasets.packaged_modules.text.text import Text

from ..utils import require_pil





@pytest.fixture
def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      List[Any] )			->						str:
						UpperCAmelCase_      :   Tuple      =	tmp_path / '''text.txt'''
						UpperCAmelCase_      :   Optional[int]      =	textwrap.dedent(
						    '''\
        Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
        Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
        Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.
        Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.

        Second paragraph:
        Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
        Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
        Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.
        Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
        ''' )
						with open(SCREAMING_SNAKE_CASE__,      '''w''',      encoding='''utf-8''' ) as f:
												f.write(SCREAMING_SNAKE_CASE__ )
						return str(SCREAMING_SNAKE_CASE__ )





@pytest.fixture
def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Tuple,      SCREAMING_SNAKE_CASE__   :      List[str] )			->						Any:
						UpperCAmelCase_      :   Tuple      =	tmp_path / '''text_with_image.txt'''
						with open(SCREAMING_SNAKE_CASE__,      '''w''',      encoding='''utf-8''' ) as f:
												f.write(SCREAMING_SNAKE_CASE__ )
						return str(SCREAMING_SNAKE_CASE__ )





@pytest.mark.parametrize('''keep_linebreaks''',      [True, False] )
def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Dict,      SCREAMING_SNAKE_CASE__   :      str )			->						int:
						with open(SCREAMING_SNAKE_CASE__,      encoding='''utf-8''' ) as f:
												UpperCAmelCase_      :   Tuple      =	f.read().splitlines(keepends=SCREAMING_SNAKE_CASE__ )
						UpperCAmelCase_      :   str      =	Text(keep_linebreaks=SCREAMING_SNAKE_CASE__,      encoding='''utf-8''' )
						UpperCAmelCase_      :   Optional[Any]      =	text._generate_tables([[text_file]] )
						UpperCAmelCase_      :   Optional[Any]      =	pa.concat_tables([table for _, table in generator] ).to_pydict()['''text''']
						assert generated_content == expected_content





@require_pil
def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      List[Any] )			->						str:
						with open(SCREAMING_SNAKE_CASE__,      encoding='''utf-8''' ) as f:
												UpperCAmelCase_      :   Tuple      =	f.read().splitlines()[0]
						UpperCAmelCase_      :   Optional[int]      =	Text(encoding='''utf-8''',      features=Features({'''image''': Image()} ) )
						UpperCAmelCase_      :   Optional[int]      =	text._generate_tables([[text_file_with_image]] )
						UpperCAmelCase_      :   List[Any]      =	pa.concat_tables([table for _, table in generator] )
						assert pa_table.schema.field('''image''' ).type == Image()()
						UpperCAmelCase_      :   Tuple      =	pa_table.to_pydict()['''image''']
						assert generated_content == [{"path": image_file, "bytes": None}]





@pytest.mark.parametrize('''sample_by''',      ['''line''', '''paragraph''', '''document'''] )
def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Optional[int],      SCREAMING_SNAKE_CASE__   :      Optional[int] )			->						Optional[int]:
						with open(SCREAMING_SNAKE_CASE__,      encoding='''utf-8''' ) as f:
												UpperCAmelCase_      :   List[Any]      =	f.read()
						if sample_by == "line":
												UpperCAmelCase_      :   Union[str, Any]      =	expected_content.splitlines()
						elif sample_by == "paragraph":
												UpperCAmelCase_      :   Optional[int]      =	expected_content.split('''\n\n''' )
						elif sample_by == "document":
												UpperCAmelCase_      :   List[str]      =	[expected_content]
						UpperCAmelCase_      :   Optional[int]      =	Text(sample_by=SCREAMING_SNAKE_CASE__,      encoding='''utf-8''',      chunksize=100 )
						UpperCAmelCase_      :   Dict      =	text._generate_tables([[text_file]] )
						UpperCAmelCase_      :   Union[str, Any]      =	pa.concat_tables([table for _, table in generator] ).to_pydict()['''text''']
						assert generated_content == expected_content













"""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.






