

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   __lowercase      (_UpperCAmelCase		,   unittest.TestCase ):
 _UpperCamelCase			  =  UnCLIPImageVariationPipeline
 _UpperCamelCase			  =  IMAGE_VARIATION_PARAMS - {"""height""", """width""", """guidance_scale"""}
 _UpperCamelCase			  =  IMAGE_VARIATION_BATCH_PARAMS

 _UpperCamelCase			  =  [
     """generator""",
     """return_dict""",
     """decoder_num_inference_steps""",
     """super_res_num_inference_steps""",
 ]
 _UpperCamelCase			  =  False
 @property
 def 	UpperCamelCase__  (							self					)					->Union[str, Any]:



       '''simple docstring'''





       return 32
 @property
 def 	UpperCamelCase__  (							self					)					->Union[str, Any]:



       '''simple docstring'''





       return 32
 @property
 def 	UpperCamelCase__  (							self					)					->Tuple:



       '''simple docstring'''





       return self.time_input_dim
 @property
 def 	UpperCamelCase__  (							self					)					->Tuple:



       '''simple docstring'''





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



       '''simple docstring'''





       return 100
 @property
 def 	UpperCamelCase__  (							self					)					->int:



       '''simple docstring'''





       __lowerCAmelCase					:						List[Any]					=					CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)
       return tokenizer
 @property
 def 	UpperCamelCase__  (							self					)					->Union[str, Any]:



       '''simple docstring'''





       torch.manual_seed(0					)
       __lowerCAmelCase					:						Optional[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=1000      ,							)
       return CLIPTextModelWithProjection(A_					)
 @property
 def 	UpperCamelCase__  (							self					)					->str:



       '''simple docstring'''





       torch.manual_seed(0					)
       __lowerCAmelCase					:						Optional[int]					=					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 	UpperCamelCase__  (							self					)					->str:



       '''simple docstring'''





       torch.manual_seed(0					)

       __lowerCAmelCase					:						Optional[Any]					=					{
           '''clip_embeddings_dim''': self.text_embedder_hidden_size,
           '''time_embed_dim''': self.time_embed_dim,
           '''cross_attention_dim''': self.cross_attention_dim,
       }

       __lowerCAmelCase					:						Any					=					UnCLIPTextProjModel(**A_					)
       return model
 @property
 def 	UpperCamelCase__  (							self					)					->Dict:



       '''simple docstring'''





       torch.manual_seed(0					)

       __lowerCAmelCase					:						int					=					{
           '''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					:						Any					=					UNetaDConditionModel(**A_					)
       return model
 @property
 def 	UpperCamelCase__  (							self					)					->Optional[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 	UpperCamelCase__  (							self					)					->str:



       '''simple docstring'''





       torch.manual_seed(0					)

       __lowerCAmelCase					:						str					=					UNetaDModel(**self.dummy_super_res_kwargs					)
       return model
 @property
 def 	UpperCamelCase__  (							self					)					->List[str]:



       '''simple docstring'''





       torch.manual_seed(1					)

       __lowerCAmelCase					:						List[Any]					=					UNetaDModel(**self.dummy_super_res_kwargs					)
       return model
 def 	UpperCamelCase__  (							self					)					->Optional[int]:



       '''simple docstring'''





       __lowerCAmelCase					:						Optional[Any]					=					self.dummy_decoder
       __lowerCAmelCase					:						Optional[Any]					=					self.dummy_text_proj
       __lowerCAmelCase					:						Dict					=					self.dummy_text_encoder
       __lowerCAmelCase					:						Union[str, Any]					=					self.dummy_tokenizer
       __lowerCAmelCase					:						Optional[int]					=					self.dummy_super_res_first
       __lowerCAmelCase					:						List[Any]					=					self.dummy_super_res_last

       __lowerCAmelCase					:						List[Any]					=					UnCLIPScheduler(
           variance_type='''learned_range'''      ,							prediction_type='''epsilon'''      ,							num_train_timesteps=1000      ,							)

       __lowerCAmelCase					:						Optional[int]					=					UnCLIPScheduler(
           variance_type='''fixed_small_log'''      ,							prediction_type='''epsilon'''      ,							num_train_timesteps=1000      ,							)

       __lowerCAmelCase					:						Optional[Any]					=					CLIPImageProcessor(crop_size=32      ,							size=32					)

       __lowerCAmelCase					:						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 	UpperCamelCase__  (							self      ,							A_      ,							A_=0      ,							A_=True					)					->Union[str, Any]:



       '''simple docstring'''





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

       if pil_image:
             __lowerCAmelCase					:						Union[str, Any]					=					input_image * 0.5 + 0.5
             __lowerCAmelCase					:						Dict					=					input_image.clamp(0      ,							1					)
             __lowerCAmelCase					:						Any					=					input_image.cpu().permute(0      ,							2      ,							3      ,							1					).float().numpy()
             __lowerCAmelCase					:						List[Any]					=					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 	UpperCamelCase__  (							self					)					->Any:



       '''simple docstring'''





       __lowerCAmelCase					:						str					=					'''cpu'''

       __lowerCAmelCase					:						List[Any]					=					self.get_dummy_components()

       __lowerCAmelCase					:						List[Any]					=					self.pipeline_class(**A_					)
       __lowerCAmelCase					:						List[str]					=					pipe.to(A_					)

       pipe.set_progress_bar_config(disable=A_					)

       __lowerCAmelCase					:						List[str]					=					self.get_dummy_inputs(A_      ,							pil_image=A_					)

       __lowerCAmelCase					:						Dict					=					pipe(**A_					)
       __lowerCAmelCase					:						List[Any]					=					output.images

       __lowerCAmelCase					:						int					=					self.get_dummy_inputs(A_      ,							pil_image=A_					)

       __lowerCAmelCase					:						List[str]					=					pipe(
           **A_      ,							return_dict=A_      ,							)[0]

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

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

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



       '''simple docstring'''





       __lowerCAmelCase					:						int					=					'''cpu'''

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

       __lowerCAmelCase					:						List[str]					=					self.pipeline_class(**A_					)
       __lowerCAmelCase					:						Optional[int]					=					pipe.to(A_					)

       pipe.set_progress_bar_config(disable=A_					)

       __lowerCAmelCase					:						Optional[int]					=					self.get_dummy_inputs(A_      ,							pil_image=A_					)

       __lowerCAmelCase					:						Union[str, Any]					=					pipe(**A_					)
       __lowerCAmelCase					:						str					=					output.images

       __lowerCAmelCase					:						Optional[Any]					=					self.get_dummy_inputs(A_      ,							pil_image=A_					)

       __lowerCAmelCase					:						Optional[Any]					=					pipe(
           **A_      ,							return_dict=A_      ,							)[0]

       __lowerCAmelCase					:						Optional[int]					=					image[0, -3:, -3:, -1]
       __lowerCAmelCase					:						List[Any]					=					image_from_tuple[0, -3:, -3:, -1]

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

       __lowerCAmelCase					:						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 	UpperCamelCase__  (							self					)					->List[Any]:



       '''simple docstring'''





       __lowerCAmelCase					:						Optional[int]					=					'''cpu'''

       __lowerCAmelCase					:						int					=					self.get_dummy_components()

       __lowerCAmelCase					:						int					=					self.pipeline_class(**A_					)
       __lowerCAmelCase					:						int					=					pipe.to(A_					)

       pipe.set_progress_bar_config(disable=A_					)

       __lowerCAmelCase					:						str					=					self.get_dummy_inputs(A_      ,							pil_image=A_					)
       __lowerCAmelCase					:						Optional[int]					=					[
           pipeline_inputs['''image'''],
           pipeline_inputs['''image'''],
       ]

       __lowerCAmelCase					:						Dict					=					pipe(**A_					)
       __lowerCAmelCase					:						Dict					=					output.images

       __lowerCAmelCase					:						Optional[Any]					=					self.get_dummy_inputs(A_      ,							pil_image=A_					)
       __lowerCAmelCase					:						int					=					[
           tuple_pipeline_inputs['''image'''],
           tuple_pipeline_inputs['''image'''],
       ]

       __lowerCAmelCase					:						Dict					=					pipe(
           **A_      ,							return_dict=A_      ,							)[0]

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

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

       __lowerCAmelCase					:						Any					=					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 	UpperCamelCase__  (							self					)					->Optional[int]:



       '''simple docstring'''





       __lowerCAmelCase					:						List[str]					=					torch.device('''cpu'''					)







       class   __lowercase      :
        _UpperCamelCase			  =  1



       __lowerCAmelCase					:						List[Any]					=					self.get_dummy_components()

       __lowerCAmelCase					:						Optional[int]					=					self.pipeline_class(**A_					)
       __lowerCAmelCase					:						Union[str, Any]					=					pipe.to(A_					)

       pipe.set_progress_bar_config(disable=A_					)

       __lowerCAmelCase					:						str					=					torch.Generator(device=A_					).manual_seed(0					)
       __lowerCAmelCase					:						List[str]					=					pipe.decoder.dtype
       __lowerCAmelCase					:						List[Any]					=					1

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

       __lowerCAmelCase					:						Optional[int]					=					(
           batch_size,
           pipe.super_res_first.config.in_channels // 2,
           pipe.super_res_first.config.sample_size,
           pipe.super_res_first.config.sample_size,
       )
       __lowerCAmelCase					:						int					=					pipe.prepare_latents(
           A_      ,							dtype=A_      ,							device=A_      ,							generator=A_      ,							latents=A_      ,							scheduler=DummyScheduler()					)

       __lowerCAmelCase					:						Tuple					=					self.get_dummy_inputs(A_      ,							pil_image=A_					)

       __lowerCAmelCase					:						Dict					=					pipe(
           **A_      ,							decoder_latents=A_      ,							super_res_latents=A_					).images

       __lowerCAmelCase					:						Any					=					self.get_dummy_inputs(A_      ,							pil_image=A_					)
       # Don't pass image, instead pass embedding
       __lowerCAmelCase					:						Dict					=					pipeline_inputs.pop('''image'''					)
       __lowerCAmelCase					:						Union[str, Any]					=					pipe.image_encoder(A_					).image_embeds

       __lowerCAmelCase					:						Optional[int]					=					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 	UpperCamelCase__  (							self					)					->List[Any]:



       '''simple docstring'''





       __lowerCAmelCase					:						List[str]					=					torch_device == '''cpu'''

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

       self._test_attention_slicing_forward_pass(
           test_max_difference=A_      ,							expected_max_diff=A_					)
 @skip_mps
 def 	UpperCamelCase__  (							self					)					->Any:



       '''simple docstring'''





       __lowerCAmelCase					:						Union[str, Any]					=					torch_device == '''cpu'''
       __lowerCAmelCase					:						str					=					True
       __lowerCAmelCase					:						int					=					[
           '''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 	UpperCamelCase__  (							self					)					->Any:



       '''simple docstring'''





       __lowerCAmelCase					:						List[Any]					=					[
           '''decoder_num_inference_steps''',
           '''super_res_num_inference_steps''',
       ]

       if torch_device == "mps":
             # TODO: MPS errors with larger batch sizes
             __lowerCAmelCase					:						Tuple					=					[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 	UpperCamelCase__  (							self					)					->str:



       '''simple docstring'''





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



       '''simple docstring'''





       return super().test_save_load_local()





 @skip_mps
 def 	UpperCamelCase__  (							self					)					->int:



       '''simple docstring'''





       return super().test_save_load_optional_components()







@slow
@require_torch_gpu
class   __lowercase      (unittest.TestCase ):
 def 	UpperCamelCase__  (							self					)					->Tuple:



       '''simple docstring'''





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





 def 	UpperCamelCase__  (							self					)					->Optional[int]:



       '''simple docstring'''





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

       __lowerCAmelCase					:						Tuple					=					UnCLIPImageVariationPipeline.from_pretrained(
           '''kakaobrain/karlo-v1-alpha-image-variations'''      ,							torch_dtype=torch.floataa					)
       __lowerCAmelCase					:						Tuple					=					pipeline.to(A_					)
       pipeline.set_progress_bar_config(disable=A_					)

       __lowerCAmelCase					:						Dict					=					torch.Generator(device='''cpu'''					).manual_seed(0					)
       __lowerCAmelCase					:						Any					=					pipeline(
           A_      ,							generator=A_      ,							output_type='''np'''      ,							)

       __lowerCAmelCase					:						List[str]					=					output.images[0]

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

       assert_mean_pixel_difference(A_      ,							A_      ,							15					)





def 		_lowercase (     lowercase__   ,	lowercase__						):
      if a < 0 or b < 0:
            raise ValueError('''the value of both inputs must be positive'''						)
      __lowerCAmelCase					:						int					=					str(bin(lowercase__						)						)[2:]  # remove the leading "0b"
      __lowerCAmelCase					:						Any					=					str(bin(lowercase__						)						)[2:]
      __lowerCAmelCase					:						List[str]					=					max(len(lowercase__						)   ,	len(lowercase__						)						)
      return "0b" + "".join(
          str(int('''1''' in (char_a, char_b)						)						)
          for char_a, char_b in zip(a_binary.zfill(lowercase__						)   ,	b_binary.zfill(lowercase__						)						)						)


if __name__ == "__main__":
 import doctest

 doctest.testmod()


