


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

  def    _UpperCamelCase( self			:  Union[str, Any]				,				lowerCamelCase__			:  List[str]				,				lowerCamelCase__			:  List[str]=0       ):
        if str(lowerCamelCase__       ).startswith("mps"       ):
              a__	:  Any     =   torch.manual_seed(lowerCamelCase__       )
        else:
              a__	:  Optional[int]     =   torch.Generator(device=lowerCamelCase__       ).manual_seed(lowerCamelCase__       )
        a__	:  int     =   {
            "class_labels": [1],
            "generator": generator,
            "num_inference_steps": 2,
            "output_type": "numpy",
        }
        return inputs

  def    _UpperCamelCase( self			:  List[str]       ):
        a__	:  int     =   "cpu"

        a__	:  Optional[int]     =   self.get_dummy_components()
        a__	:  List[Any]     =   self.pipeline_class(**lowerCamelCase__       )
        pipe.to(lowerCamelCase__       )
        pipe.set_progress_bar_config(disable=lowerCamelCase__       )

        a__	:  str     =   self.get_dummy_inputs(lowerCamelCase__       )
        a__	:  Optional[int]     =   pipe(**lowerCamelCase__       ).images
        a__	:  Optional[Any]     =   image[0, -3:, -3:, -1]

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

  def    _UpperCamelCase( self			:  str       ):
        self._test_inference_batch_single_identical(relax_max_difference=lowerCamelCase__				,				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			:  Tuple       ):
        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3       )




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



  """simple docstring"""



  def    _UpperCamelCase( self			:  Optional[int]       ):
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()

  def    _UpperCamelCase( self			:  int       ):
        a__	:  int     =   torch.manual_seed(0       )

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

        a__	:  Any     =   ["vase", "umbrella", "white shark", "white wolf"]
        a__	:  Tuple     =   pipe.get_label_ids(lowerCamelCase__       )

        a__	:  int     =   pipe(lowerCamelCase__				,				generator=lowerCamelCase__				,				num_inference_steps=40				,				output_type="np"       ).images

        for word, image in zip(lowerCamelCase__				,				lowerCamelCase__       ):
              a__	:  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    _UpperCamelCase( self			:  Dict       ):
        a__	:  List[Any]     =   DiTPipeline.from_pretrained("facebook/DiT-XL-2-512"       )
        a__	:  Union[str, Any]     =   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config       )
        pipe.to("cuda"       )

        a__	:  List[str]     =   ["vase", "umbrella"]
        a__	:  List[str]     =   pipe.get_label_ids(lowerCamelCase__       )

        a__	:  Optional[Any]     =   torch.manual_seed(0       )
        a__	:  Union[str, Any]     =   pipe(lowerCamelCase__				,				generator=lowerCamelCase__				,				num_inference_steps=25				,				output_type="np"       ).images

        for word, image in zip(lowerCamelCase__				,				lowerCamelCase__       ):
              a__	:  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








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        __magic_name__	(				unittest.TestCase				):





     """simple docstring"""



     def     SCREAMING_SNAKE_CASE		(				self							:str	):





        '''simple docstring'''




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



     def     SCREAMING_SNAKE_CASE		(				self							:int	):





        '''simple docstring'''




        A_   ,    A_       :					Optional[Any]        =      FlaxStableDiffusionPipeline.from_pretrained(
            "stabilityai/stable-diffusion-2"      ,   revision="bf16"      ,   dtype=jnp.bfloataa      ,   )

        A_       :					List[Any]        =      "A painting of a squirrel eating a burger"
        A_       :					int        =      jax.device_count()
        A_       :					Union[str, Any]        =      num_samples * [prompt]
        A_       :					Dict        =      sd_pipe.prepare_inputs(snake_case	)

        A_       :					List[Any]        =      replicate(snake_case	)
        A_       :					str        =      shard(snake_case	)

        A_       :					Any        =      jax.random.PRNGKey(0	)
        A_       :					Dict        =      jax.random.split(snake_case      ,   jax.device_count()	)

        A_       :					int        =      sd_pipe(snake_case      ,   snake_case      ,   snake_case      ,   num_inference_steps=25      ,   jit=snake_case	)[0]
        assert images.shape == (jax.device_count(), 1, 768, 768, 3)

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

        A_       :					Dict        =      jnp.asarray(jax.device_get(image_slice.flatten()	)	)
        A_       :					List[str]        =      jnp.array([0.4238, 0.4414, 0.4395, 0.4453, 0.4629, 0.4590, 0.4531, 0.45508, 0.4512]	)
        print(f"output_slice: {output_slice}"	)
        assert jnp.abs(output_slice - expected_slice	).max() < 1e-2





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





        '''simple docstring'''




        A_       :					str        =      "stabilityai/stable-diffusion-2"
        A_   ,    A_       :					Dict        =      FlaxDPMSolverMultistepScheduler.from_pretrained(snake_case      ,   subfolder="scheduler"	)
        A_   ,    A_       :					Any        =      FlaxStableDiffusionPipeline.from_pretrained(
            snake_case      ,   scheduler=snake_case      ,   revision="bf16"      ,   dtype=jnp.bfloataa      ,   )
        A_       :					Any        =      scheduler_params

        A_       :					Optional[Any]        =      "A painting of a squirrel eating a burger"
        A_       :					Dict        =      jax.device_count()
        A_       :					List[str]        =      num_samples * [prompt]
        A_       :					int        =      sd_pipe.prepare_inputs(snake_case	)

        A_       :					Optional[Any]        =      replicate(snake_case	)
        A_       :					Dict        =      shard(snake_case	)

        A_       :					Any        =      jax.random.PRNGKey(0	)
        A_       :					Union[str, Any]        =      jax.random.split(snake_case      ,   jax.device_count()	)

        A_       :					Union[str, Any]        =      sd_pipe(snake_case      ,   snake_case      ,   snake_case      ,   num_inference_steps=25      ,   jit=snake_case	)[0]
        assert images.shape == (jax.device_count(), 1, 768, 768, 3)

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

        A_       :					Tuple        =      jnp.asarray(jax.device_get(image_slice.flatten()	)	)
        A_       :					str        =      jnp.array([0.4336, 0.42969, 0.4453, 0.4199, 0.4297, 0.4531, 0.4434, 0.4434, 0.4297]	)
        print(f"output_slice: {output_slice}"	)
        assert jnp.abs(output_slice - expected_slice	).max() < 1e-2
