




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 							UpperCamelCase_       (					UpperCAmelCase__						,   unittest.TestCase	):

    '''simple docstring'''






    UpperCAmelCase__				    = DiTPipeline
    UpperCAmelCase__				    = CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
    UpperCAmelCase__				    = PipelineTesterMixin.required_optional_params - {
        '''latents''',
        '''num_images_per_prompt''',
        '''callback''',
        '''callback_steps''',
    }
    UpperCAmelCase__				    = CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
    UpperCAmelCase__				    = False


    def 			SCREAMING_SNAKE_CASE				(     self :      int)		->List[Any]:





           '''simple docstring'''

           torch.manual_seed(0)
           A__ 					=       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=UpperCAmelCase__   ,					activation_fn='''gelu-approximate'''   ,					num_embeds_ada_norm=1_000   ,					norm_type='''ada_norm_zero'''   ,					norm_elementwise_affine=UpperCAmelCase__   ,					)
           A__ 					=       AutoencoderKL()
           A__ 					=       DDIMScheduler()
           A__ 					=       {'''transformer''': transformer.eval(), '''vae''': vae.eval(), '''scheduler''': scheduler}
           return components


    def 			SCREAMING_SNAKE_CASE				(     self :      Any   ,					UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      List[Any]=0)		->str:





           '''simple docstring'''

           if str(UpperCAmelCase__).startswith('''mps'''):
                  A__ 					=       torch.manual_seed(UpperCAmelCase__)
           else:
                  A__ 					=       torch.Generator(device=UpperCAmelCase__).manual_seed(UpperCAmelCase__)
           A__ 					=       {
               '''class_labels''': [1],
               '''generator''': generator,
               '''num_inference_steps''': 2,
               '''output_type''': '''numpy''',
           }
           return inputs


    def 			SCREAMING_SNAKE_CASE				(     self :      List[Any])		->Optional[int]:





           '''simple docstring'''

           A__ 					=       '''cpu'''

           A__ 					=       self.get_dummy_components()
           A__ 					=       self.pipeline_class(**UpperCAmelCase__)
           pipe.to(UpperCAmelCase__)
           pipe.set_progress_bar_config(disable=UpperCAmelCase__)

           A__ 					=       self.get_dummy_inputs(UpperCAmelCase__)
           A__ 					=       pipe(**UpperCAmelCase__).images
           A__ 					=       image[0, -3:, -3:, -1]

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


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





           '''simple docstring'''

           self._test_inference_batch_single_identical(relax_max_difference=UpperCAmelCase__   ,					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 			SCREAMING_SNAKE_CASE				(     self :      Any)		->Union[str, Any]:





           '''simple docstring'''

           self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3)


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

    '''simple docstring'''






    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[Any])		->str:





           '''simple docstring'''

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


    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->Optional[Any]:





           '''simple docstring'''

           A__ 					=       torch.manual_seed(0)

           A__ 					=       DiTPipeline.from_pretrained('''facebook/DiT-XL-2-256''')
           pipe.to('''cuda''')

           A__ 					=       ['''vase''', '''umbrella''', '''white shark''', '''white wolf''']
           A__ 					=       pipe.get_label_ids(UpperCAmelCase__)

           A__ 					=       pipe(UpperCAmelCase__   ,					generator=UpperCAmelCase__   ,					num_inference_steps=40   ,					output_type='''np''').images

           for word, image in zip(UpperCAmelCase__   ,					UpperCAmelCase__):
                  A__ 					=       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 			SCREAMING_SNAKE_CASE				(     self :      Optional[int])		->List[str]:





           '''simple docstring'''

           A__ 					=       DiTPipeline.from_pretrained('''facebook/DiT-XL-2-512''')
           A__ 					=       DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
           pipe.to('''cuda''')

           A__ 					=       ['''vase''', '''umbrella''']
           A__ 					=       pipe.get_label_ids(UpperCAmelCase__)

           A__ 					=       torch.manual_seed(0)
           A__ 					=       pipe(UpperCAmelCase__   ,					generator=UpperCAmelCase__   ,					num_inference_steps=25   ,					output_type='''np''').images

           for word, image in zip(UpperCAmelCase__   ,					UpperCAmelCase__):
                  A__ 					=       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'''






# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from ...utils.dataclasses import (
    ComputeEnvironment,
    DistributedType,
    DynamoBackend,
    PrecisionType,
    SageMakerDistributedType,
)
from ..menu import BulletMenu


__UpperCAmelCase       							=				[
    '''EAGER''',
    '''AOT_EAGER''',
    '''INDUCTOR''',
    '''NVFUSER''',
    '''AOT_NVFUSER''',
    '''AOT_CUDAGRAPHS''',
    '''OFI''',
    '''FX2TRT''',
    '''ONNXRT''',
    '''IPEX''',
]


def    _snake_case  (	A					,       A=None					,       A=None					,       A=None			)		->       Union[str, Any]:
							lowerCAmelCase__			  =      True
							while ask_again:
														lowerCAmelCase__			  =      input(A			)
														try:
																					if default is not None and len(A			) == 0:
																												return default
																					return convert_value(A			) if convert_value is not None else result
														except Exception:
																					if error_message is not None:
																												print(A			)


def    _snake_case  (	A					,       A=[]					,       A=None					,       A=0			)		->       List[Any]:
							lowerCAmelCase__			  =      BulletMenu(A					,       A			)
							lowerCAmelCase__			  =      menu.run(default_choice=A			)
							return convert_value(A			) if convert_value is not None else result


def    _snake_case  (	A			)		->       Tuple:
							lowerCAmelCase__			  =      int(A			)
							return ComputeEnvironment(['''LOCAL_MACHINE''', '''AMAZON_SAGEMAKER'''][value]			)


def    _snake_case  (	A			)		->       Union[str, Any]:
							lowerCAmelCase__			  =      int(A			)
							return DistributedType(['''NO''', '''MULTI_CPU''', '''MULTI_XPU''', '''MULTI_GPU''', '''MULTI_NPU''', '''TPU'''][value]			)


def    _snake_case  (	A			)		->       str:
							lowerCAmelCase__			  =      int(A			)
							return DynamoBackend(DYNAMO_BACKENDS[value]			).value


def    _snake_case  (	A			)		->       Tuple:
							lowerCAmelCase__			  =      int(A			)
							return PrecisionType(['''no''', '''fp16''', '''bf16''', '''fp8'''][value]			)


def    _snake_case  (	A			)		->       Union[str, Any]:
							lowerCAmelCase__			  =      int(A			)
							return SageMakerDistributedType(['''NO''', '''DATA_PARALLEL''', '''MODEL_PARALLEL'''][value]			)


def    _snake_case  (	A			)		->       List[str]:
							return {"yes": True, "no": False}[value.lower()]






class 						a__							(				argparse.RawDescriptionHelpFormatter    ):





							'''simple docstring'''






							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )			->				Dict:
														lowerCAmelCase__			  =      super()._format_usage(lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )
														lowerCAmelCase__			  =      usage.replace('''<command> [<args>] '''   ,       ''''''  )
														return usage