
import gc
import random
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()




class 				_lowerCamelCase (      UpperCamelCase_	,							UpperCamelCase_	,							UpperCamelCase_	,							unittest.TestCase ):





     """simple docstring"""







     SCREAMING_SNAKE_CASE_		   =  StableUnCLIPImgaImgPipeline
     SCREAMING_SNAKE_CASE_		   =  TEXT_GUIDED_IMAGE_VARIATION_PARAMS
     SCREAMING_SNAKE_CASE_		   =  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
     SCREAMING_SNAKE_CASE_		   =  frozenset(
         [] )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
     SCREAMING_SNAKE_CASE_		   =  frozenset([] )






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Any:




           """simple docstring"""







           UpperCamelCase__    :   Union[str, Any]					=				3_2
           UpperCamelCase__    :   int					=				embedder_hidden_size

           # image encoding components

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

           torch.manual_seed(0  )
           UpperCamelCase__    :   Union[str, Any]					=				CLIPVisionModelWithProjection(
               CLIPVisionConfig(
                   hidden_size=__SCREAMING_SNAKE_CASE						,						projection_dim=__SCREAMING_SNAKE_CASE						,						num_hidden_layers=5						,						num_attention_heads=4						,						image_size=3_2						,						intermediate_size=3_7						,						patch_size=1						,						)  )

           # regular denoising components

           torch.manual_seed(0  )
           UpperCamelCase__    :   Any					=				StableUnCLIPImageNormalizer(embedding_dim=__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Any					=				DDPMScheduler(beta_schedule='''squaredcos_cap_v2'''  )

           torch.manual_seed(0  )
           UpperCamelCase__    :   Any					=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )

           torch.manual_seed(0  )
           UpperCamelCase__    :   Tuple					=				CLIPTextModel(
               CLIPTextConfig(
                   bos_token_id=0						,						eos_token_id=2						,						hidden_size=__SCREAMING_SNAKE_CASE						,						projection_dim=3_2						,						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						,						)  )

           torch.manual_seed(0  )
           UpperCamelCase__    :   int					=				UNetaDConditionModel(
               sample_size=3_2						,						in_channels=4						,						out_channels=4						,						down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''')						,						up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''')						,						block_out_channels=(3_2, 6_4)						,						attention_head_dim=(2, 4)						,						class_embed_type='''projection'''						,						projection_class_embeddings_input_dim=embedder_projection_dim * 2						,						cross_attention_dim=__SCREAMING_SNAKE_CASE						,						layers_per_block=1						,						upcast_attention=__SCREAMING_SNAKE_CASE						,						use_linear_projection=__SCREAMING_SNAKE_CASE						,						)

           torch.manual_seed(0  )
           UpperCamelCase__    :   int					=				DDIMScheduler(
               beta_schedule='''scaled_linear'''						,						beta_start=0.00085						,						beta_end=0.012						,						prediction_type='''v_prediction'''						,						set_alpha_to_one=__SCREAMING_SNAKE_CASE						,						steps_offset=1						,						)

           torch.manual_seed(0  )
           UpperCamelCase__    :   Optional[int]					=				AutoencoderKL()

           UpperCamelCase__    :   Optional[Any]					=				{
               # image encoding components
               '''feature_extractor''': feature_extractor,
               '''image_encoder''': image_encoder.eval(),
               # image noising components
               '''image_normalizer''': image_normalizer.eval(),
               '''image_noising_scheduler''': image_noising_scheduler,
               # regular denoising components
               '''tokenizer''': tokenizer,
               '''text_encoder''': text_encoder.eval(),
               '''unet''': unet.eval(),
               '''scheduler''': scheduler,
               '''vae''': vae.eval(),
           }

           return components






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE=0						,						__SCREAMING_SNAKE_CASE=True  )  ->   str:




           """simple docstring"""







           if str(__SCREAMING_SNAKE_CASE  ).startswith('''mps'''  ):
                 UpperCamelCase__    :   Union[str, Any]					=				torch.manual_seed(__SCREAMING_SNAKE_CASE  )
           else:
                 UpperCamelCase__    :   Dict					=				torch.Generator(device=__SCREAMING_SNAKE_CASE  ).manual_seed(__SCREAMING_SNAKE_CASE  )

           UpperCamelCase__    :   Tuple					=				floats_tensor((1, 3, 3_2, 3_2)						,						rng=random.Random(__SCREAMING_SNAKE_CASE  )  ).to(__SCREAMING_SNAKE_CASE  )

           if pil_image:
                 UpperCamelCase__    :   int					=				input_image * 0.5 + 0.5
                 UpperCamelCase__    :   str					=				input_image.clamp(0						,						1  )
                 UpperCamelCase__    :   List[str]					=				input_image.cpu().permute(0						,						2						,						3						,						1  ).float().numpy()
                 UpperCamelCase__    :   str					=				DiffusionPipeline.numpy_to_pil(__SCREAMING_SNAKE_CASE  )[0]

           return {
               "prompt": "An anime racoon running a marathon",
               "image": input_image,
               "generator": generator,
               "num_inference_steps": 2,
               "output_type": "np",
           }






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




           """simple docstring"""







           UpperCamelCase__    :   Dict					=				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
           UpperCamelCase__    :   Optional[int]					=				self.get_dummy_components()
           UpperCamelCase__    :   str					=				StableUnCLIPImgaImgPipeline(**__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Union[str, Any]					=				sd_pipe.to(__SCREAMING_SNAKE_CASE  )
           sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE  )

           UpperCamelCase__    :   Union[str, Any]					=				self.get_dummy_inputs(__SCREAMING_SNAKE_CASE  )
           inputs.update({'''image_embeds''': None}  )
           UpperCamelCase__    :   List[Any]					=				sd_pipe(**__SCREAMING_SNAKE_CASE  ).images
           UpperCamelCase__    :   List[Any]					=				image[0, -3:, -3:, -1]

           assert image.shape == (1, 3_2, 3_2, 3)
           UpperCamelCase__    :   Any					=				np.array([0.3872, 0.7224, 0.5601, 0.4741, 0.6872, 0.5814, 0.4636, 0.3867, 0.5078]  )

           assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-3






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Any:




           """simple docstring"""







           UpperCamelCase__    :   List[Any]					=				torch_device in ['''cpu''', '''mps''']

           self._test_attention_slicing_forward_pass(test_max_difference=__SCREAMING_SNAKE_CASE  )






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   List[str]:




           """simple docstring"""







           UpperCamelCase__    :   List[str]					=				torch_device in ['''cpu''', '''mps''']

           self._test_inference_batch_single_identical(test_max_difference=__SCREAMING_SNAKE_CASE  )






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




           """simple docstring"""







           self._test_xformers_attention_forwardGenerator_pass(test_max_difference=__SCREAMING_SNAKE_CASE  )




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





     """simple docstring"""







     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   List[str]:




           """simple docstring"""







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






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




           """simple docstring"""







           UpperCamelCase__    :   str					=				load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''  )

           UpperCamelCase__    :   List[str]					=				load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy'''  )

           UpperCamelCase__    :   Union[str, Any]					=				StableUnCLIPImgaImgPipeline.from_pretrained(
               '''fusing/stable-unclip-2-1-l-img2img'''						,						torch_dtype=torch.floataa  )
           pipe.to(__SCREAMING_SNAKE_CASE  )
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE  )
           # stable unclip will oom when integration tests are run on a V100,
           # so turn on memory savings
           pipe.enable_attention_slicing()
           pipe.enable_sequential_cpu_offload()

           UpperCamelCase__    :   int					=				torch.Generator(device='''cpu'''  ).manual_seed(0  )
           UpperCamelCase__    :   str					=				pipe(__SCREAMING_SNAKE_CASE						,						'''anime turle'''						,						generator=__SCREAMING_SNAKE_CASE						,						output_type='''np'''  )

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

           assert image.shape == (7_6_8, 7_6_8, 3)

           assert_mean_pixel_difference(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  )






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




           """simple docstring"""







           UpperCamelCase__    :   Dict					=				load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''  )

           UpperCamelCase__    :   Union[str, Any]					=				load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy'''  )

           UpperCamelCase__    :   int					=				StableUnCLIPImgaImgPipeline.from_pretrained(
               '''fusing/stable-unclip-2-1-h-img2img'''						,						torch_dtype=torch.floataa  )
           pipe.to(__SCREAMING_SNAKE_CASE  )
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE  )
           # stable unclip will oom when integration tests are run on a V100,
           # so turn on memory savings
           pipe.enable_attention_slicing()
           pipe.enable_sequential_cpu_offload()

           UpperCamelCase__    :   int					=				torch.Generator(device='''cpu'''  ).manual_seed(0  )
           UpperCamelCase__    :   int					=				pipe(__SCREAMING_SNAKE_CASE						,						'''anime turle'''						,						generator=__SCREAMING_SNAKE_CASE						,						output_type='''np'''  )

           UpperCamelCase__    :   Dict					=				output.images[0]

           assert image.shape == (7_6_8, 7_6_8, 3)

           assert_mean_pixel_difference(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  )






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




           """simple docstring"""







           UpperCamelCase__    :   Tuple					=				load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png'''  )

           torch.cuda.empty_cache()
           torch.cuda.reset_max_memory_allocated()
           torch.cuda.reset_peak_memory_stats()

           UpperCamelCase__    :   Optional[Any]					=				StableUnCLIPImgaImgPipeline.from_pretrained(
               '''fusing/stable-unclip-2-1-h-img2img'''						,						torch_dtype=torch.floataa  )
           UpperCamelCase__    :   List[str]					=				pipe.to(__SCREAMING_SNAKE_CASE  )
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE  )
           pipe.enable_attention_slicing()
           pipe.enable_sequential_cpu_offload()

           UpperCamelCase__    :   int					=				pipe(
               __SCREAMING_SNAKE_CASE						,						'''anime turtle'''						,						num_inference_steps=2						,						output_type='''np'''						,						)

           UpperCamelCase__    :   List[str]					=				torch.cuda.max_memory_allocated()
           # make sure that less than 7 GB is allocated
           assert mem_bytes < 7 * 1_0**9




import math


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__		):
      UpperCamelCase__    :   List[str]					=				math.loga(math.sqrt(4 * positive_integer + 1		) / 2 + 1 / 2		)

      return exponent == int(UpperCamelCase__		)


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__ = 1 / 1_2_3_4_5		):
      UpperCamelCase__    :   Dict					=				0
      UpperCamelCase__    :   Tuple					=				0

      UpperCamelCase__    :   Optional[Any]					=				3
      while True:
            UpperCamelCase__    :   Union[str, Any]					=				(integer**2 - 1) / 4
            # if candidate is an integer, then there is a partition for k
            if partition_candidate == int(UpperCamelCase__		):
                  UpperCamelCase__    :   int					=				int(UpperCamelCase__		)
                  total_partitions += 1
                  if check_partition_perfect(UpperCamelCase__		):
                        perfect_partitions += 1
            if perfect_partitions > 0:
                  if perfect_partitions / total_partitions < max_proportion:
                        return int(UpperCamelCase__		)
            integer += 1


if __name__ == "__main__":
    print(F'''{solution() = }''')


