



"""simple docstring"""




import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import (
    DiffusionPipeline,
    UnCLIPImageVariationPipeline,
    UnCLIPScheduler,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, load_image, require_torch_gpu, skip_mps

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()






class      UpperCAmelCase_  (    _lowerCAmelCase	,      unittest.TestCase						):
      __SCREAMING_SNAKE_CASE						:       List[Any]								=			UnCLIPImageVariationPipeline
      __SCREAMING_SNAKE_CASE						:       List[Any]								=			IMAGE_VARIATION_PARAMS - {"height", "width", "guidance_scale"}
      __SCREAMING_SNAKE_CASE						:       int								=			IMAGE_VARIATION_BATCH_PARAMS

      __SCREAMING_SNAKE_CASE						:       List[Any]								=			[
          "generator",
          "return_dict",
          "decoder_num_inference_steps",
          "super_res_num_inference_steps",
      ]
      __SCREAMING_SNAKE_CASE						:       Any								=			False



      @property
      def 						snake_case_ (					self					:     Tuple						):
           return 3_2



      @property
      def 						snake_case_ (					self					:     int						):
           return 3_2



      @property
      def 						snake_case_ (					self					:     int						):
           return self.time_input_dim



      @property
      def 						snake_case_ (					self					:     Dict						):
           return self.time_input_dim * 4



      @property
      def 						snake_case_ (					self					:     List[Any]						):
           return 1_0_0



      @property
      def 						snake_case_ (					self					:     List[str]						):
           _UpperCAmelCase      :     Tuple        =   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						)
           return tokenizer



      @property
      def 						snake_case_ (					self					:     Optional[Any]						):
           torch.manual_seed(0						)
           _UpperCAmelCase      :     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=3_7					,     layer_norm_eps=1e-05					,     num_attention_heads=4					,     num_hidden_layers=5					,     pad_token_id=1					,     vocab_size=1_0_0_0					,     )
           return CLIPTextModelWithProjection(_lowercase						)



      @property
      def 						snake_case_ (					self					:     List[str]						):
           torch.manual_seed(0						)
           _UpperCAmelCase      :     Union[str, Any]        =   CLIPVisionConfig(
               hidden_size=self.text_embedder_hidden_size					,     projection_dim=self.text_embedder_hidden_size					,     num_hidden_layers=5					,     num_attention_heads=4					,     image_size=3_2					,     intermediate_size=3_7					,     patch_size=1					,     )
           return CLIPVisionModelWithProjection(_lowercase						)



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

           _UpperCAmelCase      :     Optional[Any]        =   {
               "clip_embeddings_dim": self.text_embedder_hidden_size,
               "time_embed_dim": self.time_embed_dim,
               "cross_attention_dim": self.cross_attention_dim,
           }

           _UpperCAmelCase      :     int        =   UnCLIPTextProjModel(**_lowercase						)
           return model



      @property
      def 						snake_case_ (					self					:     Tuple						):
           torch.manual_seed(0						)

           _UpperCAmelCase      :     int        =   {
               "sample_size": 3_2,
               # 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      :     Any        =   UNetaDConditionModel(**_lowercase						)
           return model



      @property
      def 						snake_case_ (					self					:     Dict						):
           return {
               "sample_size": 6_4,
               "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					:     Union[str, Any]						):
           torch.manual_seed(0						)

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



      @property
      def 						snake_case_ (					self					:     List[str]						):
           # seeded differently to get different unet than `self.dummy_super_res_first`
           torch.manual_seed(1						)

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



      def 						snake_case_ (					self					:     Any						):
           _UpperCAmelCase      :     Dict        =   self.dummy_decoder
           _UpperCAmelCase      :     List[Any]        =   self.dummy_text_proj
           _UpperCAmelCase      :     Dict        =   self.dummy_text_encoder
           _UpperCAmelCase      :     Dict        =   self.dummy_tokenizer
           _UpperCAmelCase      :     Any        =   self.dummy_super_res_first
           _UpperCAmelCase      :     Union[str, Any]        =   self.dummy_super_res_last

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

           _UpperCAmelCase      :     Dict        =   UnCLIPScheduler(
               variance_type="fixed_small_log"					,     prediction_type="epsilon"					,     num_train_timesteps=1_0_0_0					,     )

           _UpperCAmelCase      :     List[str]        =   CLIPImageProcessor(crop_size=3_2					,     size=3_2						)

           _UpperCAmelCase      :     str        =   self.dummy_image_encoder

           return {
               "decoder": decoder,
               "text_encoder": text_encoder,
               "tokenizer": tokenizer,
               "text_proj": text_proj,
               "feature_extractor": feature_extractor,
               "image_encoder": image_encoder,
               "super_res_first": super_res_first,
               "super_res_last": super_res_last,
               "decoder_scheduler": decoder_scheduler,
               "super_res_scheduler": super_res_scheduler,
           }



      def 						snake_case_ (					self					:     Any					,     A					:     Optional[Any]					,     A					:     int=0					,     A					:     List[Any]=True						):
           _UpperCAmelCase      :     List[Any]        =   floats_tensor((1, 3, 3_2, 3_2)					,     rng=random.Random(_lowercase						)						).to(_lowercase						)
           if str(_lowercase						).startswith("mps"						):
                _UpperCAmelCase      :     Any        =   torch.manual_seed(_lowercase						)
           else:
                _UpperCAmelCase      :     Optional[Any]        =   torch.Generator(device=_lowercase						).manual_seed(_lowercase						)

           if pil_image:
                _UpperCAmelCase      :     Any        =   input_image * 0.5 + 0.5
                _UpperCAmelCase      :     Union[str, Any]        =   input_image.clamp(0					,     1						)
                _UpperCAmelCase      :     Any        =   input_image.cpu().permute(0					,     2					,     3					,     1						).float().numpy()
                _UpperCAmelCase      :     List[str]        =   DiffusionPipeline.numpy_to_pil(_lowercase						)[0]

           return {
               "image": input_image,
               "generator": generator,
               "decoder_num_inference_steps": 2,
               "super_res_num_inference_steps": 2,
               "output_type": "np",
           }



      def 						snake_case_ (					self					:     Dict						):
           _UpperCAmelCase      :     Optional[Any]        =   "cpu"

           _UpperCAmelCase      :     int        =   self.get_dummy_components()

           _UpperCAmelCase      :     Any        =   self.pipeline_class(**_lowercase						)
           _UpperCAmelCase      :     Optional[int]        =   pipe.to(_lowercase						)

           pipe.set_progress_bar_config(disable=_lowercase						)

           _UpperCAmelCase      :     int        =   self.get_dummy_inputs(_lowercase					,     pil_image=_lowercase						)

           _UpperCAmelCase      :     List[Any]        =   pipe(**_lowercase						)
           _UpperCAmelCase      :     int        =   output.images

           _UpperCAmelCase      :     Tuple        =   self.get_dummy_inputs(_lowercase					,     pil_image=_lowercase						)

           _UpperCAmelCase      :     Dict        =   pipe(
               **_lowercase					,     return_dict=_lowercase					,     )[0]

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

           assert image.shape == (1, 6_4, 6_4, 3)

           _UpperCAmelCase      :     Optional[Any]        =   np.array(
               [
                   0.9_997,
                   0.0_002,
                   0.9_997,
                   0.9_997,
                   0.9_969,
                   0.0_023,
                   0.9_997,
                   0.9_969,
                   0.9_970,
               ]						)

           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					:     int						):
           _UpperCAmelCase      :     str        =   "cpu"

           _UpperCAmelCase      :     Optional[Any]        =   self.get_dummy_components()

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

           pipe.set_progress_bar_config(disable=_lowercase						)

           _UpperCAmelCase      :     str        =   self.get_dummy_inputs(_lowercase					,     pil_image=_lowercase						)

           _UpperCAmelCase      :     Tuple        =   pipe(**_lowercase						)
           _UpperCAmelCase      :     str        =   output.images

           _UpperCAmelCase      :     List[Any]        =   self.get_dummy_inputs(_lowercase					,     pil_image=_lowercase						)

           _UpperCAmelCase      :     str        =   pipe(
               **_lowercase					,     return_dict=_lowercase					,     )[0]

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

           assert image.shape == (1, 6_4, 6_4, 3)

           _UpperCAmelCase      :     List[str]        =   np.array([0.9_997, 0.0_003, 0.9_997, 0.9_997, 0.9_970, 0.0_024, 0.9_997, 0.9_971, 0.9_971]						)

           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					:     List[Any]						):
           _UpperCAmelCase      :     List[str]        =   "cpu"

           _UpperCAmelCase      :     Optional[Any]        =   self.get_dummy_components()

           _UpperCAmelCase      :     Dict        =   self.pipeline_class(**_lowercase						)
           _UpperCAmelCase      :     Tuple        =   pipe.to(_lowercase						)

           pipe.set_progress_bar_config(disable=_lowercase						)

           _UpperCAmelCase      :     Optional[int]        =   self.get_dummy_inputs(_lowercase					,     pil_image=_lowercase						)
           _UpperCAmelCase      :     Tuple        =   [
               pipeline_inputs["image"],
               pipeline_inputs["image"],
           ]

           _UpperCAmelCase      :     str        =   pipe(**_lowercase						)
           _UpperCAmelCase      :     Optional[Any]        =   output.images

           _UpperCAmelCase      :     Tuple        =   self.get_dummy_inputs(_lowercase					,     pil_image=_lowercase						)
           _UpperCAmelCase      :     Any        =   [
               tuple_pipeline_inputs["image"],
               tuple_pipeline_inputs["image"],
           ]

           _UpperCAmelCase      :     Dict        =   pipe(
               **_lowercase					,     return_dict=_lowercase					,     )[0]

           _UpperCAmelCase      :     Optional[Any]        =   image[0, -3:, -3:, -1]
           _UpperCAmelCase      :     Tuple        =   image_from_tuple[0, -3:, -3:, -1]

           assert image.shape == (2, 6_4, 6_4, 3)

           _UpperCAmelCase      :     Dict        =   np.array(
               [
                   0.9_997,
                   0.9_989,
                   0.0_008,
                   0.0_021,
                   0.9_960,
                   0.0_018,
                   0.0_014,
                   0.0_002,
                   0.9_933,
               ]						)

           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					:     Dict						):
           _UpperCAmelCase      :     Dict        =   torch.device("cpu"						)






           class      UpperCAmelCase_  :
                 __SCREAMING_SNAKE_CASE						:       Union[str, Any]								=			1







           _UpperCAmelCase      :     Tuple        =   self.get_dummy_components()

           _UpperCAmelCase      :     Optional[int]        =   self.pipeline_class(**_lowercase						)
           _UpperCAmelCase      :     Optional[Any]        =   pipe.to(_lowercase						)

           pipe.set_progress_bar_config(disable=_lowercase						)

           _UpperCAmelCase      :     List[str]        =   torch.Generator(device=_lowercase						).manual_seed(0						)
           _UpperCAmelCase      :     List[str]        =   pipe.decoder.dtype
           _UpperCAmelCase      :     Dict        =   1

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

           _UpperCAmelCase      :     str        =   (
               batch_size,
               pipe.super_res_first.config.in_channels // 2,
               pipe.super_res_first.config.sample_size,
               pipe.super_res_first.config.sample_size,
           )
           _UpperCAmelCase      :     str        =   pipe.prepare_latents(
               _lowercase					,     dtype=_lowercase					,     device=_lowercase					,     generator=_lowercase					,     latents=_lowercase					,     scheduler=DummyScheduler()						)

           _UpperCAmelCase      :     Union[str, Any]        =   self.get_dummy_inputs(_lowercase					,     pil_image=_lowercase						)

           _UpperCAmelCase      :     List[Any]        =   pipe(
               **_lowercase					,     decoder_latents=_lowercase					,     super_res_latents=_lowercase						).images

           _UpperCAmelCase      :     Union[str, Any]        =   self.get_dummy_inputs(_lowercase					,     pil_image=_lowercase						)
           # Don't pass image, instead pass embedding
           _UpperCAmelCase      :     List[Any]        =   pipeline_inputs.pop("image"						)
           _UpperCAmelCase      :     List[Any]        =   pipe.image_encoder(_lowercase						).image_embeds

           _UpperCAmelCase      :     Dict        =   pipe(
               **_lowercase					,     decoder_latents=_lowercase					,     super_res_latents=_lowercase					,     image_embeddings=_lowercase					,     ).images

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



      @skip_mps
      def 						snake_case_ (					self					:     Dict						):
           _UpperCAmelCase      :     Optional[Any]        =   torch_device == "cpu"

           # Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
           _UpperCAmelCase      :     List[str]        =   1e-2

           self._test_attention_slicing_forward_pass(
               test_max_difference=_lowercase					,     expected_max_diff=_lowercase						)



      @skip_mps
      def 						snake_case_ (					self					:     Union[str, Any]						):
           _UpperCAmelCase      :     Dict        =   torch_device == "cpu"
           _UpperCAmelCase      :     Optional[Any]        =   True
           _UpperCAmelCase      :     int        =   [
               "decoder_num_inference_steps",
               "super_res_num_inference_steps",
           ]

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



      def 						snake_case_ (					self					:     Optional[int]						):
           _UpperCAmelCase      :     List[str]        =   [
               "decoder_num_inference_steps",
               "super_res_num_inference_steps",
           ]

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



      @skip_mps
      def 						snake_case_ (					self					:     List[Any]						):
           return super().test_dict_tuple_outputs_equivalent()



      @skip_mps
      def 						snake_case_ (					self					:     str						):
           return super().test_save_load_local()



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







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



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



      def 						snake_case_ (					self					:     Any						):
           _UpperCAmelCase      :     Optional[Any]        =   load_image(
               "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png"						)
           _UpperCAmelCase      :     List[Any]        =   load_numpy(
               "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
               "/unclip/karlo_v1_alpha_cat_variation_fp16.npy"						)

           _UpperCAmelCase      :     Union[str, Any]        =   UnCLIPImageVariationPipeline.from_pretrained(
               "kakaobrain/karlo-v1-alpha-image-variations"					,     torch_dtype=torch.floataa						)
           _UpperCAmelCase      :     Dict        =   pipeline.to(_lowercase						)
           pipeline.set_progress_bar_config(disable=_lowercase						)

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

           _UpperCAmelCase      :     Any        =   output.images[0]

           assert image.shape == (2_5_6, 2_5_6, 3)

           assert_mean_pixel_difference(_lowercase					,     _lowercase					,     1_5						)





"""simple docstring"""




from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available


_lowerCAmelCase    : Union[str, Any]							=							{}

try:
 if not is_sentencepiece_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _lowerCAmelCase    : Optional[Any]							=							["GPTSw3Tokenizer"]


if TYPE_CHECKING:
 try:
  if not is_sentencepiece_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_gpt_swa import GPTSwaTokenizer

else:
 import sys

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