


'''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 			lowerCamelCase_   (	__a     ,       unittest.TestCase      ):
      lowerCAmelCase__		     =     UnCLIPImageVariationPipeline
      lowerCAmelCase__		     =     IMAGE_VARIATION_PARAMS - {'height', 'width', 'guidance_scale'}
      lowerCAmelCase__		     =     IMAGE_VARIATION_BATCH_PARAMS

      lowerCAmelCase__		     =     [
          'generator',
          'return_dict',
          'decoder_num_inference_steps',
          'super_res_num_inference_steps',
      ]
      lowerCAmelCase__		     =     False






      @property
      def lowercase_							(					self				:    List[Any] ):


          '''simple docstring'''


          return 32






      @property
      def lowercase_							(					self				:    Union[str, Any] ):


          '''simple docstring'''


          return 32






      @property
      def lowercase_							(					self				:    List[Any] ):


          '''simple docstring'''


          return self.time_input_dim






      @property
      def lowercase_							(					self				:    Tuple ):


          '''simple docstring'''


          return self.time_input_dim * 4






      @property
      def lowercase_							(					self				:    Optional[int] ):


          '''simple docstring'''


          return 100






      @property
      def lowercase_							(					self				:    List[str] ):


          '''simple docstring'''


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






      @property
      def lowercase_							(					self				:    Tuple ):


          '''simple docstring'''


          torch.manual_seed(0 )
          UpperCAmelCase__			:		int													=				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=1_000						,						)
          return CLIPTextModelWithProjection(_A )






      @property
      def lowercase_							(					self				:    Optional[int] ):


          '''simple docstring'''


          torch.manual_seed(0 )
          UpperCAmelCase__			:		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=32						,						intermediate_size=37						,						patch_size=1						,						)
          return CLIPVisionModelWithProjection(_A )






      @property
      def lowercase_							(					self				:    Tuple ):


          '''simple docstring'''


          torch.manual_seed(0 )

          UpperCAmelCase__			:		Optional[int]													=				{
              '''clip_embeddings_dim''': self.text_embedder_hidden_size,
              '''time_embed_dim''': self.time_embed_dim,
              '''cross_attention_dim''': self.cross_attention_dim,
          }

          UpperCAmelCase__			:		Any													=				UnCLIPTextProjModel(**_A )
          return model






      @property
      def lowercase_							(					self				:    Tuple ):


          '''simple docstring'''


          torch.manual_seed(0 )

          UpperCAmelCase__			:		str													=				{
              '''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__			:		int													=				UNetaDConditionModel(**_A )
          return model






      @property
      def lowercase_							(					self				:    Union[str, 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 lowercase_							(					self				:    Any ):


          '''simple docstring'''


          torch.manual_seed(0 )

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






      @property
      def lowercase_							(					self				:    Any ):


          '''simple docstring'''


          torch.manual_seed(1 )

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






      def lowercase_							(					self				:    Tuple ):


          '''simple docstring'''


          UpperCAmelCase__			:		Any													=				self.dummy_decoder
          UpperCAmelCase__			:		Any													=				self.dummy_text_proj
          UpperCAmelCase__			:		int													=				self.dummy_text_encoder
          UpperCAmelCase__			:		Tuple													=				self.dummy_tokenizer
          UpperCAmelCase__			:		Optional[Any]													=				self.dummy_super_res_first
          UpperCAmelCase__			:		Dict													=				self.dummy_super_res_last

          UpperCAmelCase__			:		int													=				UnCLIPScheduler(
              variance_type='''learned_range'''						,						prediction_type='''epsilon'''						,						num_train_timesteps=1_000						,						)

          UpperCAmelCase__			:		Union[str, Any]													=				UnCLIPScheduler(
              variance_type='''fixed_small_log'''						,						prediction_type='''epsilon'''						,						num_train_timesteps=1_000						,						)

          UpperCAmelCase__			:		List[str]													=				CLIPImageProcessor(crop_size=32						,						size=32 )

          UpperCAmelCase__			:		Tuple													=				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 lowercase_							(					self				:    Tuple						,						_A				:    Tuple						,						_A				:    Dict=0						,						_A				:    Optional[Any]=True ):


          '''simple docstring'''


          UpperCAmelCase__			:		List[Any]													=				floats_tensor((1, 3, 32, 32)						,						rng=random.Random(_A ) ).to(_A )
          if str(_A ).startswith('''mps''' ):
              UpperCAmelCase__			:		int													=				torch.manual_seed(_A )
          else:
              UpperCAmelCase__			:		List[str]													=				torch.Generator(device=_A ).manual_seed(_A )

          if pil_image:
              UpperCAmelCase__			:		Dict													=				input_image * 0.5 + 0.5
              UpperCAmelCase__			:		Optional[int]													=				input_image.clamp(0						,						1 )
              UpperCAmelCase__			:		Any													=				input_image.cpu().permute(0						,						2						,						3						,						1 ).float().numpy()
              UpperCAmelCase__			:		str													=				DiffusionPipeline.numpy_to_pil(_A )[0]

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






      def lowercase_							(					self				:    Tuple ):


          '''simple docstring'''


          UpperCAmelCase__			:		Union[str, Any]													=				'''cpu'''

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

          UpperCAmelCase__			:		List[Any]													=				self.pipeline_class(**_A )
          UpperCAmelCase__			:		List[str]													=				pipe.to(_A )

          pipe.set_progress_bar_config(disable=_A )

          UpperCAmelCase__			:		str													=				self.get_dummy_inputs(_A						,						pil_image=_A )

          UpperCAmelCase__			:		Optional[int]													=				pipe(**_A )
          UpperCAmelCase__			:		Optional[int]													=				output.images

          UpperCAmelCase__			:		int													=				self.get_dummy_inputs(_A						,						pil_image=_A )

          UpperCAmelCase__			:		Any													=				pipe(
              **_A						,						return_dict=_A						,						)[0]

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

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

          UpperCAmelCase__			:		Optional[int]													=				np.array(
              [
                  0.9_9_9_7,
                  0.0_0_0_2,
                  0.9_9_9_7,
                  0.9_9_9_7,
                  0.9_9_6_9,
                  0.0_0_2_3,
                  0.9_9_9_7,
                  0.9_9_6_9,
                  0.9_9_7_0,
              ] )

          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 lowercase_							(					self				:    Tuple ):


          '''simple docstring'''


          UpperCAmelCase__			:		Tuple													=				'''cpu'''

          UpperCAmelCase__			:		int													=				self.get_dummy_components()

          UpperCAmelCase__			:		Optional[int]													=				self.pipeline_class(**_A )
          UpperCAmelCase__			:		Any													=				pipe.to(_A )

          pipe.set_progress_bar_config(disable=_A )

          UpperCAmelCase__			:		List[Any]													=				self.get_dummy_inputs(_A						,						pil_image=_A )

          UpperCAmelCase__			:		Union[str, Any]													=				pipe(**_A )
          UpperCAmelCase__			:		Optional[Any]													=				output.images

          UpperCAmelCase__			:		Dict													=				self.get_dummy_inputs(_A						,						pil_image=_A )

          UpperCAmelCase__			:		Dict													=				pipe(
              **_A						,						return_dict=_A						,						)[0]

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

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

          UpperCAmelCase__			:		List[str]													=				np.array([0.9_9_9_7, 0.0_0_0_3, 0.9_9_9_7, 0.9_9_9_7, 0.9_9_7_0, 0.0_0_2_4, 0.9_9_9_7, 0.9_9_7_1, 0.9_9_7_1] )

          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 lowercase_							(					self				:    Optional[Any] ):


          '''simple docstring'''


          UpperCAmelCase__			:		Tuple													=				'''cpu'''

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

          UpperCAmelCase__			:		Dict													=				self.pipeline_class(**_A )
          UpperCAmelCase__			:		Optional[Any]													=				pipe.to(_A )

          pipe.set_progress_bar_config(disable=_A )

          UpperCAmelCase__			:		Optional[Any]													=				self.get_dummy_inputs(_A						,						pil_image=_A )
          UpperCAmelCase__			:		Union[str, Any]													=				[
              pipeline_inputs['''image'''],
              pipeline_inputs['''image'''],
          ]

          UpperCAmelCase__			:		Any													=				pipe(**_A )
          UpperCAmelCase__			:		List[str]													=				output.images

          UpperCAmelCase__			:		Any													=				self.get_dummy_inputs(_A						,						pil_image=_A )
          UpperCAmelCase__			:		List[str]													=				[
              tuple_pipeline_inputs['''image'''],
              tuple_pipeline_inputs['''image'''],
          ]

          UpperCAmelCase__			:		List[Any]													=				pipe(
              **_A						,						return_dict=_A						,						)[0]

          UpperCAmelCase__			:		Union[str, Any]													=				image[0, -3:, -3:, -1]
          UpperCAmelCase__			:		Optional[int]													=				image_from_tuple[0, -3:, -3:, -1]

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

          UpperCAmelCase__			:		List[Any]													=				np.array(
              [
                  0.9_9_9_7,
                  0.9_9_8_9,
                  0.0_0_0_8,
                  0.0_0_2_1,
                  0.9_9_6_0,
                  0.0_0_1_8,
                  0.0_0_1_4,
                  0.0_0_0_2,
                  0.9_9_3_3,
              ] )

          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 lowercase_							(					self				:    str ):


          '''simple docstring'''


          UpperCAmelCase__			:		str													=				torch.device('''cpu''' )




          class 			lowerCamelCase_   :
                lowerCAmelCase__		     =     1

          UpperCAmelCase__			:		Union[str, Any]													=				self.get_dummy_components()

          UpperCAmelCase__			:		str													=				self.pipeline_class(**_A )
          UpperCAmelCase__			:		int													=				pipe.to(_A )

          pipe.set_progress_bar_config(disable=_A )

          UpperCAmelCase__			:		Optional[Any]													=				torch.Generator(device=_A ).manual_seed(0 )
          UpperCAmelCase__			:		Optional[Any]													=				pipe.decoder.dtype
          UpperCAmelCase__			:		Any													=				1

          UpperCAmelCase__			:		Tuple													=				(
              batch_size,
              pipe.decoder.config.in_channels,
              pipe.decoder.config.sample_size,
              pipe.decoder.config.sample_size,
          )
          UpperCAmelCase__			:		Any													=				pipe.prepare_latents(
              _A						,						dtype=_A						,						device=_A						,						generator=_A						,						latents=_A						,						scheduler=DummyScheduler() )

          UpperCAmelCase__			:		List[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__			:		Optional[Any]													=				pipe.prepare_latents(
              _A						,						dtype=_A						,						device=_A						,						generator=_A						,						latents=_A						,						scheduler=DummyScheduler() )

          UpperCAmelCase__			:		Optional[Any]													=				self.get_dummy_inputs(_A						,						pil_image=_A )

          UpperCAmelCase__			:		List[Any]													=				pipe(
              **_A						,						decoder_latents=_A						,						super_res_latents=_A ).images

          UpperCAmelCase__			:		List[Any]													=				self.get_dummy_inputs(_A						,						pil_image=_A )
          # Don't pass image, instead pass embedding
          UpperCAmelCase__			:		Optional[int]													=				pipeline_inputs.pop('''image''' )
          UpperCAmelCase__			:		str													=				pipe.image_encoder(_A ).image_embeds

          UpperCAmelCase__			:		Dict													=				pipe(
              **_A						,						decoder_latents=_A						,						super_res_latents=_A						,						image_embeddings=_A						,						).images

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






      @skip_mps
      def lowercase_							(					self				:    Union[str, Any] ):


          '''simple docstring'''


          UpperCAmelCase__			:		Any													=				torch_device == '''cpu'''

          # Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
          UpperCAmelCase__			:		Dict													=				1e-2

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






      @skip_mps
      def lowercase_							(					self				:    Optional[int] ):


          '''simple docstring'''


          UpperCAmelCase__			:		Optional[int]													=				torch_device == '''cpu'''
          UpperCAmelCase__			:		Dict													=				True
          UpperCAmelCase__			:		List[str]													=				[
              '''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 lowercase_							(					self				:    Union[str, Any] ):


          '''simple docstring'''


          UpperCAmelCase__			:		int													=				[
              '''decoder_num_inference_steps''',
              '''super_res_num_inference_steps''',
          ]

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


          '''simple docstring'''


          return super().test_dict_tuple_outputs_equivalent()






      @skip_mps
      def lowercase_							(					self				:    int ):


          '''simple docstring'''


          return super().test_save_load_local()






      @skip_mps
      def lowercase_							(					self				:    Optional[Any] ):


          '''simple docstring'''


          return super().test_save_load_optional_components()




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






      def lowercase_							(					self				:    int ):


          '''simple docstring'''


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






      def lowercase_							(					self				:    Union[str, Any] ):


          '''simple docstring'''


          UpperCAmelCase__			:		Any													=				load_image(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png''' )
          UpperCAmelCase__			:		int													=				load_numpy(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
              '''/unclip/karlo_v1_alpha_cat_variation_fp16.npy''' )

          UpperCAmelCase__			:		Optional[int]													=				UnCLIPImageVariationPipeline.from_pretrained(
              '''kakaobrain/karlo-v1-alpha-image-variations'''						,						torch_dtype=torch.floataa )
          UpperCAmelCase__			:		Optional[Any]													=				pipeline.to(_A )
          pipeline.set_progress_bar_config(disable=_A )

          UpperCAmelCase__			:		List[str]													=				torch.Generator(device='''cpu''' ).manual_seed(0 )
          UpperCAmelCase__			:		Tuple													=				pipeline(
              _A						,						generator=_A						,						output_type='''np'''						,						)

          UpperCAmelCase__			:		List[Any]													=				output.images[0]

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

          assert_mean_pixel_difference(_A						,						_A						,						15 )




'''simple docstring'''
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCamelCase__    					=      {
    '''configuration_rembert''': ['''REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''RemBertConfig''', '''RemBertOnnxConfig''']
}

try:
  if not is_sentencepiece_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      ['''RemBertTokenizer''']

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      ['''RemBertTokenizerFast''']

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      [
      '''REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''RemBertForCausalLM''',
      '''RemBertForMaskedLM''',
      '''RemBertForMultipleChoice''',
      '''RemBertForQuestionAnswering''',
      '''RemBertForSequenceClassification''',
      '''RemBertForTokenClassification''',
      '''RemBertLayer''',
      '''RemBertModel''',
      '''RemBertPreTrainedModel''',
      '''load_tf_weights_in_rembert''',
  ]


try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      [
      '''TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''TFRemBertForCausalLM''',
      '''TFRemBertForMaskedLM''',
      '''TFRemBertForMultipleChoice''',
      '''TFRemBertForQuestionAnswering''',
      '''TFRemBertForSequenceClassification''',
      '''TFRemBertForTokenClassification''',
      '''TFRemBertLayer''',
      '''TFRemBertModel''',
      '''TFRemBertPreTrainedModel''',
  ]


if TYPE_CHECKING:
  from .configuration_rembert import REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RemBertConfig, RemBertOnnxConfig

  try:
    if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_rembert import RemBertTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_rembert_fast import RemBertTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_rembert import (
        REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        RemBertForCausalLM,
        RemBertForMaskedLM,
        RemBertForMultipleChoice,
        RemBertForQuestionAnswering,
        RemBertForSequenceClassification,
        RemBertForTokenClassification,
        RemBertLayer,
        RemBertModel,
        RemBertPreTrainedModel,
        load_tf_weights_in_rembert,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_rembert import (
        TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFRemBertForCausalLM,
        TFRemBertForMaskedLM,
        TFRemBertForMultipleChoice,
        TFRemBertForQuestionAnswering,
        TFRemBertForSequenceClassification,
        TFRemBertForTokenClassification,
        TFRemBertLayer,
        TFRemBertModel,
        TFRemBertPreTrainedModel,
    )


else:
  import sys

  UpperCamelCase__    					=      _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)
