
"""simple docstring"""




import gc
import unittest

import numpy as np
import torch

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, DPMSolverMultistepScheduler, TransformeraDModel
from diffusers.utils import is_xformers_available, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS,
    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class 	lowerCAmelCase_						(a__ ,					unittest.TestCase  ):


				"""simple docstring"""



				__UpperCamelCase      :						Dict					       =   DiTPipeline
				__UpperCamelCase      :						List[Any]					       =   CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
				__UpperCamelCase      :						Union[str, Any]					       =   PipelineTesterMixin.required_optional_params - {
				    '''latents''',
				    '''num_images_per_prompt''',
				    '''callback''',
				    '''callback_steps''',
				}
				__UpperCamelCase      :						int					       =   CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
				__UpperCamelCase      :						List[str]					       =   False



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



											"""simple docstring"""




											torch.manual_seed(0    )
											SCREAMING_SNAKE_CASE__						:       Dict           =   TransformeraDModel(
											    sample_size=16   ,       num_layers=2   ,       patch_size=4   ,       attention_head_dim=8   ,       num_attention_heads=2   ,       in_channels=4   ,       out_channels=8   ,       attention_bias=SCREAMING_SNAKE_CASE__   ,       activation_fn="""gelu-approximate"""   ,       num_embeds_ada_norm=10_00   ,       norm_type="""ada_norm_zero"""   ,       norm_elementwise_affine=SCREAMING_SNAKE_CASE__   ,       )
											SCREAMING_SNAKE_CASE__						:       Dict           =   AutoencoderKL()
											SCREAMING_SNAKE_CASE__						:       Tuple           =   DDIMScheduler()
											SCREAMING_SNAKE_CASE__						:       List[Any]           =   {"""transformer""": transformer.eval(), """vae""": vae.eval(), """scheduler""": scheduler}
											return components



				def 	__magic_name__    (self   ,       SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__=0    )					->		Dict:



											"""simple docstring"""




											if str(SCREAMING_SNAKE_CASE__    ).startswith("""mps"""    ):
																		SCREAMING_SNAKE_CASE__						:       Optional[int]           =   torch.manual_seed(SCREAMING_SNAKE_CASE__    )
											else:
																		SCREAMING_SNAKE_CASE__						:       Optional[int]           =   torch.Generator(device=SCREAMING_SNAKE_CASE__    ).manual_seed(SCREAMING_SNAKE_CASE__    )
											SCREAMING_SNAKE_CASE__						:       List[Any]           =   {
											    """class_labels""": [1],
											    """generator""": generator,
											    """num_inference_steps""": 2,
											    """output_type""": """numpy""",
											}
											return inputs



				def 	__magic_name__    (self    )					->		Dict:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:       Tuple           =   """cpu"""

											SCREAMING_SNAKE_CASE__						:       Dict           =   self.get_dummy_components()
											SCREAMING_SNAKE_CASE__						:       Optional[int]           =   self.pipeline_class(**SCREAMING_SNAKE_CASE__    )
											pipe.to(SCREAMING_SNAKE_CASE__    )
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__    )

											SCREAMING_SNAKE_CASE__						:       List[str]           =   self.get_dummy_inputs(SCREAMING_SNAKE_CASE__    )
											SCREAMING_SNAKE_CASE__						:       Dict           =   pipe(**SCREAMING_SNAKE_CASE__    ).images
											SCREAMING_SNAKE_CASE__						:       Dict           =   image[0, -3:, -3:, -1]

											self.assertEqual(image.shape   ,       (1, 16, 16, 3)    )
											SCREAMING_SNAKE_CASE__						:       Optional[int]           =   np.array([0.2946, 0.6601, 0.4329, 0.3296, 0.4144, 0.5319, 0.7273, 0.5013, 0.4457]    )
											SCREAMING_SNAKE_CASE__						:       Optional[int]           =   np.abs(image_slice.flatten() - expected_slice    ).max()
											self.assertLessEqual(SCREAMING_SNAKE_CASE__   ,       1E-3    )



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



											"""simple docstring"""




											self._test_inference_batch_single_identical(relax_max_difference=SCREAMING_SNAKE_CASE__   ,       expected_max_diff=1E-3    )



				@unittest.skipIf(
				    torch_device != """cuda""" or not is_xformers_available()   ,       reason="""XFormers attention is only available with CUDA and `xformers` installed"""   ,       )
				def 	__magic_name__    (self    )					->		Tuple:



											"""simple docstring"""




											self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3    )






