



import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





class 		UpperCAmelCase						(						__lowercase       ,__lowercase       ,unittest.TestCase				):
   SCREAMING_SNAKE_CASE__							=     StableDiffusionXLImgaImgPipeline
   SCREAMING_SNAKE_CASE__							=     TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
   SCREAMING_SNAKE_CASE__							=     PipelineTesterMixin.required_optional_params - {'''latents'''}
   SCREAMING_SNAKE_CASE__							=     TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
   SCREAMING_SNAKE_CASE__							=     IMAGE_TO_IMAGE_IMAGE_PARAMS
   SCREAMING_SNAKE_CASE__							=     IMAGE_TO_IMAGE_IMAGE_PARAMS





   def 	__lowerCAmelCase							(	self			):
         torch.manual_seed(0			)
         _lowerCAmelCase 			=  UNetaDConditionModel(
             block_out_channels=(32, 64)       ,     layers_per_block=2       ,     sample_size=32       ,     in_channels=4       ,     out_channels=4       ,     down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')       ,     up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')       ,     attention_head_dim=(2, 4)       ,     use_linear_projection=__a       ,     addition_embed_type='''text_time'''       ,     addition_time_embed_dim=8       ,     transformer_layers_per_block=(1, 2)       ,     projection_class_embeddings_input_dim=80       ,     cross_attention_dim=64       ,     )
         _lowerCAmelCase 			=  EulerDiscreteScheduler(
             beta_start=0.00_085       ,     beta_end=0.012       ,     steps_offset=1       ,     beta_schedule='''scaled_linear'''       ,     timestep_spacing='''leading'''       ,     )
         torch.manual_seed(0			)
         _lowerCAmelCase 			=  AutoencoderKL(
             block_out_channels=[32, 64]       ,     in_channels=3       ,     out_channels=3       ,     down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']       ,     up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']       ,     latent_channels=4       ,     sample_size=128       ,     )
         torch.manual_seed(0			)
         _lowerCAmelCase 			=  CLIPTextConfig(
             bos_token_id=0       ,     eos_token_id=2       ,     hidden_size=32       ,     intermediate_size=37       ,     layer_norm_eps=1E-05       ,     num_attention_heads=4       ,     num_hidden_layers=5       ,     pad_token_id=1       ,     vocab_size=1_000       ,     hidden_act='''gelu'''       ,     projection_dim=32       ,     )
         _lowerCAmelCase 			=  CLIPTextModel(__a			)
         _lowerCAmelCase 			=  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       ,     local_files_only=__a			)

         _lowerCAmelCase 			=  CLIPTextModelWithProjection(__a			)
         _lowerCAmelCase 			=  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       ,     local_files_only=__a			)

         _lowerCAmelCase 			=  {
             """unet""": unet,
             """scheduler""": scheduler,
             """vae""": vae,
             """text_encoder""": text_encoder,
             """tokenizer""": tokenizer,
             """text_encoder_2""": text_encoder_a,
             """tokenizer_2""": tokenizer_a,
             # "safety_checker": None,
             # "feature_extractor": None,
         }
         return components





   def 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase       ,     _lowerCAmelCase=0			):
         _lowerCAmelCase 			=  floats_tensor((1, 3, 32, 32)       ,     rng=random.Random(__a			)			).to(__a			)
         _lowerCAmelCase 			=  image / 2 + 0.5
         if str(__a			).startswith('''mps'''			):
               _lowerCAmelCase 			=  torch.manual_seed(__a			)
         else:
               _lowerCAmelCase 			=  torch.Generator(device=__a			).manual_seed(__a			)
         _lowerCAmelCase 			=  {
             """prompt""": """A painting of a squirrel eating a burger""",
             """image""": image,
             """generator""": generator,
             """num_inference_steps""": 2,
             """guidance_scale""": 5.0,
             """output_type""": """numpy""",
             """strength""": 0.75,
         }
         return inputs





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  """cpu"""  # ensure determinism for the device-dependent torch.Generator
         _lowerCAmelCase 			=  self.get_dummy_components()
         _lowerCAmelCase 			=  StableDiffusionXLImgaImgPipeline(**__a			)
         _lowerCAmelCase 			=  sd_pipe.to(__a			)
         sd_pipe.set_progress_bar_config(disable=__a			)

         _lowerCAmelCase 			=  self.get_dummy_inputs(__a			)
         _lowerCAmelCase 			=  sd_pipe(**__a			).images
         _lowerCAmelCase 			=  image[0, -3:, -3:, -1]

         assert image.shape == (1, 32, 32, 3)

         _lowerCAmelCase 			=  np.array([0.4_656, 0.4_840, 0.4_439, 0.6_698, 0.5_574, 0.4_524, 0.5_799, 0.5_943, 0.5_165]			)

         assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2





   def 	__lowerCAmelCase							(	self			):
         super().test_attention_slicing_forward_pass(expected_max_diff=3E-3			)





   def 	__lowerCAmelCase							(	self			):
         super().test_inference_batch_single_identical(expected_max_diff=3E-3			)





   def 	__lowerCAmelCase							(	self			):
         pass





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  self.get_dummy_components()
         _lowerCAmelCase 			=  StableDiffusionXLImgaImgPipeline(**__a			)
         _lowerCAmelCase 			=  sd_pipe.to(__a			)
         _lowerCAmelCase 			=  sd_pipe.to(__a			)
         sd_pipe.set_progress_bar_config(disable=__a			)

         # forward without prompt embeds
         _lowerCAmelCase 			=  self.get_dummy_inputs(__a			)
         _lowerCAmelCase 			=  3 * ["""this is a negative prompt"""]
         _lowerCAmelCase 			=  negative_prompt
         _lowerCAmelCase 			=  3 * [inputs["""prompt"""]]

         _lowerCAmelCase 			=  sd_pipe(**__a			)
         _lowerCAmelCase 			=  output.images[0, -3:, -3:, -1]

         # forward with prompt embeds
         _lowerCAmelCase 			=  self.get_dummy_inputs(__a			)
         _lowerCAmelCase 			=  3 * ["""this is a negative prompt"""]
         _lowerCAmelCase 			=  3 * [inputs.pop('''prompt'''			)]

         (
             _lowerCAmelCase			
         ) 			=  sd_pipe.encode_prompt(__a       ,     negative_prompt=__a			)

         _lowerCAmelCase 			=  sd_pipe(
             **__a       ,     prompt_embeds=__a       ,     negative_prompt_embeds=__a       ,     pooled_prompt_embeds=__a       ,     negative_pooled_prompt_embeds=__a       ,     )
         _lowerCAmelCase 			=  output.images[0, -3:, -3:, -1]

         # make sure that it's equal
         assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()			).max() < 1E-4






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





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





   def 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase       ,     _lowerCAmelCase="cpu"       ,     _lowerCAmelCase=torch.floataa       ,     _lowerCAmelCase=0			):
         _lowerCAmelCase 			=  torch.Generator(device=__a			).manual_seed(__a			)
         _lowerCAmelCase 			=  np.random.RandomState(__a			).standard_normal((1, 4, 64, 64)			)
         _lowerCAmelCase 			=  torch.from_numpy(__a			).to(device=__a       ,     dtype=__a			)
         _lowerCAmelCase 			=  {
             """prompt""": """a photograph of an astronaut riding a horse""",
             """latents""": latents,
             """generator""": generator,
             """num_inference_steps""": 3,
             """guidance_scale""": 7.5,
             """output_type""": """numpy""",
         }
         return inputs





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''			)
         pipe.to(__a			)
         pipe.set_progress_bar_config(disable=__a			)

         _lowerCAmelCase 			=  self.get_inputs(__a			)
         _lowerCAmelCase 			=  pipe(**__a			).images
         _lowerCAmelCase 			=  image[0, -3:, -3:, -1].flatten()

         assert image.shape == (1, 512, 512, 3)
         _lowerCAmelCase 			=  np.array([0.49_493, 0.47_896, 0.40_798, 0.54_214, 0.53_212, 0.48_202, 0.47_656, 0.46_329, 0.48_506]			)
         assert np.abs(image_slice - expected_slice			).max() < 7E-3











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						(						snake_case_       ,unittest.TestCase				):
   SCREAMING_SNAKE_CASE__							=     DiTPipeline
   SCREAMING_SNAKE_CASE__							=     CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
   SCREAMING_SNAKE_CASE__							=     PipelineTesterMixin.required_optional_params - {
       '''latents''',
       '''num_images_per_prompt''',
       '''callback''',
       '''callback_steps''',
   }
   SCREAMING_SNAKE_CASE__							=     CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
   SCREAMING_SNAKE_CASE__							=     False





   def 	__lowerCAmelCase							(	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=_lowerCAmelCase       ,     activation_fn='''gelu-approximate'''       ,     num_embeds_ada_norm=1_000       ,     norm_type='''ada_norm_zero'''       ,     norm_elementwise_affine=_lowerCAmelCase       ,     )
         _lowerCAmelCase 			=  AutoencoderKL()
         _lowerCAmelCase 			=  DDIMScheduler()
         _lowerCAmelCase 			=  {'''transformer''': transformer.eval(), '''vae''': vae.eval(), '''scheduler''': scheduler}
         return components





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





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  '''cpu'''

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

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

         self.assertEqual(image.shape       ,     (1, 16, 16, 3)			)
         _lowerCAmelCase 			=  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]			)
         _lowerCAmelCase 			=  np.abs(image_slice.flatten() - expected_slice			).max()
         self.assertLessEqual(_lowerCAmelCase       ,     1E-3			)





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






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





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





   def 	__lowerCAmelCase							(	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(_lowerCAmelCase			)

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

         for word, image in zip(_lowerCAmelCase       ,     _lowerCAmelCase			):
               _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 	__lowerCAmelCase							(	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(_lowerCAmelCase			)

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

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