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





def    _snake_case  (  _snake_case :	int	):

      if not isinstance(_snake_case       ,      _snake_case	):
            raise TypeError('''only integers accepted as input'''	)
      else:
            lowerCAmelCase    :		List[str]				=      str(abs(_snake_case	)	)
            lowerCAmelCase    :		Optional[Any]				=      [list(_snake_case	) for char in range(len(_snake_case	)	)]
            for index in range(len(_snake_case	)	):
                  num_transpositions[index].pop(_snake_case	)
            return max(
                int(''''''.join(list(_snake_case	)	)	) for transposition in num_transpositions	)


if __name__ == "__main__":
 __import__('''doctest''').testmod()


