


'''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    UpperCAmelCase_  (_UpperCAmelCase					,		unittest.TestCase  ):





    """simple docstring"""






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







    def       lowercase_				(							self	)				->     Any:
         torch.manual_seed(0	)
         __lowerCamelCase   :	Optional[int]								=		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_			,     )
         __lowerCamelCase   :	List[str]								=		AutoencoderKL()
         __lowerCamelCase   :	Any								=		DDIMScheduler()
         __lowerCamelCase   :	Union[str, Any]								=		{'transformer': transformer.eval(), 'vae': vae.eval(), 'scheduler': scheduler}
         return components







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=0	)				->     List[str]:
         if str(SCREAMING_SNAKE_CASE_	).startswith('mps'	):
              __lowerCamelCase   :	Any								=		torch.manual_seed(SCREAMING_SNAKE_CASE_	)
         else:
              __lowerCamelCase   :	List[str]								=		torch.Generator(device=SCREAMING_SNAKE_CASE_	).manual_seed(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Tuple								=		{
             'class_labels': [1],
             'generator': generator,
             'num_inference_steps': 2,
             'output_type': 'numpy',
         }
         return inputs







    def       lowercase_				(							self	)				->     Tuple:
         __lowerCamelCase   :	List[str]								=		'cpu'

         __lowerCamelCase   :	List[Any]								=		self.get_dummy_components()
         __lowerCamelCase   :	str								=		self.pipeline_class(**SCREAMING_SNAKE_CASE_	)
         pipe.to(SCREAMING_SNAKE_CASE_	)
         pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)

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

         self.assertEqual(image.shape			,     (1, 16, 16, 3)	)
         __lowerCamelCase   :	Optional[int]								=		np.array([0.2_9_4_6, 0.6_6_0_1, 0.4_3_2_9, 0.3_2_9_6, 0.4_1_4_4, 0.5_3_1_9, 0.7_2_7_3, 0.5_0_1_3, 0.4_4_5_7]	)
         __lowerCamelCase   :	List[str]								=		np.abs(image_slice.flatten() - expected_slice	).max()
         self.assertLessEqual(SCREAMING_SNAKE_CASE_			,     1E-3	)







    def       lowercase_				(							self	)				->     Union[str, Any]:
         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       lowercase_				(							self	)				->     Tuple:
         self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3	)







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





    """simple docstring"""






    def       lowercase_				(							self	)				->     Dict:
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()







    def       lowercase_				(							self	)				->     List[Any]:
         __lowerCamelCase   :	str								=		torch.manual_seed(0	)

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

         __lowerCamelCase   :	Dict								=		['vase', 'umbrella', 'white shark', 'white wolf']
         __lowerCamelCase   :	int								=		pipe.get_label_ids(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	Optional[Any]								=		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_	):
              __lowerCamelCase   :	List[str]								=		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       lowercase_				(							self	)				->     List[Any]:
         __lowerCamelCase   :	str								=		DiTPipeline.from_pretrained('facebook/DiT-XL-2-512'	)
         __lowerCamelCase   :	Optional[int]								=		DPMSolverMultistepScheduler.from_config(pipe.scheduler.config	)
         pipe.to('cuda'	)

         __lowerCamelCase   :	Any								=		['vase', 'umbrella']
         __lowerCamelCase   :	Tuple								=		pipe.get_label_ids(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	Dict								=		torch.manual_seed(0	)
         __lowerCamelCase   :	Optional[int]								=		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_	):
              __lowerCamelCase   :	Optional[int]								=		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'''



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


A__		:					List[Any]												=  {
    """configuration_megatron_bert""": ["""MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MegatronBertConfig"""],
}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      A__		:					int												=  [
          """MEGATRON_BERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """MegatronBertForCausalLM""",
          """MegatronBertForMaskedLM""",
          """MegatronBertForMultipleChoice""",
          """MegatronBertForNextSentencePrediction""",
          """MegatronBertForPreTraining""",
          """MegatronBertForQuestionAnswering""",
          """MegatronBertForSequenceClassification""",
          """MegatronBertForTokenClassification""",
          """MegatronBertModel""",
          """MegatronBertPreTrainedModel""",
      ]

if TYPE_CHECKING:
      from .configuration_megatron_bert import MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, MegatronBertConfig

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_megatron_bert import (
                MEGATRON_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                MegatronBertForCausalLM,
                MegatronBertForMaskedLM,
                MegatronBertForMultipleChoice,
                MegatronBertForNextSentencePrediction,
                MegatronBertForPreTraining,
                MegatronBertForQuestionAnswering,
                MegatronBertForSequenceClassification,
                MegatronBertForTokenClassification,
                MegatronBertModel,
                MegatronBertPreTrainedModel,
            )

else:
      import sys

      A__		:					List[Any]												=  _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)






