



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							(			lowerCAmelCase_  ,							unittest.TestCase				):
		__lowerCAmelCase	:    Tuple											=					DiTPipeline
		__lowerCAmelCase	:    int											=					CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
		__lowerCAmelCase	:    List[str]											=					PipelineTesterMixin.required_optional_params - {
		    'latents',
		    'num_images_per_prompt',
		    'callback',
		    'callback_steps',
		}
		__lowerCAmelCase	:    Tuple											=					CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
		__lowerCAmelCase	:    int											=					False



		def 	__lowerCamelCase			(     self     :str  ):
						torch.manual_seed(0  )
						snake_case__	:      str     		=	TransformeraDModel(
						    sample_size=1_6   ,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_0_0_0   ,norm_type='''ada_norm_zero'''   ,norm_elementwise_affine=__lowerCAmelCase   ,)
						snake_case__	:      List[str]     		=	AutoencoderKL()
						snake_case__	:      List[Any]     		=	DDIMScheduler()
						snake_case__	:      List[str]     		=	{'''transformer''': transformer.eval(), '''vae''': vae.eval(), '''scheduler''': scheduler}
						return components



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



		def 	__lowerCamelCase			(     self     :Any  ):
						snake_case__	:      Union[str, Any]     		=	'''cpu'''

						snake_case__	:      Dict     		=	self.get_dummy_components()
						snake_case__	:      Dict     		=	self.pipeline_class(**__lowerCAmelCase  )
						pipe.to(__lowerCAmelCase  )
						pipe.set_progress_bar_config(disable=__lowerCAmelCase  )

						snake_case__	:      List[str]     		=	self.get_dummy_inputs(__lowerCAmelCase  )
						snake_case__	:      Optional[int]     		=	pipe(**__lowerCAmelCase  ).images
						snake_case__	:      str     		=	image[0, -3:, -3:, -1]

						self.assertEqual(image.shape   ,(1, 1_6, 1_6, 3)  )
						snake_case__	:      Any     		=	np.array([0.2946, 0.6601, 0.4329, 0.3296, 0.4144, 0.5319, 0.7273, 0.5013, 0.4457]  )
						snake_case__	:      List[str]     		=	np.abs(image_slice.flatten() - expected_slice  ).max()
						self.assertLessEqual(__lowerCAmelCase   ,1e-3  )



		def 	__lowerCamelCase			(     self     :Optional[int]  ):
						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     :Optional[int]  ):
						self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3  )


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



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



		def 	__lowerCamelCase			(     self     :List[Any]  ):
						snake_case__	:      int     		=	torch.manual_seed(0  )

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

						snake_case__	:      Optional[Any]     		=	['''vase''', '''umbrella''', '''white shark''', '''white wolf''']
						snake_case__	:      Tuple     		=	pipe.get_label_ids(__lowerCAmelCase  )

						snake_case__	:      Optional[int]     		=	pipe(__lowerCAmelCase   ,generator=__lowerCAmelCase   ,num_inference_steps=4_0   ,output_type='''np'''  ).images

						for word, image in zip(__lowerCAmelCase   ,__lowerCAmelCase  ):
										snake_case__	:      Optional[int]     		=	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     :int  ):
						snake_case__	:      int     		=	DiTPipeline.from_pretrained('''facebook/DiT-XL-2-512'''  )
						snake_case__	:      Union[str, Any]     		=	DPMSolverMultistepScheduler.from_config(pipe.scheduler.config  )
						pipe.to('''cuda'''  )

						snake_case__	:      List[Any]     		=	['''vase''', '''umbrella''']
						snake_case__	:      List[str]     		=	pipe.get_label_ids(__lowerCAmelCase  )

						snake_case__	:      Optional[Any]     		=	torch.manual_seed(0  )
						snake_case__	:      Tuple     		=	pipe(__lowerCAmelCase   ,generator=__lowerCAmelCase   ,num_inference_steps=2_5   ,output_type='''np'''  ).images

						for word, image in zip(__lowerCAmelCase   ,__lowerCAmelCase  ):
										snake_case__	:      Tuple     		=	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

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()


