



'''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 							lowercase (						_lowerCamelCase      ,      unittest.TestCase       ):


 """simple docstring"""


 UpperCAmelCase					=						UnCLIPPipeline
 UpperCAmelCase					=						TEXT_TO_IMAGE_PARAMS - {
     """negative_prompt""",
     """height""",
     """width""",
     """negative_prompt_embeds""",
     """guidance_scale""",
     """prompt_embeds""",
     """cross_attention_kwargs""",
 }
 UpperCAmelCase					=						TEXT_TO_IMAGE_BATCH_PARAMS
 UpperCAmelCase					=						[
     """generator""",
     """return_dict""",
     """prior_num_inference_steps""",
     """decoder_num_inference_steps""",
     """super_res_num_inference_steps""",
 ]
 UpperCAmelCase					=						False







 @property
 def _snake_case				(							self							)	->     Dict:
        return 32







 @property
 def _snake_case				(							self							)	->     Optional[int]:
        return 32







 @property
 def _snake_case				(							self							)	->     int:
        return self.time_input_dim







 @property
 def _snake_case				(							self							)	->     Optional[int]:
        return self.time_input_dim * 4







 @property
 def _snake_case				(							self							)	->     List[Any]:
        return 100







 @property
 def _snake_case				(							self							)	->     Optional[int]:
        _UpperCAmelCase						:		str      =				CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""							)
        return tokenizer







 @property
 def _snake_case				(							self							)	->     Tuple:
        torch.manual_seed(0							)
        _UpperCAmelCase						:		Union[str, 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=1_000						,)
        return CLIPTextModelWithProjection(a_							)







 @property
 def _snake_case				(							self							)	->     Tuple:
        torch.manual_seed(0							)

        _UpperCAmelCase						:		Optional[Any]      =				{
            """num_attention_heads""": 2,
            """attention_head_dim""": 12,
            """embedding_dim""": self.text_embedder_hidden_size,
            """num_layers""": 1,
        }

        _UpperCAmelCase						:		Tuple      =				PriorTransformer(**a_							)
        return model







 @property
 def _snake_case				(							self							)	->     str:
        torch.manual_seed(0							)

        _UpperCAmelCase						:		List[str]      =				{
            """clip_embeddings_dim""": self.text_embedder_hidden_size,
            """time_embed_dim""": self.time_embed_dim,
            """cross_attention_dim""": self.cross_attention_dim,
        }

        _UpperCAmelCase						:		str      =				UnCLIPTextProjModel(**a_							)
        return model







 @property
 def _snake_case				(							self							)	->     Tuple:
        torch.manual_seed(0							)

        _UpperCAmelCase						:		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""",
        }

        _UpperCAmelCase						:		List[str]      =				UNetaDConditionModel(**a_							)
        return model







 @property
 def _snake_case				(							self							)	->     Union[str, Any]:
        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 _snake_case				(							self							)	->     str:
        torch.manual_seed(0							)

        _UpperCAmelCase						:		List[Any]      =				UNetaDModel(**self.dummy_super_res_kwargs							)
        return model







 @property
 def _snake_case				(							self							)	->     int:
        # seeded differently to get different unet than `self.dummy_super_res_first`
        torch.manual_seed(1							)

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







 def _snake_case				(							self							)	->     Dict:
        _UpperCAmelCase						:		List[str]      =				self.dummy_prior
        _UpperCAmelCase						:		Optional[Any]      =				self.dummy_decoder
        _UpperCAmelCase						:		int      =				self.dummy_text_proj
        _UpperCAmelCase						:		List[Any]      =				self.dummy_text_encoder
        _UpperCAmelCase						:		List[str]      =				self.dummy_tokenizer
        _UpperCAmelCase						:		Optional[Any]      =				self.dummy_super_res_first
        _UpperCAmelCase						:		List[str]      =				self.dummy_super_res_last

        _UpperCAmelCase						:		Tuple      =				UnCLIPScheduler(
            variance_type="""fixed_small_log"""						,prediction_type="""sample"""						,num_train_timesteps=1_000						,clip_sample_range=5.0						,)

        _UpperCAmelCase						:		List[Any]      =				UnCLIPScheduler(
            variance_type="""learned_range"""						,prediction_type="""epsilon"""						,num_train_timesteps=1_000						,)

        _UpperCAmelCase						:		Tuple      =				UnCLIPScheduler(
            variance_type="""fixed_small_log"""						,prediction_type="""epsilon"""						,num_train_timesteps=1_000						,)

        _UpperCAmelCase						:		Dict      =				{
            """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 _snake_case				(							self						,a_						,a_=0							)	->     Union[str, Any]:
        if str(a_							).startswith("""mps"""							):
               _UpperCAmelCase						:		Tuple      =				torch.manual_seed(a_							)
        else:
               _UpperCAmelCase						:		Tuple      =				torch.Generator(device=a_							).manual_seed(a_							)
        _UpperCAmelCase						:		List[Any]      =				{
            """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 _snake_case				(							self							)	->     Optional[Any]:
        _UpperCAmelCase						:		Tuple      =				"""cpu"""

        _UpperCAmelCase						:		Dict      =				self.get_dummy_components()

        _UpperCAmelCase						:		str      =				self.pipeline_class(**a_							)
        _UpperCAmelCase						:		Any      =				pipe.to(a_							)

        pipe.set_progress_bar_config(disable=a_							)

        _UpperCAmelCase						:		Optional[int]      =				pipe(**self.get_dummy_inputs(a_							)							)
        _UpperCAmelCase						:		str      =				output.images

        _UpperCAmelCase						:		str      =				pipe(
            **self.get_dummy_inputs(a_							)						,return_dict=a_						,)[0]

        _UpperCAmelCase						:		str      =				image[0, -3:, -3:, -1]
        _UpperCAmelCase						:		str      =				image_from_tuple[0, -3:, -3:, -1]

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

        _UpperCAmelCase						:		Any      =				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 _snake_case				(							self							)	->     Union[str, Any]:
        _UpperCAmelCase						:		Union[str, Any]      =				torch.device("""cpu"""							)



        class 							lowercase :


         """simple docstring"""


         UpperCAmelCase					=						1







        _UpperCAmelCase						:		Any      =				self.get_dummy_components()

        _UpperCAmelCase						:		List[str]      =				self.pipeline_class(**a_							)
        _UpperCAmelCase						:		int      =				pipe.to(a_							)

        _UpperCAmelCase						:		List[Any]      =				components["""prior"""]
        _UpperCAmelCase						:		List[Any]      =				components["""decoder"""]
        _UpperCAmelCase						:		List[str]      =				components["""super_res_first"""]
        _UpperCAmelCase						:		List[Any]      =				components["""tokenizer"""]
        _UpperCAmelCase						:		List[Any]      =				components["""text_encoder"""]

        _UpperCAmelCase						:		Optional[Any]      =				torch.Generator(device=a_							).manual_seed(0							)
        _UpperCAmelCase						:		List[str]      =				prior.dtype
        _UpperCAmelCase						:		Union[str, Any]      =				1

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

        _UpperCAmelCase						:		List[str]      =				(
            batch_size,
            super_res_first.config.in_channels // 2,
            super_res_first.config.sample_size,
            super_res_first.config.sample_size,
        )
        _UpperCAmelCase						:		int      =				pipe.prepare_latents(
            a_						,dtype=a_						,device=a_						,generator=a_						,latents=a_						,scheduler=DummyScheduler()							)

        pipe.set_progress_bar_config(disable=a_							)

        _UpperCAmelCase						:		str      =				"""this is a prompt example"""

        _UpperCAmelCase						:		Optional[Any]      =				torch.Generator(device=a_							).manual_seed(0							)
        _UpperCAmelCase						:		Union[str, Any]      =				pipe(
            [prompt]						,generator=a_						,prior_num_inference_steps=2						,decoder_num_inference_steps=2						,super_res_num_inference_steps=2						,prior_latents=a_						,decoder_latents=a_						,super_res_latents=a_						,output_type="""np"""						,)
        _UpperCAmelCase						:		Optional[int]      =				output.images

        _UpperCAmelCase						:		Any      =				tokenizer(
            a_						,padding="""max_length"""						,max_length=tokenizer.model_max_length						,return_tensors="""pt"""						,)
        _UpperCAmelCase						:		Tuple      =				text_encoder(text_inputs.input_ids							)
        _UpperCAmelCase						:		str      =				text_inputs.attention_mask

        _UpperCAmelCase						:		Any      =				torch.Generator(device=a_							).manual_seed(0							)
        _UpperCAmelCase						:		Optional[int]      =				pipe(
            generator=a_						,prior_num_inference_steps=2						,decoder_num_inference_steps=2						,super_res_num_inference_steps=2						,prior_latents=a_						,decoder_latents=a_						,super_res_latents=a_						,text_model_output=a_						,text_attention_mask=a_						,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 _snake_case				(							self							)	->     List[Any]:
        _UpperCAmelCase						:		Union[str, Any]      =				torch_device == """cpu"""

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







 @skip_mps
 def _snake_case				(							self							)	->     Optional[Any]:
        _UpperCAmelCase						:		Tuple      =				torch_device == """cpu"""
        _UpperCAmelCase						:		Any      =				True
        _UpperCAmelCase						:		Any      =				[
            """prior_num_inference_steps""",
            """decoder_num_inference_steps""",
            """super_res_num_inference_steps""",
        ]

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







 def _snake_case				(							self							)	->     List[Any]:
        _UpperCAmelCase						:		Optional[int]      =				[
            """prior_num_inference_steps""",
            """decoder_num_inference_steps""",
            """super_res_num_inference_steps""",
        ]

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







 @skip_mps
 def _snake_case				(							self							)	->     int:
        return super().test_dict_tuple_outputs_equivalent()







 @skip_mps
 def _snake_case				(							self							)	->     Tuple:
        return super().test_save_load_local()







 @skip_mps
 def _snake_case				(							self							)	->     Any:
        return super().test_save_load_optional_components()







@nightly
class 							lowercase (						unittest.TestCase       ):


 """simple docstring"""


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







 def _snake_case				(							self							)	->     List[str]:
        _UpperCAmelCase						:		str      =				load_numpy(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
            """/unclip/karlo_v1_alpha_horse_cpu.npy"""							)

        _UpperCAmelCase						:		Optional[int]      =				UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""							)
        pipeline.set_progress_bar_config(disable=a_							)

        _UpperCAmelCase						:		List[Any]      =				torch.manual_seed(0							)
        _UpperCAmelCase						:		Union[str, Any]      =				pipeline(
            """horse"""						,num_images_per_prompt=1						,generator=a_						,output_type="""np"""						,)

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

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







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


 """simple docstring"""


 def _snake_case				(							self							)	->     Any:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







 def _snake_case				(							self							)	->     Optional[Any]:
        _UpperCAmelCase						:		List[str]      =				load_numpy(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
            """/unclip/karlo_v1_alpha_horse_fp16.npy"""							)

        _UpperCAmelCase						:		Tuple      =				UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""						,torch_dtype=torch.floataa							)
        _UpperCAmelCase						:		Union[str, Any]      =				pipeline.to(a_							)
        pipeline.set_progress_bar_config(disable=a_							)

        _UpperCAmelCase						:		Dict      =				torch.Generator(device="""cpu"""							).manual_seed(0							)
        _UpperCAmelCase						:		Tuple      =				pipeline(
            """horse"""						,generator=a_						,output_type="""np"""						,)

        _UpperCAmelCase						:		str      =				output.images[0]

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

        assert_mean_pixel_difference(a_						,a_							)







 def _snake_case				(							self							)	->     List[str]:
        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        _UpperCAmelCase						:		Dict      =				UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""						,torch_dtype=torch.floataa							)
        _UpperCAmelCase						:		Optional[Any]      =				pipe.to(a_							)
        pipe.set_progress_bar_config(disable=a_							)
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        _UpperCAmelCase						:		Tuple      =				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"""						,)

        _UpperCAmelCase						:		Union[str, Any]      =				torch.cuda.max_memory_allocated()
        # make sure that less than 7 GB is allocated
        assert mem_bytes < 7 * 10**9








'''simple docstring'''






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


A_							:				List[Any]						=			{
    """configuration_deta""": ["""DETA_PRETRAINED_CONFIG_ARCHIVE_MAP""", """DetaConfig"""],
}

try:
   if not is_vision_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   A_							:				Optional[int]						=			["""DetaImageProcessor"""]

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   A_							:				Optional[Any]						=			[
       """DETA_PRETRAINED_MODEL_ARCHIVE_LIST""",
       """DetaForObjectDetection""",
       """DetaModel""",
       """DetaPreTrainedModel""",
   ]


if TYPE_CHECKING:
   from .configuration_deta import DETA_PRETRAINED_CONFIG_ARCHIVE_MAP, DetaConfig

   try:
      if not is_vision_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .image_processing_deta import DetaImageProcessor

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_deta import (
          DETA_PRETRAINED_MODEL_ARCHIVE_LIST,
          DetaForObjectDetection,
          DetaModel,
          DetaPreTrainedModel,
      )

else:
   import sys

   A_							:				List[Any]						=			_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