@require_torch_gpu
@slow
class 	lowerCAmelCase_						(unittest.TestCase  ):


				"""simple docstring"""



				def 	__magic_name__    (self    )					->		str:



											"""simple docstring"""




											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()



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



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   torch.manual_seed(0    )

											SCREAMING_SNAKE_CASE__						:       str           =   DiTPipeline.from_pretrained("""facebook/DiT-XL-2-256"""    )
											pipe.to("""cuda"""    )

											SCREAMING_SNAKE_CASE__						:       Dict           =   ["""vase""", """umbrella""", """white shark""", """white wolf"""]
											SCREAMING_SNAKE_CASE__						:       int           =   pipe.get_label_ids(SCREAMING_SNAKE_CASE__    )

											SCREAMING_SNAKE_CASE__						:       Tuple           =   pipe(SCREAMING_SNAKE_CASE__   ,       generator=SCREAMING_SNAKE_CASE__   ,       num_inference_steps=40   ,       output_type="""np"""    ).images

											for word, image in zip(SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__    ):
																		SCREAMING_SNAKE_CASE__						:       Tuple           =   load_numpy(
																		    F'''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/dit/{word}.npy'''    )
																		assert np.abs((expected_image - image).max()    ) < 1E-2



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



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:       str           =   DiTPipeline.from_pretrained("""facebook/DiT-XL-2-512"""    )
											SCREAMING_SNAKE_CASE__						:       str           =   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config    )
											pipe.to("""cuda"""    )

											SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   ["""vase""", """umbrella"""]
											SCREAMING_SNAKE_CASE__						:       int           =   pipe.get_label_ids(SCREAMING_SNAKE_CASE__    )

											SCREAMING_SNAKE_CASE__						:       Tuple           =   torch.manual_seed(0    )
											SCREAMING_SNAKE_CASE__						:       List[str]           =   pipe(SCREAMING_SNAKE_CASE__   ,       generator=SCREAMING_SNAKE_CASE__   ,       num_inference_steps=25   ,       output_type="""np"""    ).images

											for word, image in zip(SCREAMING_SNAKE_CASE__   ,       SCREAMING_SNAKE_CASE__    ):
																		SCREAMING_SNAKE_CASE__						:       List[str]           =   load_numpy(
																		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
																		    F'''/dit/{word}_512.npy'''    )

																		assert np.abs((expected_image - image).max()    ) < 1E-1






"""simple docstring"""




import mpmath  # for roots of unity
import numpy as np



