




"""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 					lowercase(     __a       ,			unittest.TestCase ):







 '''simple docstring'''





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


 def  UpperCamelCase_   (				self:		Any		):




   '''simple docstring'''



   torch.manual_seed(0		)
   _snake_case							:		Tuple								=  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=a_,      activation_fn="""gelu-approximate""",      num_embeds_ada_norm=1_000,      norm_type="""ada_norm_zero""",      norm_elementwise_affine=a_,      )
   _snake_case							:		List[Any]								=  AutoencoderKL()
   _snake_case							:		Any								=  DDIMScheduler()
   _snake_case							:		Optional[Any]								=  {"""transformer""": transformer.eval(), """vae""": vae.eval(), """scheduler""": scheduler}
   return components


 def  UpperCamelCase_   (				self:		Any,      a_:		Dict,      a_:		List[Any]=0		):




   '''simple docstring'''



   if str(a_		).startswith("""mps"""		):
     _snake_case							:		str								=  torch.manual_seed(a_		)
   else:
     _snake_case							:		Dict								=  torch.Generator(device=a_		).manual_seed(a_		)
   _snake_case							:		Union[str, Any]								=  {
       """class_labels""": [1],
       """generator""": generator,
       """num_inference_steps""": 2,
       """output_type""": """numpy""",
   }
   return inputs


 def  UpperCamelCase_   (				self:		Dict		):




   '''simple docstring'''



   _snake_case							:		Tuple								=  """cpu"""

   _snake_case							:		str								=  self.get_dummy_components()
   _snake_case							:		Union[str, Any]								=  self.pipeline_class(**a_		)
   pipe.to(a_		)
   pipe.set_progress_bar_config(disable=a_		)

   _snake_case							:		int								=  self.get_dummy_inputs(a_		)
   _snake_case							:		Any								=  pipe(**a_		).images
   _snake_case							:		Tuple								=  image[0, -3:, -3:, -1]

   self.assertEqual(image.shape,      (1, 16, 16, 3)		)
   _snake_case							:		List[str]								=  np.array([0.2_946, 0.6_601, 0.4_329, 0.3_296, 0.4_144, 0.5_319, 0.7_273, 0.5_013, 0.4_457]		)
   _snake_case							:		Optional[Any]								=  np.abs(image_slice.flatten() - expected_slice		).max()
   self.assertLessEqual(a_,      1E-3		)


 def  UpperCamelCase_   (				self:		Dict		):




   '''simple docstring'''



   self._test_inference_batch_single_identical(relax_max_difference=a_,      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  UpperCamelCase_   (				self:		Any		):




   '''simple docstring'''



   self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3		)







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







 '''simple docstring'''





 def  UpperCamelCase_   (				self:		Optional[int]		):




   '''simple docstring'''



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


 def  UpperCamelCase_   (				self:		str		):




   '''simple docstring'''



   _snake_case							:		List[Any]								=  torch.manual_seed(0		)

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

   _snake_case							:		Optional[Any]								=  ["""vase""", """umbrella""", """white shark""", """white wolf"""]
   _snake_case							:		List[str]								=  pipe.get_label_ids(a_		)

   _snake_case							:		List[str]								=  pipe(a_,      generator=a_,      num_inference_steps=40,      output_type="""np"""		).images

   for word, image in zip(a_,      a_		):
     _snake_case							:		Dict								=  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  UpperCamelCase_   (				self:		Dict		):




   '''simple docstring'''



   _snake_case							:		Dict								=  DiTPipeline.from_pretrained("""facebook/DiT-XL-2-512"""		)
   _snake_case							:		List[Any]								=  DPMSolverMultistepScheduler.from_config(pipe.scheduler.config		)
   pipe.to("""cuda"""		)

   _snake_case							:		Dict								=  ["""vase""", """umbrella"""]
   _snake_case							:		List[str]								=  pipe.get_label_ids(a_		)

   _snake_case							:		Tuple								=  torch.manual_seed(0		)
   _snake_case							:		Dict								=  pipe(a_,      generator=a_,      num_inference_steps=25,      output_type="""np"""		).images

   for word, image in zip(a_,      a_		):
     _snake_case							:		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_tf_available, is_torch_available


_snake_case          =   {
    'configuration_transfo_xl': ['TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP', 'TransfoXLConfig'],
    'tokenization_transfo_xl': ['TransfoXLCorpus', 'TransfoXLTokenizer'],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _snake_case          =   [
     'TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST',
     'AdaptiveEmbedding',
     'TransfoXLForSequenceClassification',
     'TransfoXLLMHeadModel',
     'TransfoXLModel',
     'TransfoXLPreTrainedModel',
     'load_tf_weights_in_transfo_xl',
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _snake_case          =   [
     'TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST',
     'TFAdaptiveEmbedding',
     'TFTransfoXLForSequenceClassification',
     'TFTransfoXLLMHeadModel',
     'TFTransfoXLMainLayer',
     'TFTransfoXLModel',
     'TFTransfoXLPreTrainedModel',
 ]


if TYPE_CHECKING:
 from .configuration_transfo_xl import TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP, TransfoXLConfig
 from .tokenization_transfo_xl import TransfoXLCorpus, TransfoXLTokenizer

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_transfo_xl import (
      TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
      AdaptiveEmbedding,
      TransfoXLForSequenceClassification,
      TransfoXLLMHeadModel,
      TransfoXLModel,
      TransfoXLPreTrainedModel,
      load_tf_weights_in_transfo_xl,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_transfo_xl import (
      TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFAdaptiveEmbedding,
      TFTransfoXLForSequenceClassification,
      TFTransfoXLLMHeadModel,
      TFTransfoXLMainLayer,
      TFTransfoXLModel,
      TFTransfoXLPreTrainedModel,
  )

else:
 import sys

 _snake_case          =   _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)

