




'''simple docstring'''



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 			lowerCAmelCase      (       __lowerCAmelCase ,      unittest.TestCase):
       __lowercase     :					Any       				=	UnCLIPPipeline
       __lowercase     :					Optional[int]       				=	TEXT_TO_IMAGE_PARAMS - {
           '''negative_prompt''',
           '''height''',
           '''width''',
           '''negative_prompt_embeds''',
           '''guidance_scale''',
           '''prompt_embeds''',
           '''cross_attention_kwargs''',
       }
       __lowercase     :					Tuple       				=	TEXT_TO_IMAGE_BATCH_PARAMS
       __lowercase     :					Tuple       				=	[
           '''generator''',
           '''return_dict''',
           '''prior_num_inference_steps''',
           '''decoder_num_inference_steps''',
           '''super_res_num_inference_steps''',
       ]
       __lowercase     :					Union[str, Any]       				=	False


       @property
       def 					lowerCAmelCase					(      self     )					->     Union[str, Any]:



            '''simple docstring'''




            return 32


       @property
       def 					lowerCAmelCase					(      self     )					->     Any:



            '''simple docstring'''




            return 32


       @property
       def 					lowerCAmelCase					(      self     )					->     Dict:



            '''simple docstring'''




            return self.time_input_dim


       @property
       def 					lowerCAmelCase					(      self     )					->     str:



            '''simple docstring'''




            return self.time_input_dim * 4


       @property
       def 					lowerCAmelCase					(      self     )					->     Any:



            '''simple docstring'''




            return 100


       @property
       def 					lowerCAmelCase					(      self     )					->     str:



            '''simple docstring'''




            __snake_case			    =     CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )
            return tokenizer


       @property
       def 					lowerCAmelCase					(      self     )					->     str:



            '''simple docstring'''




            torch.manual_seed(0     )
            __snake_case			    =     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-05			,     num_attention_heads=4			,     num_hidden_layers=5			,     pad_token_id=1			,     vocab_size=1000			,     )
            return CLIPTextModelWithProjection(__SCREAMING_SNAKE_CASE     )


       @property
       def 					lowerCAmelCase					(      self     )					->     Union[str, Any]:



            '''simple docstring'''




            torch.manual_seed(0     )

            __snake_case			    =     {
                '''num_attention_heads''': 2,
                '''attention_head_dim''': 12,
                '''embedding_dim''': self.text_embedder_hidden_size,
                '''num_layers''': 1,
            }

            __snake_case			    =     PriorTransformer(**__SCREAMING_SNAKE_CASE     )
            return model


       @property
       def 					lowerCAmelCase					(      self     )					->     Optional[Any]:



            '''simple docstring'''




            torch.manual_seed(0     )

            __snake_case			    =     {
                '''clip_embeddings_dim''': self.text_embedder_hidden_size,
                '''time_embed_dim''': self.time_embed_dim,
                '''cross_attention_dim''': self.cross_attention_dim,
            }

            __snake_case			    =     UnCLIPTextProjModel(**__SCREAMING_SNAKE_CASE     )
            return model


       @property
       def 					lowerCAmelCase					(      self     )					->     List[Any]:



            '''simple docstring'''




            torch.manual_seed(0     )

            __snake_case			    =     {
                '''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''',
            }

            __snake_case			    =     UNetaDConditionModel(**__SCREAMING_SNAKE_CASE     )
            return model


       @property
       def 					lowerCAmelCase					(      self     )					->     List[Any]:



            '''simple docstring'''




            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 					lowerCAmelCase					(      self     )					->     int:



            '''simple docstring'''




            torch.manual_seed(0     )

            __snake_case			    =     UNetaDModel(**self.dummy_super_res_kwargs     )
            return model


       @property
       def 					lowerCAmelCase					(      self     )					->     Dict:



            '''simple docstring'''




            torch.manual_seed(1     )

            __snake_case			    =     UNetaDModel(**self.dummy_super_res_kwargs     )
            return model


       def 					lowerCAmelCase					(      self     )					->     Dict:



            '''simple docstring'''




            __snake_case			    =     self.dummy_prior
            __snake_case			    =     self.dummy_decoder
            __snake_case			    =     self.dummy_text_proj
            __snake_case			    =     self.dummy_text_encoder
            __snake_case			    =     self.dummy_tokenizer
            __snake_case			    =     self.dummy_super_res_first
            __snake_case			    =     self.dummy_super_res_last

            __snake_case			    =     UnCLIPScheduler(
                variance_type='''fixed_small_log'''			,     prediction_type='''sample'''			,     num_train_timesteps=1000			,     clip_sample_range=5.0			,     )

            __snake_case			    =     UnCLIPScheduler(
                variance_type='''learned_range'''			,     prediction_type='''epsilon'''			,     num_train_timesteps=1000			,     )

            __snake_case			    =     UnCLIPScheduler(
                variance_type='''fixed_small_log'''			,     prediction_type='''epsilon'''			,     num_train_timesteps=1000			,     )

            __snake_case			    =     {
                '''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 					lowerCAmelCase					(      self			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE=0     )					->     Optional[Any]:



            '''simple docstring'''




            if str(__SCREAMING_SNAKE_CASE     ).startswith('''mps'''     ):
                 __snake_case			    =     torch.manual_seed(__SCREAMING_SNAKE_CASE     )
            else:
                 __snake_case			    =     torch.Generator(device=__SCREAMING_SNAKE_CASE     ).manual_seed(__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     {
                '''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 					lowerCAmelCase					(      self     )					->     Dict:



            '''simple docstring'''




            __snake_case			    =     '''cpu'''

            __snake_case			    =     self.get_dummy_components()

            __snake_case			    =     self.pipeline_class(**__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     pipe.to(__SCREAMING_SNAKE_CASE     )

            pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

            __snake_case			    =     pipe(**self.get_dummy_inputs(__SCREAMING_SNAKE_CASE     )     )
            __snake_case			    =     output.images

            __snake_case			    =     pipe(
                **self.get_dummy_inputs(__SCREAMING_SNAKE_CASE     )			,     return_dict=__SCREAMING_SNAKE_CASE			,     )[0]

            __snake_case			    =     image[0, -3:, -3:, -1]
            __snake_case			    =     image_from_tuple[0, -3:, -3:, -1]

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

            __snake_case			    =     np.array(
                [
                    0.9_997,
                    0.9_988,
                    0.0_028,
                    0.9_997,
                    0.9_984,
                    0.9_965,
                    0.0_029,
                    0.9_986,
                    0.0_025,
                ]     )

            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 					lowerCAmelCase					(      self     )					->     List[str]:



            '''simple docstring'''




            __snake_case			    =     torch.device('''cpu'''     )





            class 			lowerCAmelCase      :
                   __lowercase     :					List[str]       				=	1





            __snake_case			    =     self.get_dummy_components()

            __snake_case			    =     self.pipeline_class(**__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     pipe.to(__SCREAMING_SNAKE_CASE     )

            __snake_case			    =     components['''prior''']
            __snake_case			    =     components['''decoder''']
            __snake_case			    =     components['''super_res_first''']
            __snake_case			    =     components['''tokenizer''']
            __snake_case			    =     components['''text_encoder''']

            __snake_case			    =     torch.Generator(device=__SCREAMING_SNAKE_CASE     ).manual_seed(0     )
            __snake_case			    =     prior.dtype
            __snake_case			    =     1

            __snake_case			    =     (batch_size, prior.config.embedding_dim)
            __snake_case			    =     pipe.prepare_latents(
                __SCREAMING_SNAKE_CASE			,     dtype=__SCREAMING_SNAKE_CASE			,     device=__SCREAMING_SNAKE_CASE			,     generator=__SCREAMING_SNAKE_CASE			,     latents=__SCREAMING_SNAKE_CASE			,     scheduler=DummyScheduler()     )
            __snake_case			    =     (batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
            __snake_case			    =     pipe.prepare_latents(
                __SCREAMING_SNAKE_CASE			,     dtype=__SCREAMING_SNAKE_CASE			,     device=__SCREAMING_SNAKE_CASE			,     generator=__SCREAMING_SNAKE_CASE			,     latents=__SCREAMING_SNAKE_CASE			,     scheduler=DummyScheduler()     )

            __snake_case			    =     (
                batch_size,
                super_res_first.config.in_channels // 2,
                super_res_first.config.sample_size,
                super_res_first.config.sample_size,
            )
            __snake_case			    =     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     )

            __snake_case			    =     '''this is a prompt example'''

            __snake_case			    =     torch.Generator(device=__SCREAMING_SNAKE_CASE     ).manual_seed(0     )
            __snake_case			    =     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'''			,     )
            __snake_case			    =     output.images

            __snake_case			    =     tokenizer(
                __SCREAMING_SNAKE_CASE			,     padding='''max_length'''			,     max_length=tokenizer.model_max_length			,     return_tensors='''pt'''			,     )
            __snake_case			    =     text_encoder(text_inputs.input_ids     )
            __snake_case			    =     text_inputs.attention_mask

            __snake_case			    =     torch.Generator(device=__SCREAMING_SNAKE_CASE     ).manual_seed(0     )
            __snake_case			    =     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 					lowerCAmelCase					(      self     )					->     Union[str, Any]:



            '''simple docstring'''




            __snake_case			    =     torch_device == '''cpu'''

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


       @skip_mps
       def 					lowerCAmelCase					(      self     )					->     str:



            '''simple docstring'''




            __snake_case			    =     torch_device == '''cpu'''
            __snake_case			    =     True
            __snake_case			    =     [
                '''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 					lowerCAmelCase					(      self     )					->     List[Any]:



            '''simple docstring'''




            __snake_case			    =     [
                '''prior_num_inference_steps''',
                '''decoder_num_inference_steps''',
                '''super_res_num_inference_steps''',
            ]

            if torch_device == "mps":
                 # TODO: MPS errors with larger batch sizes
                 __snake_case			    =     [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 					lowerCAmelCase					(      self     )					->     List[Any]:



            '''simple docstring'''




            return super().test_dict_tuple_outputs_equivalent()


       @skip_mps
       def 					lowerCAmelCase					(      self     )					->     int:



            '''simple docstring'''




            return super().test_save_load_local()




       @skip_mps
       def 					lowerCAmelCase					(      self     )					->     List[str]:



            '''simple docstring'''




            return super().test_save_load_optional_components()





@nightly
class 			lowerCAmelCase      (       unittest.TestCase):


       def 					lowerCAmelCase					(      self     )					->     List[str]:



            '''simple docstring'''




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




       def 					lowerCAmelCase					(      self     )					->     Union[str, Any]:



            '''simple docstring'''




            __snake_case			    =     load_numpy(
                '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
                '''/unclip/karlo_v1_alpha_horse_cpu.npy'''     )

            __snake_case			    =     UnCLIPPipeline.from_pretrained('''kakaobrain/karlo-v1-alpha'''     )
            pipeline.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

            __snake_case			    =     torch.manual_seed(0     )
            __snake_case			    =     pipeline(
                '''horse'''			,     num_images_per_prompt=1			,     generator=__SCREAMING_SNAKE_CASE			,     output_type='''np'''			,     )

            __snake_case			    =     output.images[0]

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





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


       def 					lowerCAmelCase					(      self     )					->     Optional[Any]:



            '''simple docstring'''




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


       def 					lowerCAmelCase					(      self     )					->     Any:



            '''simple docstring'''




            __snake_case			    =     load_numpy(
                '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
                '''/unclip/karlo_v1_alpha_horse_fp16.npy'''     )

            __snake_case			    =     UnCLIPPipeline.from_pretrained('''kakaobrain/karlo-v1-alpha'''			,     torch_dtype=torch.floataa     )
            __snake_case			    =     pipeline.to(__SCREAMING_SNAKE_CASE     )
            pipeline.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )

            __snake_case			    =     torch.Generator(device='''cpu'''     ).manual_seed(0     )
            __snake_case			    =     pipeline(
                '''horse'''			,     generator=__SCREAMING_SNAKE_CASE			,     output_type='''np'''			,     )

            __snake_case			    =     output.images[0]

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

            assert_mean_pixel_difference(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE     )




       def 					lowerCAmelCase					(      self     )					->     Tuple:



            '''simple docstring'''




            torch.cuda.empty_cache()
            torch.cuda.reset_max_memory_allocated()
            torch.cuda.reset_peak_memory_stats()

            __snake_case			    =     UnCLIPPipeline.from_pretrained('''kakaobrain/karlo-v1-alpha'''			,     torch_dtype=torch.floataa     )
            __snake_case			    =     pipe.to(__SCREAMING_SNAKE_CASE     )
            pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE     )
            pipe.enable_attention_slicing()
            pipe.enable_sequential_cpu_offload()

            __snake_case			    =     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'''			,     )

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












'''simple docstring'''



import warnings

from ...utils import logging
from .image_processing_detr import DetrImageProcessor


UpperCAmelCase_				:     Dict     			=						logging.get_logger(__name__)





class 			lowerCAmelCase      (       __lowerCAmelCase):


       def __init__(      self			,     *__SCREAMING_SNAKE_CASE			,     **__SCREAMING_SNAKE_CASE     )					->     None:



            '''simple docstring'''




            warnings.warn(
                '''The class DetrFeatureExtractor is deprecated and will be removed in version 5 of Transformers.'''
                ''' Please use DetrImageProcessor instead.'''			,     __SCREAMING_SNAKE_CASE			,     )
            super().__init__(*__SCREAMING_SNAKE_CASE			,     **__SCREAMING_SNAKE_CASE     )






