





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

      '''simple docstring'''



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

      def       A_					(							self       :      Dict ):


        '''simple docstring'''




        torch.manual_seed(0 )
        __UpperCAmelCase    :						Optional[Any]       =						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=__lowercase						,				activation_fn='''gelu-approximate'''						,				num_embeds_ada_norm=1_000						,				norm_type='''ada_norm_zero'''						,				norm_elementwise_affine=__lowercase						,				)
        __UpperCAmelCase    :						Union[str, Any]       =						AutoencoderKL()
        __UpperCAmelCase    :						Dict       =						DDIMScheduler()
        __UpperCAmelCase    :						int       =						{'''transformer''': transformer.eval(), '''vae''': vae.eval(), '''scheduler''': scheduler}
        return components

      def       A_					(							self       :      Optional[Any]						,				__lowercase       :      Union[str, Any]						,				__lowercase       :      Any=0 ):


        '''simple docstring'''




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

      def       A_					(							self       :      str ):


        '''simple docstring'''




        __UpperCAmelCase    :						Tuple       =						'''cpu'''

        __UpperCAmelCase    :						Optional[int]       =						self.get_dummy_components()
        __UpperCAmelCase    :						Dict       =						self.pipeline_class(**__lowercase )
        pipe.to(__lowercase )
        pipe.set_progress_bar_config(disable=__lowercase )

        __UpperCAmelCase    :						Tuple       =						self.get_dummy_inputs(__lowercase )
        __UpperCAmelCase    :						List[Any]       =						pipe(**__lowercase ).images
        __UpperCAmelCase    :						List[Any]       =						image[0, -3:, -3:, -1]

        self.assertEqual(image.shape						,				(1, 16, 16, 3) )
        __UpperCAmelCase    :						Any       =						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] )
        __UpperCAmelCase    :						Tuple       =						np.abs(image_slice.flatten() - expected_slice ).max()
        self.assertLessEqual(__lowercase						,				1e-3 )

      def       A_					(							self       :      List[str] ):


        '''simple docstring'''




        self._test_inference_batch_single_identical(relax_max_difference=__lowercase						,				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       A_					(							self       :      Tuple ):


        '''simple docstring'''




        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3 )




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

      '''simple docstring'''



      def       A_					(							self       :      Optional[int] ):


        '''simple docstring'''




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

      def       A_					(							self       :      Tuple ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[int]       =						torch.manual_seed(0 )

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

        __UpperCAmelCase    :						Optional[Any]       =						['''vase''', '''umbrella''', '''white shark''', '''white wolf''']
        __UpperCAmelCase    :						Union[str, Any]       =						pipe.get_label_ids(__lowercase )

        __UpperCAmelCase    :						Optional[Any]       =						pipe(__lowercase						,				generator=__lowercase						,				num_inference_steps=40						,				output_type='''np''' ).images

        for word, image in zip(__lowercase						,				__lowercase ):
          __UpperCAmelCase    :						Union[str, Any]       =						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       A_					(							self       :      Dict ):


        '''simple docstring'''




        __UpperCAmelCase    :						List[str]       =						DiTPipeline.from_pretrained('''facebook/DiT-XL-2-512''' )
        __UpperCAmelCase    :						Union[str, Any]       =						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config )
        pipe.to('''cuda''' )

        __UpperCAmelCase    :						List[str]       =						['''vase''', '''umbrella''']
        __UpperCAmelCase    :						Dict       =						pipe.get_label_ids(__lowercase )

        __UpperCAmelCase    :						Optional[int]       =						torch.manual_seed(0 )
        __UpperCAmelCase    :						List[Any]       =						pipe(__lowercase						,				generator=__lowercase						,				num_inference_steps=25						,				output_type='''np''' ).images

        for word, image in zip(__lowercase						,				__lowercase ):
          __UpperCAmelCase    :						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"""





import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
  import jax
  import jax.numpy as jnp
  from flax.jax_utils import replicate
  from flax.training.common_utils import shard




@slow
@require_flax
class 							snake_case							(			unittest.TestCase ):

      '''simple docstring'''



      def       A_					(							self       :      str ):


        '''simple docstring'''




        super().tearDown()
        gc.collect()

      def       A_					(							self       :      List[Any] ):


        '''simple docstring'''




        __UpperCAmelCase       ,     __UpperCAmelCase    :						List[Any]       =						FlaxStableDiffusionPipeline.from_pretrained(
            '''stabilityai/stable-diffusion-2'''						,				revision='''bf16'''						,				dtype=jnp.bfloataa						,				)

        __UpperCAmelCase    :						Dict       =						'''A painting of a squirrel eating a burger'''
        __UpperCAmelCase    :						Any       =						jax.device_count()
        __UpperCAmelCase    :						Optional[int]       =						num_samples * [prompt]
        __UpperCAmelCase    :						Tuple       =						sd_pipe.prepare_inputs(__lowercase )

        __UpperCAmelCase    :						List[Any]       =						replicate(__lowercase )
        __UpperCAmelCase    :						Optional[Any]       =						shard(__lowercase )

        __UpperCAmelCase    :						str       =						jax.random.PRNGKey(0 )
        __UpperCAmelCase    :						Dict       =						jax.random.split(__lowercase						,				jax.device_count() )

        __UpperCAmelCase    :						Tuple       =						sd_pipe(__lowercase						,				__lowercase						,				__lowercase						,				num_inference_steps=25						,				jit=__lowercase )[0]
        assert images.shape == (jax.device_count(), 1, 768, 768, 3)

        __UpperCAmelCase    :						Any       =						images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:] )
        __UpperCAmelCase    :						Union[str, Any]       =						images[0, 253:256, 253:256, -1]

        __UpperCAmelCase    :						str       =						jnp.asarray(jax.device_get(image_slice.flatten() ) )
        __UpperCAmelCase    :						List[Any]       =						jnp.array([0.4_2_3_8, 0.4_4_1_4, 0.4_3_9_5, 0.4_4_5_3, 0.4_6_2_9, 0.4_5_9_0, 0.4_5_3_1, 0.4_5_5_0_8, 0.4_5_1_2] )
        print(f'''output_slice: {output_slice}''' )
        assert jnp.abs(output_slice - expected_slice ).max() < 1e-2







      def       A_					(							self       :      Union[str, Any] ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[int]       =						'''stabilityai/stable-diffusion-2'''
        __UpperCAmelCase       ,     __UpperCAmelCase    :						Any       =						FlaxDPMSolverMultistepScheduler.from_pretrained(__lowercase						,				subfolder='''scheduler''' )
        __UpperCAmelCase       ,     __UpperCAmelCase    :						Optional[int]       =						FlaxStableDiffusionPipeline.from_pretrained(
            __lowercase						,				scheduler=__lowercase						,				revision='''bf16'''						,				dtype=jnp.bfloataa						,				)
        __UpperCAmelCase    :						Dict       =						scheduler_params

        __UpperCAmelCase    :						int       =						'''A painting of a squirrel eating a burger'''
        __UpperCAmelCase    :						List[Any]       =						jax.device_count()
        __UpperCAmelCase    :						Optional[int]       =						num_samples * [prompt]
        __UpperCAmelCase    :						Tuple       =						sd_pipe.prepare_inputs(__lowercase )

        __UpperCAmelCase    :						Optional[int]       =						replicate(__lowercase )
        __UpperCAmelCase    :						List[Any]       =						shard(__lowercase )

        __UpperCAmelCase    :						Union[str, Any]       =						jax.random.PRNGKey(0 )
        __UpperCAmelCase    :						List[Any]       =						jax.random.split(__lowercase						,				jax.device_count() )

        __UpperCAmelCase    :						List[Any]       =						sd_pipe(__lowercase						,				__lowercase						,				__lowercase						,				num_inference_steps=25						,				jit=__lowercase )[0]
        assert images.shape == (jax.device_count(), 1, 768, 768, 3)

        __UpperCAmelCase    :						Tuple       =						images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:] )
        __UpperCAmelCase    :						Any       =						images[0, 253:256, 253:256, -1]

        __UpperCAmelCase    :						List[str]       =						jnp.asarray(jax.device_get(image_slice.flatten() ) )
        __UpperCAmelCase    :						Any       =						jnp.array([0.4_3_3_6, 0.4_2_9_6_9, 0.4_4_5_3, 0.4_1_9_9, 0.4_2_9_7, 0.4_5_3_1, 0.4_4_3_4, 0.4_4_3_4, 0.4_2_9_7] )
        print(f'''output_slice: {output_slice}''' )
        assert jnp.abs(output_slice - expected_slice ).max() < 1e-2