'''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       ( _SCREAMING_SNAKE_CASE						,       unittest.TestCase		):
  __A :				Optional[Any]											= DiTPipeline
  __A :				List[Any]											= CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
  __A :				str											= PipelineTesterMixin.required_optional_params - {
      """latents""",
      """num_images_per_prompt""",
      """callback""",
      """callback_steps""",
  }
  __A :				Optional[Any]											= CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
  __A :				Any											= False





  def 						UpperCAmelCase_	(	self ):
     torch.manual_seed(0 )
     lowerCAmelCase_			=      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=__snake_case    ,     activation_fn='''gelu-approximate'''    ,     num_embeds_ada_norm=1000    ,     norm_type='''ada_norm_zero'''    ,     norm_elementwise_affine=__snake_case    ,     )
     lowerCAmelCase_			=      AutoencoderKL()
     lowerCAmelCase_			=      DDIMScheduler()
     lowerCAmelCase_			=      {'''transformer''': transformer.eval(), '''vae''': vae.eval(), '''scheduler''': scheduler}
     return components





  def 						UpperCAmelCase_	(	self    ,     _lowerCamelCase    ,     _lowerCamelCase=0 ):





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





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      '''cpu'''

     lowerCAmelCase_			=      self.get_dummy_components()
     lowerCAmelCase_			=      self.pipeline_class(**__snake_case )
     pipe.to(__snake_case )
     pipe.set_progress_bar_config(disable=__snake_case )

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

     self.assertEqual(image.shape    ,     (1, 16, 16, 3) )
     lowerCAmelCase_			=      np.array([0.29_46, 0.66_01, 0.43_29, 0.32_96, 0.41_44, 0.53_19, 0.72_73, 0.50_13, 0.44_57] )
     lowerCAmelCase_			=      np.abs(image_slice.flatten() - expected_slice ).max()
     self.assertLessEqual(__snake_case    ,     1E-3 )





  def 						UpperCAmelCase_	(	self ):
     self._test_inference_batch_single_identical(relax_max_difference=__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 						UpperCAmelCase_	(	self ):
     self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3 )






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





  def 						UpperCAmelCase_	(	self ):
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      torch.manual_seed(0 )

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

     lowerCAmelCase_			=      ['''vase''', '''umbrella''', '''white shark''', '''white wolf''']
     lowerCAmelCase_			=      pipe.get_label_ids(__snake_case )

     lowerCAmelCase_			=      pipe(__snake_case    ,     generator=__snake_case    ,     num_inference_steps=40    ,     output_type='''np''' ).images

     for word, image in zip(__snake_case    ,     __snake_case ):
        lowerCAmelCase_			=      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 ):
     lowerCAmelCase_			=      DiTPipeline.from_pretrained('''facebook/DiT-XL-2-512''' )
     lowerCAmelCase_			=      DPMSolverMultistepScheduler.from_config(pipe.scheduler.config )
     pipe.to('''cuda''' )

     lowerCAmelCase_			=      ['''vase''', '''umbrella''']
     lowerCAmelCase_			=      pipe.get_label_ids(__snake_case )

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

     for word, image in zip(__snake_case    ,     __snake_case ):
        lowerCAmelCase_			=      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












import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest





class UpperCAmelCase_				( _SCREAMING_SNAKE_CASE ):
  __lowerCAmelCase   :						int							=   (KDPMaDiscreteScheduler,)
  __lowerCAmelCase   :						Union[str, Any]							=   10


  def 						__UpperCAmelCase     (							self ,**__snake_case ):



       """simple docstring"""





       A_          =       {
           '''num_train_timesteps''': 1_1_0_0,
           '''beta_start''': 0.0001,
           '''beta_end''': 0.02,
           '''beta_schedule''': '''linear''',
       }

       config.update(**__snake_case )
       return config


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       for timesteps in [1_0, 5_0, 1_0_0, 1_0_0_0]:
            self.check_over_configs(num_train_timesteps=__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       for beta_start, beta_end in zip([0.0_0001, 0.0001, 0.001] ,[0.0002, 0.002, 0.02] ):
            self.check_over_configs(beta_start=__snake_case ,beta_end=__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       for schedule in ["linear", "scaled_linear"]:
            self.check_over_configs(beta_schedule=__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       for prediction_type in ["epsilon", "v_prediction"]:
            self.check_over_configs(prediction_type=__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       self.scheduler_classes[0]
       A_          =       self.get_scheduler_config(prediction_type='''v_prediction''' )
       A_          =       scheduler_class(**__snake_case )

       scheduler.set_timesteps(self.num_inference_steps )

       A_          =       self.dummy_model()
       A_          =       self.dummy_sample_deter * scheduler.init_noise_sigma
       A_          =       sample.to(__snake_case )

       for i, t in enumerate(scheduler.timesteps ):
            A_          =       scheduler.scale_model_input(__snake_case ,__snake_case )

            A_          =       model(__snake_case ,__snake_case )

            A_          =       scheduler.step(__snake_case ,__snake_case ,__snake_case )
            A_          =       output.prev_sample

       A_          =       torch.sum(torch.abs(__snake_case ) )
       A_          =       torch.mean(torch.abs(__snake_case ) )

       if torch_device in ["cpu", "mps"]:
            assert abs(result_sum.item() - 4.69_34E-07 ) < 1E-2
            assert abs(result_mean.item() - 6.11_12E-10 ) < 1E-3
       else:
            # CUDA
            assert abs(result_sum.item() - 4.6_93_42_86_50_17_09_72E-07 ) < 1E-2
            assert abs(result_mean.item() - 0.0002 ) < 1E-3


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       if torch_device == "mps":
            return
       A_          =       self.scheduler_classes[0]
       A_          =       self.get_scheduler_config()
       A_          =       scheduler_class(**__snake_case )

       scheduler.set_timesteps(self.num_inference_steps )

       A_          =       self.dummy_model()
       A_          =       self.dummy_sample_deter * scheduler.init_noise_sigma
       A_          =       sample.to(__snake_case )

       for i, t in enumerate(scheduler.timesteps ):
            A_          =       scheduler.scale_model_input(__snake_case ,__snake_case )

            A_          =       model(__snake_case ,__snake_case )

            A_          =       scheduler.step(__snake_case ,__snake_case ,__snake_case )
            A_          =       output.prev_sample

       A_          =       torch.sum(torch.abs(__snake_case ) )
       A_          =       torch.mean(torch.abs(__snake_case ) )

       if torch_device in ["cpu", "mps"]:
            assert abs(result_sum.item() - 20.4125 ) < 1E-2
            assert abs(result_mean.item() - 0.0266 ) < 1E-3
       else:
            # CUDA
            assert abs(result_sum.item() - 20.4125 ) < 1E-2
            assert abs(result_mean.item() - 0.0266 ) < 1E-3




  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       if torch_device == "mps":
            return
       A_          =       self.scheduler_classes[0]
       A_          =       self.get_scheduler_config()
       A_          =       scheduler_class(**__snake_case )

       scheduler.set_timesteps(self.num_inference_steps ,device=__snake_case )

       A_          =       self.dummy_model()
       A_          =       self.dummy_sample_deter.to(__snake_case ) * scheduler.init_noise_sigma

       for t in scheduler.timesteps:
            A_          =       scheduler.scale_model_input(__snake_case ,__snake_case )

            A_          =       model(__snake_case ,__snake_case )

            A_          =       scheduler.step(__snake_case ,__snake_case ,__snake_case )
            A_          =       output.prev_sample

       A_          =       torch.sum(torch.abs(__snake_case ) )
       A_          =       torch.mean(torch.abs(__snake_case ) )

       if str(__snake_case ).startswith('''cpu''' ):
            # The following sum varies between 148 and 156 on mps. Why?
            assert abs(result_sum.item() - 20.4125 ) < 1E-2
            assert abs(result_mean.item() - 0.0266 ) < 1E-3
       else:
            # CUDA
            assert abs(result_sum.item() - 20.4125 ) < 1E-2
            assert abs(result_mean.item() - 0.0266 ) < 1E-3