class 	lowerCAmelCase_						:


				"""simple docstring"""



				def __init__(self   ,       SCREAMING_SNAKE_CASE__=None   ,       SCREAMING_SNAKE_CASE__=None    )					->		Dict:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   list(poly_a or [0]    )[:]
											SCREAMING_SNAKE_CASE__						:       Tuple           =   list(poly_b or [0]    )[:]

											# Remove leading zero coefficients
											while self.polyA[-1] == 0:
																		self.polyA.pop()
											SCREAMING_SNAKE_CASE__						:       int           =   len(self.polyA    )

											while self.polyB[-1] == 0:
																		self.polyB.pop()
											SCREAMING_SNAKE_CASE__						:       List[str]           =   len(self.polyB    )

											# Add 0 to make lengths equal a power of 2
											SCREAMING_SNAKE_CASE__						:       Optional[int]           =   int(
											    2 ** np.ceil(np.loga(len(self.polyA    ) + len(self.polyB    ) - 1    )    )    )

											while len(self.polyA    ) < self.c_max_length:
																		self.polyA.append(0    )
											while len(self.polyB    ) < self.c_max_length:
																		self.polyB.append(0    )
											# A complex root used for the fourier transform
											SCREAMING_SNAKE_CASE__						:       List[str]           =   complex(mpmath.root(x=1   ,       n=self.c_max_length   ,       k=1    )    )

											# The product
											SCREAMING_SNAKE_CASE__						:       Tuple           =   self.__multiply()



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



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:       str           =   [[x] for x in self.polyA] if which == """A""" else [[x] for x in self.polyB]
											# Corner case
											if len(SCREAMING_SNAKE_CASE__    ) <= 1:
																		return dft[0]
											#
											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   self.c_max_length // 2
											while next_ncol > 0:
																		SCREAMING_SNAKE_CASE__						:       Any           =   [[] for i in range(SCREAMING_SNAKE_CASE__    )]
																		SCREAMING_SNAKE_CASE__						:       Tuple           =   self.root**next_ncol

																		# First half of next step
																		SCREAMING_SNAKE_CASE__						:       str           =   1
																		for j in range(self.c_max_length // (next_ncol * 2)    ):
																									for i in range(SCREAMING_SNAKE_CASE__    ):
																																new_dft[i].append(dft[i][j] + current_root * dft[i + next_ncol][j]    )
																									current_root *= root
																		# Second half of next step
																		SCREAMING_SNAKE_CASE__						:       int           =   1
																		for j in range(self.c_max_length // (next_ncol * 2)    ):
																									for i in range(SCREAMING_SNAKE_CASE__    ):
																																new_dft[i].append(dft[i][j] - current_root * dft[i + next_ncol][j]    )
																									current_root *= root
																		# Update
																		SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   new_dft
																		SCREAMING_SNAKE_CASE__						:       Tuple           =   next_ncol // 2
											return dft[0]



				def 	__magic_name__    (self    )					->		int:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   self.__dft("""A"""    )
											SCREAMING_SNAKE_CASE__						:       Dict           =   self.__dft("""B"""    )
											SCREAMING_SNAKE_CASE__						:       List[Any]           =   [[dft_a[i] * dft_b[i] for i in range(self.c_max_length    )]]
											del dft_a
											del dft_b

											# Corner Case
											if len(inverce_c[0]    ) <= 1:
																		return inverce_c[0]
											# Inverse DFT
											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   2
											while next_ncol <= self.c_max_length:
																		SCREAMING_SNAKE_CASE__						:       List[str]           =   [[] for i in range(SCREAMING_SNAKE_CASE__    )]
																		SCREAMING_SNAKE_CASE__						:       Tuple           =   self.root ** (next_ncol // 2)
																		SCREAMING_SNAKE_CASE__						:       Any           =   1
																		# First half of next step
																		for j in range(self.c_max_length // next_ncol    ):
																									for i in range(next_ncol // 2    ):
																																# Even positions
																																new_inverse_c[i].append(
																																    (
																																        inverce_c[i][j]
																																        + inverce_c[i][j + self.c_max_length // next_ncol]
																																    )
																																    / 2    )
																																# Odd positions
																																new_inverse_c[i + next_ncol // 2].append(
																																    (
																																        inverce_c[i][j]
																																        - inverce_c[i][j + self.c_max_length // next_ncol]
																																    )
																																    / (2 * current_root)    )
																									current_root *= root
																		# Update
																		SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   new_inverse_c
																		next_ncol *= 2
											# Unpack
											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   [round(x[0].real   ,       8    ) + round(x[0].imag   ,       8    ) * 1j for x in inverce_c]

											# Remove leading 0's
											while inverce_c[-1] == 0:
																		inverce_c.pop()
											return inverce_c



				def __str__(self    )					->		List[str]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:       List[str]           =   """A = """ + """ + """.join(
											    F'''{coef}*x^{i}''' for coef, i in enumerate(self.polyA[: self.len_A]    )    )
											SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   """B = """ + """ + """.join(
											    F'''{coef}*x^{i}''' for coef, i in enumerate(self.polyB[: self.len_B]    )    )
											SCREAMING_SNAKE_CASE__						:       int           =   """A*B = """ + """ + """.join(
											    F'''{coef}*x^{i}''' for coef, i in enumerate(self.product    )    )

											return F'''{a}\n{b}\n{c}'''


# Unit tests
if __name__ == "__main__":
					import doctest

					doctest.testmod()