class 				a (						lowerCAmelCase_		,							lowerCAmelCase_		,							lowerCAmelCase_		,							unittest.TestCase      ):
    _snake_case   :     int   							=							StableUnCLIPPipeline
    _snake_case   :     str   							=							TEXT_TO_IMAGE_PARAMS
    _snake_case   :     Any   							=							TEXT_TO_IMAGE_BATCH_PARAMS
    _snake_case   :     Optional[Any]   							=							TEXT_TO_IMAGE_IMAGE_PARAMS
    _snake_case   :     str   							=							TEXT_TO_IMAGE_IMAGE_PARAMS

    # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
    _snake_case   :     str   							=							False


    def    lowerCAmelCase_     (    self		:     Optional[int]    ):
      _UpperCAmelCase         =					32
      _UpperCAmelCase         =					embedder_hidden_size

      # prior components

      torch.manual_seed(0    )
      _UpperCAmelCase         =					CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""    )

      torch.manual_seed(0    )
      _UpperCAmelCase         =					CLIPTextModelWithProjection(
          CLIPTextConfig(
              bos_token_id=0      , eos_token_id=2      , hidden_size=__lowerCAmelCase      , projection_dim=__lowerCAmelCase      , intermediate_size=37      , layer_norm_eps=1e-0_5      , num_attention_heads=4      , num_hidden_layers=5      , pad_token_id=1      , vocab_size=1000      , )    )

      torch.manual_seed(0    )
      _UpperCAmelCase         =					PriorTransformer(
          num_attention_heads=2      , attention_head_dim=12      , embedding_dim=__lowerCAmelCase      , num_layers=1      , )

      torch.manual_seed(0    )
      _UpperCAmelCase         =					DDPMScheduler(
          variance_type="""fixed_small_log"""      , prediction_type="""sample"""      , num_train_timesteps=1000      , clip_sample=__lowerCAmelCase      , clip_sample_range=5.0      , beta_schedule="""squaredcos_cap_v2"""      , )

      # regular denoising components

      torch.manual_seed(0    )
      _UpperCAmelCase         =					StableUnCLIPImageNormalizer(embedding_dim=__lowerCAmelCase    )
      _UpperCAmelCase         =					DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""    )

      torch.manual_seed(0    )
      _UpperCAmelCase         =					CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""    )

      torch.manual_seed(0    )
      _UpperCAmelCase         =					CLIPTextModel(
          CLIPTextConfig(
              bos_token_id=0      , eos_token_id=2      , hidden_size=__lowerCAmelCase      , projection_dim=32      , intermediate_size=37      , layer_norm_eps=1e-0_5      , num_attention_heads=4      , num_hidden_layers=5      , pad_token_id=1      , vocab_size=1000      , )    )

      torch.manual_seed(0    )
      _UpperCAmelCase         =					UNetaDConditionModel(
          sample_size=32      , in_channels=4      , out_channels=4      , down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")      , up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")      , block_out_channels=(32, 64)      , attention_head_dim=(2, 4)      , class_embed_type="""projection"""      , projection_class_embeddings_input_dim=embedder_projection_dim * 2      , cross_attention_dim=__lowerCAmelCase      , layers_per_block=1      , upcast_attention=__lowerCAmelCase      , use_linear_projection=__lowerCAmelCase      , )

      torch.manual_seed(0    )
      _UpperCAmelCase         =					DDIMScheduler(
          beta_schedule="""scaled_linear"""      , beta_start=0.00_085      , beta_end=0.012      , prediction_type="""v_prediction"""      , set_alpha_to_one=__lowerCAmelCase      , steps_offset=1      , )

      torch.manual_seed(0    )
      _UpperCAmelCase         =					AutoencoderKL()

      _UpperCAmelCase         =					{
          # prior components
          """prior_tokenizer""": prior_tokenizer,
          """prior_text_encoder""": prior_text_encoder,
          """prior""": prior,
          """prior_scheduler""": prior_scheduler,
          # image noising components
          """image_normalizer""": image_normalizer,
          """image_noising_scheduler""": image_noising_scheduler,
          # regular denoising components
          """tokenizer""": tokenizer,
          """text_encoder""": text_encoder,
          """unet""": unet,
          """scheduler""": scheduler,
          """vae""": vae,
      }

      return components


    def    lowerCAmelCase_     (    self		:     Optional[int]      , __lowerCAmelCase		:     List[Any]      , __lowerCAmelCase		:     str=0    ):
      if str(__lowerCAmelCase    ).startswith("""mps"""    ):
        _UpperCAmelCase         =					torch.manual_seed(__lowerCAmelCase    )
      else:
        _UpperCAmelCase         =					torch.Generator(device=__lowerCAmelCase    ).manual_seed(__lowerCAmelCase    )
      _UpperCAmelCase         =					{
          """prompt""": """A painting of a squirrel eating a burger""",
          """generator""": generator,
          """num_inference_steps""": 2,
          """prior_num_inference_steps""": 2,
          """output_type""": """numpy""",
      }
      return inputs


    def    lowerCAmelCase_     (    self		:     Optional[int]    ):
      _UpperCAmelCase         =					torch_device == """cpu"""

      self._test_attention_slicing_forward_pass(test_max_difference=__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     List[str]    ):
      _UpperCAmelCase         =					torch_device in ["""cpu""", """mps"""]

      self._test_inference_batch_single_identical(test_max_difference=__lowerCAmelCase    )



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


    def    lowerCAmelCase_     (    self		:     str    ):
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()


    def    lowerCAmelCase_     (    self		:     List[Any]    ):
      _UpperCAmelCase         =					load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"""    )

      _UpperCAmelCase         =					StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""      , torch_dtype=torch.floataa    )
      pipe.to(__lowerCAmelCase    )
      pipe.set_progress_bar_config(disable=__lowerCAmelCase    )
      # stable unclip will oom when integration tests are run on a V100,
      # so turn on memory savings
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      _UpperCAmelCase         =					torch.Generator(device="""cpu"""    ).manual_seed(0    )
      _UpperCAmelCase         =					pipe("""anime turle"""      , generator=__lowerCAmelCase      , output_type="""np"""    )

      _UpperCAmelCase         =					output.images[0]

      assert image.shape == (768, 768, 3)

      assert_mean_pixel_difference(__lowerCAmelCase      , __lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     Any    ):
      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()

      _UpperCAmelCase         =					StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""      , torch_dtype=torch.floataa    )
      _UpperCAmelCase         =					pipe.to(__lowerCAmelCase    )
      pipe.set_progress_bar_config(disable=__lowerCAmelCase    )
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      _UpperCAmelCase         =					pipe(
          """anime turtle"""      , prior_num_inference_steps=2      , num_inference_steps=2      , output_type="""np"""      , )

      _UpperCAmelCase         =					torch.cuda.max_memory_allocated()
      # make sure that less than 7 GB is allocated
      assert mem_bytes < 7 * 10**9
