

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 						__magic_name__			(							lowerCamelCase__		,       unittest.TestCase							):

   '''simple docstring'''




   lowerCamelCase__ :	int    =   UnCLIPImageVariationPipeline
   lowerCamelCase__ :	Any    =   IMAGE_VARIATION_PARAMS - {'height', 'width', 'guidance_scale'}
   lowerCamelCase__ :	Union[str, Any]    =   IMAGE_VARIATION_BATCH_PARAMS

   lowerCamelCase__ :	List[str]    =   [
       'generator',
       'return_dict',
       'decoder_num_inference_steps',
       'super_res_num_inference_steps',
   ]
   lowerCamelCase__ :	Union[str, Any]    =   False



   @property
   def  _UpperCAmelCase (				self       )							->    Dict:






     """simple docstring"""






     return 32



   @property
   def  _UpperCAmelCase (				self       )							->    Tuple:






     """simple docstring"""






     return 32



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






     """simple docstring"""






     return self.time_input_dim



   @property
   def  _UpperCAmelCase (				self       )							->    Any:






     """simple docstring"""






     return self.time_input_dim * 4



   @property
   def  _UpperCAmelCase (				self       )							->    Any:






     """simple docstring"""






     return 100



   @property
   def  _UpperCAmelCase (				self       )							->    Any:






     """simple docstring"""






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



   @property
   def  _UpperCAmelCase (				self       )							->    int:






     """simple docstring"""






     torch.manual_seed(0       )
     a__  =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(lowercase_       )



   @property
   def  _UpperCAmelCase (				self       )							->    Optional[int]:






     """simple docstring"""






     torch.manual_seed(0       )
     a__  =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(lowercase_       )



   @property
   def  _UpperCAmelCase (				self       )							->    Any:






     """simple docstring"""






     torch.manual_seed(0       )

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

     a__  =UnCLIPTextProjModel(**lowercase_       )
     return model



   @property
   def  _UpperCAmelCase (				self       )							->    Tuple:






     """simple docstring"""






     torch.manual_seed(0       )

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

     a__  =UNetaDConditionModel(**lowercase_       )
     return model



   @property
   def  _UpperCAmelCase (				self       )							->    List[str]:






     """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       )							->    Optional[Any]:






     """simple docstring"""






     torch.manual_seed(0       )

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



   @property
   def  _UpperCAmelCase (				self       )							->    Optional[int]:






     """simple docstring"""






     torch.manual_seed(1       )

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



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






     """simple docstring"""






     a__  =self.dummy_decoder
     a__  =self.dummy_text_proj
     a__  =self.dummy_text_encoder
     a__  =self.dummy_tokenizer
     a__  =self.dummy_super_res_first
     a__  =self.dummy_super_res_last

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

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

     a__  =CLIPImageProcessor(crop_size=32,    size=32       )

     a__  =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,    lowercase_,    lowercase_=0,    lowercase_=True       )							->    int:






     """simple docstring"""






     a__  =floats_tensor((1, 3, 32, 32),    rng=random.Random(lowercase_       )       ).to(lowercase_       )
     if str(lowercase_       ).startswith('''mps'''       ):
       a__  =torch.manual_seed(lowercase_       )
     else:
       a__  =torch.Generator(device=lowercase_       ).manual_seed(lowercase_       )

     if pil_image:
       a__  =input_image * 0.5 + 0.5
       a__  =input_image.clamp(0,    1       )
       a__  =input_image.cpu().permute(0,    2,    3,    1       ).float().numpy()
       a__  =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  _UpperCAmelCase (				self       )							->    Optional[int]:






     """simple docstring"""






     a__  ='''cpu'''

     a__  =self.get_dummy_components()

     a__  =self.pipeline_class(**lowercase_       )
     a__  =pipe.to(lowercase_       )

     pipe.set_progress_bar_config(disable=lowercase_       )

     a__  =self.get_dummy_inputs(lowercase_,    pil_image=lowercase_       )

     a__  =pipe(**lowercase_       )
     a__  =output.images

     a__  =self.get_dummy_inputs(lowercase_,    pil_image=lowercase_       )

     a__  =pipe(
         **lowercase_,    return_dict=lowercase_,    )[0]

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

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

     a__  =np.array(
         [
             0.9997,
             0.0002,
             0.9997,
             0.9997,
             0.9969,
             0.0023,
             0.9997,
             0.9969,
             0.9970,
         ]       )

     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"""






     a__  ='''cpu'''

     a__  =self.get_dummy_components()

     a__  =self.pipeline_class(**lowercase_       )
     a__  =pipe.to(lowercase_       )

     pipe.set_progress_bar_config(disable=lowercase_       )

     a__  =self.get_dummy_inputs(lowercase_,    pil_image=lowercase_       )

     a__  =pipe(**lowercase_       )
     a__  =output.images

     a__  =self.get_dummy_inputs(lowercase_,    pil_image=lowercase_       )

     a__  =pipe(
         **lowercase_,    return_dict=lowercase_,    )[0]

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

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

     a__  =np.array([0.9997, 0.0003, 0.9997, 0.9997, 0.9970, 0.0024, 0.9997, 0.9971, 0.9971]       )

     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"""






     a__  ='''cpu'''

     a__  =self.get_dummy_components()

     a__  =self.pipeline_class(**lowercase_       )
     a__  =pipe.to(lowercase_       )

     pipe.set_progress_bar_config(disable=lowercase_       )

     a__  =self.get_dummy_inputs(lowercase_,    pil_image=lowercase_       )
     a__  =[
         pipeline_inputs['''image'''],
         pipeline_inputs['''image'''],
     ]

     a__  =pipe(**lowercase_       )
     a__  =output.images

     a__  =self.get_dummy_inputs(lowercase_,    pil_image=lowercase_       )
     a__  =[
         tuple_pipeline_inputs['''image'''],
         tuple_pipeline_inputs['''image'''],
     ]

     a__  =pipe(
         **lowercase_,    return_dict=lowercase_,    )[0]

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

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

     a__  =np.array(
         [
             0.9997,
             0.9989,
             0.0008,
             0.0021,
             0.9960,
             0.0018,
             0.0014,
             0.0002,
             0.9933,
         ]       )

     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       )							->    str:






     """simple docstring"""






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







     class 						__magic_name__			:

        '''simple docstring'''




        lowerCamelCase__ :	Union[str, Any]    =   1





     a__  =self.get_dummy_components()

     a__  =self.pipeline_class(**lowercase_       )
     a__  =pipe.to(lowercase_       )

     pipe.set_progress_bar_config(disable=lowercase_       )

     a__  =torch.Generator(device=lowercase_       ).manual_seed(0       )
     a__  =pipe.decoder.dtype
     a__  =1

     a__  =(
         batch_size,
         pipe.decoder.config.in_channels,
         pipe.decoder.config.sample_size,
         pipe.decoder.config.sample_size,
     )
     a__  =pipe.prepare_latents(
         lowercase_,    dtype=lowercase_,    device=lowercase_,    generator=lowercase_,    latents=lowercase_,    scheduler=DummyScheduler()       )

     a__  =(
         batch_size,
         pipe.super_res_first.config.in_channels // 2,
         pipe.super_res_first.config.sample_size,
         pipe.super_res_first.config.sample_size,
     )
     a__  =pipe.prepare_latents(
         lowercase_,    dtype=lowercase_,    device=lowercase_,    generator=lowercase_,    latents=lowercase_,    scheduler=DummyScheduler()       )

     a__  =self.get_dummy_inputs(lowercase_,    pil_image=lowercase_       )

     a__  =pipe(
         **lowercase_,    decoder_latents=lowercase_,    super_res_latents=lowercase_       ).images

     a__  =self.get_dummy_inputs(lowercase_,    pil_image=lowercase_       )
     # Don't pass image, instead pass embedding
     a__  =pipeline_inputs.pop('''image'''       )
     a__  =pipe.image_encoder(lowercase_       ).image_embeds

     a__  =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  _UpperCAmelCase (				self       )							->    Dict:






     """simple docstring"""






     a__  =torch_device == '''cpu'''

     # Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
     a__  =1E-2

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



   @skip_mps
   def  _UpperCAmelCase (				self       )							->    Dict:






     """simple docstring"""






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






     """simple docstring"""






     a__  =[
         '''decoder_num_inference_steps''',
         '''super_res_num_inference_steps''',
     ]

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






     """simple docstring"""






     return super().test_dict_tuple_outputs_equivalent()



   @skip_mps
   def  _UpperCAmelCase (				self       )							->    Union[str, Any]:






     """simple docstring"""






     return super().test_save_load_local()



   @skip_mps
   def  _UpperCAmelCase (				self       )							->    List[Any]:






     """simple docstring"""






     return super().test_save_load_optional_components()







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

   '''simple docstring'''




   def  _UpperCAmelCase (				self       )							->    str:






     """simple docstring"""






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



   def  _UpperCAmelCase (				self       )							->    List[Any]:






     """simple docstring"""






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

     a__  =UnCLIPImageVariationPipeline.from_pretrained(
         '''kakaobrain/karlo-v1-alpha-image-variations''',    torch_dtype=torch.floataa       )
     a__  =pipeline.to(lowercase_       )
     pipeline.set_progress_bar_config(disable=lowercase_       )

     a__  =torch.Generator(device='''cpu'''       ).manual_seed(0       )
     a__  =pipeline(
         lowercase_,    generator=lowercase_,    output_type='''np''',    )

     a__  =output.images[0]

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

     assert_mean_pixel_difference(lowercase_,    lowercase_,    15       )



import random
import unittest

import torch

from diffusers import IFInpaintingSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin







@skip_mps
class 						__magic_name__			(							lowerCamelCase__		,       lowerCamelCase__		,       unittest.TestCase							):

   '''simple docstring'''




   lowerCamelCase__ :	List[Any]    =   IFInpaintingSuperResolutionPipeline
   lowerCamelCase__ :	str    =   TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {'width', 'height'}
   lowerCamelCase__ :	List[Any]    =   TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS.union({'original_image'}							)
   lowerCamelCase__ :	str    =   PipelineTesterMixin.required_optional_params - {'latents'}



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






     """simple docstring"""






     return self._get_superresolution_dummy_components()



   def  _UpperCAmelCase (				self,    lowercase_,    lowercase_=0       )							->    Tuple:






     """simple docstring"""






     if str(lowercase_       ).startswith('''mps'''       ):
       a__  =torch.manual_seed(lowercase_       )
     else:
       a__  =torch.Generator(device=lowercase_       ).manual_seed(lowercase_       )

     a__  =floats_tensor((1, 3, 16, 16),    rng=random.Random(lowercase_       )       ).to(lowercase_       )
     a__  =floats_tensor((1, 3, 32, 32),    rng=random.Random(lowercase_       )       ).to(lowercase_       )
     a__  =floats_tensor((1, 3, 32, 32),    rng=random.Random(lowercase_       )       ).to(lowercase_       )

     a__  ={
         '''prompt''': '''A painting of a squirrel eating a burger''',
         '''image''': image,
         '''original_image''': original_image,
         '''mask_image''': mask_image,
         '''generator''': generator,
         '''num_inference_steps''': 2,
         '''output_type''': '''numpy''',
     }

     return inputs



   @unittest.skipIf(
       torch_device != '''cuda''' or not is_xformers_available(),    reason='''XFormers attention is only available with CUDA and `xformers` installed''',    )
   def  _UpperCAmelCase (				self       )							->    List[str]:






     """simple docstring"""






     self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3       )



   def  _UpperCAmelCase (				self       )							->    int:






     """simple docstring"""






     self._test_save_load_optional_components()



   @unittest.skipIf(torch_device != '''cuda''',    reason='''float16 requires CUDA'''       )
   def  _UpperCAmelCase (				self       )							->    List[Any]:






     """simple docstring"""






     super().test_save_load_floataa(expected_max_diff=1E-1       )



   def  _UpperCAmelCase (				self       )							->    int:






     """simple docstring"""






     self._test_attention_slicing_forward_pass(expected_max_diff=1E-2       )



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






     """simple docstring"""






     self._test_save_load_local()



   def  _UpperCAmelCase (				self       )							->    str:






     """simple docstring"""






     self._test_inference_batch_single_identical(
         expected_max_diff=1E-2,    )
