


import gc
import unittest

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

from diffusers import PriorTransformer, UnCLIPPipeline, UnCLIPScheduler, UNetaDConditionModel, UNetaDModel
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()



class 							a       (      _UpperCAmelCase  ,unittest.TestCase ):
     UpperCAmelCase__			:	str       =						UnCLIPPipeline
     UpperCAmelCase__			:	Union[str, Any]       =						TEXT_TO_IMAGE_PARAMS - {
         "negative_prompt",
         "height",
         "width",
         "negative_prompt_embeds",
         "guidance_scale",
         "prompt_embeds",
         "cross_attention_kwargs",
     }
     UpperCAmelCase__			:	Optional[int]       =						TEXT_TO_IMAGE_BATCH_PARAMS
     UpperCAmelCase__			:	Any       =						[
         "generator",
         "return_dict",
         "prior_num_inference_steps",
         "decoder_num_inference_steps",
         "super_res_num_inference_steps",
     ]
     UpperCAmelCase__			:	Optional[Any]       =						False





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					str				):
        return 32





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					Dict				):
        return 32





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					Union[str, Any]				):
        return self.time_input_dim





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					Any				):
        return self.time_input_dim * 4





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					Any				):
        return 100





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					Tuple				):
        __lowerCamelCase:					Optional[int]	 =						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)
        return tokenizer





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					Union[str, Any]				):
        torch.manual_seed(0				)
        __lowerCamelCase:					Optional[Any]	 =						CLIPTextConfig(
            bos_token_id=0     ,						eos_token_id=2     ,						hidden_size=self.text_embedder_hidden_size     ,						projection_dim=self.text_embedder_hidden_size     ,						intermediate_size=37     ,						layer_norm_eps=1E-0_5     ,						num_attention_heads=4     ,						num_hidden_layers=5     ,						pad_token_id=1     ,						vocab_size=1000     ,						)
        return CLIPTextModelWithProjection(SCREAMING_SNAKE_CASE_				)





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					List[Any]				):
        torch.manual_seed(0				)

        __lowerCamelCase:					int	 =						{
            """num_attention_heads""": 2,
            """attention_head_dim""": 12,
            """embedding_dim""": self.text_embedder_hidden_size,
            """num_layers""": 1,
        }

        __lowerCamelCase:					Any	 =						PriorTransformer(**SCREAMING_SNAKE_CASE_				)
        return model





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					str				):
        torch.manual_seed(0				)

        __lowerCamelCase:					Optional[int]	 =						{
            """clip_embeddings_dim""": self.text_embedder_hidden_size,
            """time_embed_dim""": self.time_embed_dim,
            """cross_attention_dim""": self.cross_attention_dim,
        }

        __lowerCamelCase:					Optional[int]	 =						UnCLIPTextProjModel(**SCREAMING_SNAKE_CASE_				)
        return model





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					Optional[Any]				):
        torch.manual_seed(0				)

        __lowerCamelCase:					Optional[Any]	 =						{
            """sample_size""": 32,
            # RGB in channels
            """in_channels""": 3,
            # Out channels is double in channels because predicts mean and variance
            """out_channels""": 6,
            """down_block_types""": ("""ResnetDownsampleBlock2D""", """SimpleCrossAttnDownBlock2D"""),
            """up_block_types""": ("""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""),
            """mid_block_type""": """UNetMidBlock2DSimpleCrossAttn""",
            """block_out_channels""": (self.block_out_channels_a, self.block_out_channels_a * 2),
            """layers_per_block""": 1,
            """cross_attention_dim""": self.cross_attention_dim,
            """attention_head_dim""": 4,
            """resnet_time_scale_shift""": """scale_shift""",
            """class_embed_type""": """identity""",
        }

        __lowerCamelCase:					List[str]	 =						UNetaDConditionModel(**SCREAMING_SNAKE_CASE_				)
        return model





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					Tuple				):
        return {
            "sample_size": 64,
            "layers_per_block": 1,
            "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
            "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
            "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
            "in_channels": 6,
            "out_channels": 3,
        }





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					List[Any]				):
        torch.manual_seed(0				)

        __lowerCamelCase:					Tuple	 =						UNetaDModel(**self.dummy_super_res_kwargs				)
        return model





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					Dict				):
        # seeded differently to get different unet than `self.dummy_super_res_first`
        torch.manual_seed(1				)

        __lowerCamelCase:					Dict	 =						UNetaDModel(**self.dummy_super_res_kwargs				)
        return model





     def 			SCREAMING_SNAKE_CASE__							(     self       :					List[str]				):
        __lowerCamelCase:					str	 =						self.dummy_prior
        __lowerCamelCase:					List[str]	 =						self.dummy_decoder
        __lowerCamelCase:					List[str]	 =						self.dummy_text_proj
        __lowerCamelCase:					Union[str, Any]	 =						self.dummy_text_encoder
        __lowerCamelCase:					List[str]	 =						self.dummy_tokenizer
        __lowerCamelCase:					Any	 =						self.dummy_super_res_first
        __lowerCamelCase:					List[str]	 =						self.dummy_super_res_last

        __lowerCamelCase:					Optional[Any]	 =						UnCLIPScheduler(
            variance_type="""fixed_small_log"""     ,						prediction_type="""sample"""     ,						num_train_timesteps=1000     ,						clip_sample_range=5.0     ,						)

        __lowerCamelCase:					List[str]	 =						UnCLIPScheduler(
            variance_type="""learned_range"""     ,						prediction_type="""epsilon"""     ,						num_train_timesteps=1000     ,						)

        __lowerCamelCase:					Union[str, Any]	 =						UnCLIPScheduler(
            variance_type="""fixed_small_log"""     ,						prediction_type="""epsilon"""     ,						num_train_timesteps=1000     ,						)

        __lowerCamelCase:					Union[str, Any]	 =						{
            """prior""": prior,
            """decoder""": decoder,
            """text_proj""": text_proj,
            """text_encoder""": text_encoder,
            """tokenizer""": tokenizer,
            """super_res_first""": super_res_first,
            """super_res_last""": super_res_last,
            """prior_scheduler""": prior_scheduler,
            """decoder_scheduler""": decoder_scheduler,
            """super_res_scheduler""": super_res_scheduler,
        }

        return components





     def 			SCREAMING_SNAKE_CASE__							(     self       :					List[Any]     ,						SCREAMING_SNAKE_CASE_       :					Optional[Any]     ,						SCREAMING_SNAKE_CASE_       :					int=0				):
        if str(SCREAMING_SNAKE_CASE_				).startswith("""mps"""				):
           __lowerCamelCase:					str	 =						torch.manual_seed(SCREAMING_SNAKE_CASE_				)
        else:
           __lowerCamelCase:					Optional[Any]	 =						torch.Generator(device=SCREAMING_SNAKE_CASE_				).manual_seed(SCREAMING_SNAKE_CASE_				)
        __lowerCamelCase:					str	 =						{
            """prompt""": """horse""",
            """generator""": generator,
            """prior_num_inference_steps""": 2,
            """decoder_num_inference_steps""": 2,
            """super_res_num_inference_steps""": 2,
            """output_type""": """numpy""",
        }
        return inputs





     def 			SCREAMING_SNAKE_CASE__							(     self       :					Dict				):
        __lowerCamelCase:					str	 =						"""cpu"""

        __lowerCamelCase:					Optional[int]	 =						self.get_dummy_components()

        __lowerCamelCase:					Dict	 =						self.pipeline_class(**SCREAMING_SNAKE_CASE_				)
        __lowerCamelCase:					int	 =						pipe.to(SCREAMING_SNAKE_CASE_				)

        pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

        __lowerCamelCase:					Union[str, Any]	 =						pipe(**self.get_dummy_inputs(SCREAMING_SNAKE_CASE_				)				)
        __lowerCamelCase:					int	 =						output.images

        __lowerCamelCase:					str	 =						pipe(
            **self.get_dummy_inputs(SCREAMING_SNAKE_CASE_				)     ,						return_dict=SCREAMING_SNAKE_CASE_     ,						)[0]

        __lowerCamelCase:					List[Any]	 =						image[0, -3:, -3:, -1]
        __lowerCamelCase:					List[str]	 =						image_from_tuple[0, -3:, -3:, -1]

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

        __lowerCamelCase:					Dict	 =						np.array(
            [
                0.9997,
                0.9988,
                0.0028,
                0.9997,
                0.9984,
                0.9965,
                0.0029,
                0.9986,
                0.0025,
            ]				)

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





     def 			SCREAMING_SNAKE_CASE__							(     self       :					Union[str, Any]				):
        __lowerCamelCase:					Dict	 =						torch.device("""cpu"""				)



        class 							a       :
             UpperCAmelCase__			:	Union[str, Any]       =						1



        __lowerCamelCase:					Dict	 =						self.get_dummy_components()

        __lowerCamelCase:					List[Any]	 =						self.pipeline_class(**SCREAMING_SNAKE_CASE_				)
        __lowerCamelCase:					Optional[Any]	 =						pipe.to(SCREAMING_SNAKE_CASE_				)

        __lowerCamelCase:					Optional[int]	 =						components["""prior"""]
        __lowerCamelCase:					int	 =						components["""decoder"""]
        __lowerCamelCase:					Union[str, Any]	 =						components["""super_res_first"""]
        __lowerCamelCase:					Tuple	 =						components["""tokenizer"""]
        __lowerCamelCase:					Optional[int]	 =						components["""text_encoder"""]

        __lowerCamelCase:					Tuple	 =						torch.Generator(device=SCREAMING_SNAKE_CASE_				).manual_seed(0				)
        __lowerCamelCase:					Optional[Any]	 =						prior.dtype
        __lowerCamelCase:					Dict	 =						1

        __lowerCamelCase:					List[str]	 =						(batch_size, prior.config.embedding_dim)
        __lowerCamelCase:					Optional[Any]	 =						pipe.prepare_latents(
            SCREAMING_SNAKE_CASE_     ,						dtype=SCREAMING_SNAKE_CASE_     ,						device=SCREAMING_SNAKE_CASE_     ,						generator=SCREAMING_SNAKE_CASE_     ,						latents=SCREAMING_SNAKE_CASE_     ,						scheduler=DummyScheduler()				)
        __lowerCamelCase:					List[Any]	 =						(batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
        __lowerCamelCase:					Optional[Any]	 =						pipe.prepare_latents(
            SCREAMING_SNAKE_CASE_     ,						dtype=SCREAMING_SNAKE_CASE_     ,						device=SCREAMING_SNAKE_CASE_     ,						generator=SCREAMING_SNAKE_CASE_     ,						latents=SCREAMING_SNAKE_CASE_     ,						scheduler=DummyScheduler()				)

        __lowerCamelCase:					Tuple	 =						(
            batch_size,
            super_res_first.config.in_channels // 2,
            super_res_first.config.sample_size,
            super_res_first.config.sample_size,
        )
        __lowerCamelCase:					Dict	 =						pipe.prepare_latents(
            SCREAMING_SNAKE_CASE_     ,						dtype=SCREAMING_SNAKE_CASE_     ,						device=SCREAMING_SNAKE_CASE_     ,						generator=SCREAMING_SNAKE_CASE_     ,						latents=SCREAMING_SNAKE_CASE_     ,						scheduler=DummyScheduler()				)

        pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

        __lowerCamelCase:					Any	 =						"""this is a prompt example"""

        __lowerCamelCase:					Tuple	 =						torch.Generator(device=SCREAMING_SNAKE_CASE_				).manual_seed(0				)
        __lowerCamelCase:					Tuple	 =						pipe(
            [prompt]     ,						generator=SCREAMING_SNAKE_CASE_     ,						prior_num_inference_steps=2     ,						decoder_num_inference_steps=2     ,						super_res_num_inference_steps=2     ,						prior_latents=SCREAMING_SNAKE_CASE_     ,						decoder_latents=SCREAMING_SNAKE_CASE_     ,						super_res_latents=SCREAMING_SNAKE_CASE_     ,						output_type="""np"""     ,						)
        __lowerCamelCase:					Any	 =						output.images

        __lowerCamelCase:					Any	 =						tokenizer(
            SCREAMING_SNAKE_CASE_     ,						padding="""max_length"""     ,						max_length=tokenizer.model_max_length     ,						return_tensors="""pt"""     ,						)
        __lowerCamelCase:					List[Any]	 =						text_encoder(text_inputs.input_ids				)
        __lowerCamelCase:					List[str]	 =						text_inputs.attention_mask

        __lowerCamelCase:					Optional[int]	 =						torch.Generator(device=SCREAMING_SNAKE_CASE_				).manual_seed(0				)
        __lowerCamelCase:					Dict	 =						pipe(
            generator=SCREAMING_SNAKE_CASE_     ,						prior_num_inference_steps=2     ,						decoder_num_inference_steps=2     ,						super_res_num_inference_steps=2     ,						prior_latents=SCREAMING_SNAKE_CASE_     ,						decoder_latents=SCREAMING_SNAKE_CASE_     ,						super_res_latents=SCREAMING_SNAKE_CASE_     ,						text_model_output=SCREAMING_SNAKE_CASE_     ,						text_attention_mask=SCREAMING_SNAKE_CASE_     ,						output_type="""np"""     ,						)[0]

        # make sure passing text embeddings manually is identical
        assert np.abs(image - image_from_text				).max() < 1E-4





     @skip_mps
     def 			SCREAMING_SNAKE_CASE__							(     self       :					List[Any]				):
        __lowerCamelCase:					Dict	 =						torch_device == """cpu"""

        self._test_attention_slicing_forward_pass(test_max_difference=SCREAMING_SNAKE_CASE_     ,						expected_max_diff=0.01				)





     @skip_mps
     def 			SCREAMING_SNAKE_CASE__							(     self       :					str				):
        __lowerCamelCase:					List[str]	 =						torch_device == """cpu"""
        __lowerCamelCase:					List[Any]	 =						True
        __lowerCamelCase:					int	 =						[
            """prior_num_inference_steps""",
            """decoder_num_inference_steps""",
            """super_res_num_inference_steps""",
        ]

        self._test_inference_batch_single_identical(
            test_max_difference=SCREAMING_SNAKE_CASE_     ,						relax_max_difference=SCREAMING_SNAKE_CASE_     ,						additional_params_copy_to_batched_inputs=SCREAMING_SNAKE_CASE_     ,						)





     def 			SCREAMING_SNAKE_CASE__							(     self       :					Union[str, Any]				):
        __lowerCamelCase:					Dict	 =						[
            """prior_num_inference_steps""",
            """decoder_num_inference_steps""",
            """super_res_num_inference_steps""",
        ]

        if torch_device == "mps":
           # TODO: MPS errors with larger batch sizes
           __lowerCamelCase:					str	 =						[2, 3]
           self._test_inference_batch_consistent(
               batch_sizes=SCREAMING_SNAKE_CASE_     ,						additional_params_copy_to_batched_inputs=SCREAMING_SNAKE_CASE_     ,						)
        else:
           self._test_inference_batch_consistent(
               additional_params_copy_to_batched_inputs=SCREAMING_SNAKE_CASE_				)





     @skip_mps
     def 			SCREAMING_SNAKE_CASE__							(     self       :					Optional[Any]				):
        return super().test_dict_tuple_outputs_equivalent()





     @skip_mps
     def 			SCREAMING_SNAKE_CASE__							(     self       :					Optional[Any]				):
        return super().test_save_load_local()







     @skip_mps
     def 			SCREAMING_SNAKE_CASE__							(     self       :					Dict				):
        return super().test_save_load_optional_components()



@nightly
class 							a       (      unittest.TestCase ):





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







     def 			SCREAMING_SNAKE_CASE__							(     self       :					int				):
        __lowerCamelCase:					Union[str, Any]	 =						load_numpy(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
            """/unclip/karlo_v1_alpha_horse_cpu.npy"""				)

        __lowerCamelCase:					Dict	 =						UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""				)
        pipeline.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

        __lowerCamelCase:					Dict	 =						torch.manual_seed(0				)
        __lowerCamelCase:					Tuple	 =						pipeline(
            """horse"""     ,						num_images_per_prompt=1     ,						generator=SCREAMING_SNAKE_CASE_     ,						output_type="""np"""     ,						)

        __lowerCamelCase:					str	 =						output.images[0]

        assert image.shape == (256, 256, 3)
        assert np.abs(expected_image - image				).max() < 1E-1



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





     def 			SCREAMING_SNAKE_CASE__							(     self       :					Optional[Any]				):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()





     def 			SCREAMING_SNAKE_CASE__							(     self       :					int				):
        __lowerCamelCase:					Any	 =						load_numpy(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
            """/unclip/karlo_v1_alpha_horse_fp16.npy"""				)

        __lowerCamelCase:					int	 =						UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""     ,						torch_dtype=torch.floataa				)
        __lowerCamelCase:					str	 =						pipeline.to(SCREAMING_SNAKE_CASE_				)
        pipeline.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

        __lowerCamelCase:					int	 =						torch.Generator(device="""cpu"""				).manual_seed(0				)
        __lowerCamelCase:					Dict	 =						pipeline(
            """horse"""     ,						generator=SCREAMING_SNAKE_CASE_     ,						output_type="""np"""     ,						)

        __lowerCamelCase:					Optional[Any]	 =						output.images[0]

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

        assert_mean_pixel_difference(SCREAMING_SNAKE_CASE_     ,						SCREAMING_SNAKE_CASE_				)







     def 			SCREAMING_SNAKE_CASE__							(     self       :					int				):
        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        __lowerCamelCase:					str	 =						UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""     ,						torch_dtype=torch.floataa				)
        __lowerCamelCase:					Union[str, Any]	 =						pipe.to(SCREAMING_SNAKE_CASE_				)
        pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        __lowerCamelCase:					Optional[Any]	 =						pipe(
            """horse"""     ,						num_images_per_prompt=1     ,						prior_num_inference_steps=2     ,						decoder_num_inference_steps=2     ,						super_res_num_inference_steps=2     ,						output_type="""np"""     ,						)

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






_A		:					Union[str, Any]          =							1_000



def      __lowerCAmelCase    (				snake_case						:     float			)							->				float:
   return x * x



def      __lowerCAmelCase    (				snake_case						:     Optional[int]    ,   snake_case						:     float    ,   snake_case						:     float    ,   snake_case						:     int = 4			)							->				float:
   assert callable(
       snake_case			), f'the function(object) passed should be callable your input : {function}'
   assert isinstance(snake_case    ,   (float, int)			), f'a should be float or integer your input : {a}'
   assert isinstance(function(snake_case			)    ,   (float, int)			), (
       "the function should return integer or float return type of your function, "
       f'{type(snake_case			)}'
   )
   assert isinstance(snake_case    ,   (float, int)			), f'b should be float or integer your input : {b}'
   assert (
       isinstance(snake_case    ,   snake_case			) and precision > 0
   ), f'precision should be positive integer your input : {precision}'

   # just applying the formula of simpson for approximate integration written in
   # mentioned article in first comment of this file and above this function

   __lowerCamelCase:					List[str]	 =						(b - a) / N_STEPS
   __lowerCamelCase:					Optional[Any]	 =						function(snake_case			) + function(snake_case			)

   for i in range(1    ,   snake_case			):
      __lowerCamelCase:					Dict	 =						a + h * i
      result += function(snake_case			) * (4 if i % 2 else 2)

   result *= h / 3
   return round(snake_case    ,   snake_case			)


if __name__ == "__main__":
   import doctest

   doctest.testmod()


