






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	):
	a_         =UnCLIPImageVariationPipeline
	a_         =IMAGE_VARIATION_PARAMS - {"""height""", """width""", """guidance_scale"""}
	a_         =IMAGE_VARIATION_BATCH_PARAMS

	a_         =[
	    """generator""",
	    """return_dict""",
	    """decoder_num_inference_steps""",
	    """super_res_num_inference_steps""",
	]
	a_         =False

	@property
	def 		UpperCAmelCase  (							self       )->					List[Any]:







							'''simple docstring'''





							return 32

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







							'''simple docstring'''





							return 32

	@property
	def 		UpperCAmelCase  (							self       )->					Dict:







							'''simple docstring'''





							return self.time_input_dim

	@property
	def 		UpperCAmelCase  (							self       )->					List[Any]:







							'''simple docstring'''





							return self.time_input_dim * 4

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







							'''simple docstring'''





							return 100

	@property
	def 		UpperCAmelCase  (							self       )->					Any:







							'''simple docstring'''





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

	@property
	def 		UpperCAmelCase  (							self       )->					int:







							'''simple docstring'''





							torch.manual_seed(0       )
							lowerCAmelCase__       =						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-0_5		,    num_attention_heads=4		,    num_hidden_layers=5		,    pad_token_id=1		,    vocab_size=1000		,    )
							return CLIPTextModelWithProjection(__UpperCAmelCase       )

	@property
	def 		UpperCAmelCase  (							self       )->					str:







							'''simple docstring'''





							torch.manual_seed(0       )
							lowerCAmelCase__       =						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(__UpperCAmelCase       )

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







							'''simple docstring'''





							torch.manual_seed(0       )

							lowerCAmelCase__       =						{
							    "clip_embeddings_dim": self.text_embedder_hidden_size,
							    "time_embed_dim": self.time_embed_dim,
							    "cross_attention_dim": self.cross_attention_dim,
							}

							lowerCAmelCase__       =						UnCLIPTextProjModel(**__UpperCAmelCase       )
							return model

	@property
	def 		UpperCAmelCase  (							self       )->					int:







							'''simple docstring'''





							torch.manual_seed(0       )

							lowerCAmelCase__       =						{
							    "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__       =						UNetaDConditionModel(**__UpperCAmelCase       )
							return model

	@property
	def 		UpperCAmelCase  (							self       )->					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       )->					Tuple:







							'''simple docstring'''





							torch.manual_seed(0       )

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

	@property
	def 		UpperCAmelCase  (							self       )->					Union[str, Any]:







							'''simple docstring'''





							torch.manual_seed(1       )

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

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







							'''simple docstring'''





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

							lowerCAmelCase__       =						UnCLIPScheduler(
							    variance_type="learned_range"		,    prediction_type="epsilon"		,    num_train_timesteps=1000		,    )

							lowerCAmelCase__       =						UnCLIPScheduler(
							    variance_type="fixed_small_log"		,    prediction_type="epsilon"		,    num_train_timesteps=1000		,    )

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

							lowerCAmelCase__       =						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		,    __UpperCAmelCase		,    __UpperCAmelCase=0		,    __UpperCAmelCase=True       )->					List[str]:







							'''simple docstring'''





							lowerCAmelCase__       =						floats_tensor((1, 3, 32, 32)		,    rng=random.Random(__UpperCAmelCase       )       ).to(__UpperCAmelCase       )
							if str(__UpperCAmelCase       ).startswith("mps"       ):
													lowerCAmelCase__       =						torch.manual_seed(__UpperCAmelCase       )
							else:
													lowerCAmelCase__       =						torch.Generator(device=__UpperCAmelCase       ).manual_seed(__UpperCAmelCase       )

							if pil_image:
													lowerCAmelCase__       =						input_image * 0.5 + 0.5
													lowerCAmelCase__       =						input_image.clamp(0		,    1       )
													lowerCAmelCase__       =						input_image.cpu().permute(0		,    2		,    3		,    1       ).float().numpy()
													lowerCAmelCase__       =						DiffusionPipeline.numpy_to_pil(__UpperCAmelCase       )[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'''





							lowerCAmelCase__       =						"cpu"

							lowerCAmelCase__       =						self.get_dummy_components()

							lowerCAmelCase__       =						self.pipeline_class(**__UpperCAmelCase       )
							lowerCAmelCase__       =						pipe.to(__UpperCAmelCase       )

							pipe.set_progress_bar_config(disable=__UpperCAmelCase       )

							lowerCAmelCase__       =						self.get_dummy_inputs(__UpperCAmelCase		,    pil_image=__UpperCAmelCase       )

							lowerCAmelCase__       =						pipe(**__UpperCAmelCase       )
							lowerCAmelCase__       =						output.images

							lowerCAmelCase__       =						self.get_dummy_inputs(__UpperCAmelCase		,    pil_image=__UpperCAmelCase       )

							lowerCAmelCase__       =						pipe(
							    **__UpperCAmelCase		,    return_dict=__UpperCAmelCase		,    )[0]

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

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

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







							'''simple docstring'''





							lowerCAmelCase__       =						"cpu"

							lowerCAmelCase__       =						self.get_dummy_components()

							lowerCAmelCase__       =						self.pipeline_class(**__UpperCAmelCase       )
							lowerCAmelCase__       =						pipe.to(__UpperCAmelCase       )

							pipe.set_progress_bar_config(disable=__UpperCAmelCase       )

							lowerCAmelCase__       =						self.get_dummy_inputs(__UpperCAmelCase		,    pil_image=__UpperCAmelCase       )

							lowerCAmelCase__       =						pipe(**__UpperCAmelCase       )
							lowerCAmelCase__       =						output.images

							lowerCAmelCase__       =						self.get_dummy_inputs(__UpperCAmelCase		,    pil_image=__UpperCAmelCase       )

							lowerCAmelCase__       =						pipe(
							    **__UpperCAmelCase		,    return_dict=__UpperCAmelCase		,    )[0]

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

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

							lowerCAmelCase__       =						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       )->					Tuple:







							'''simple docstring'''





							lowerCAmelCase__       =						"cpu"

							lowerCAmelCase__       =						self.get_dummy_components()

							lowerCAmelCase__       =						self.pipeline_class(**__UpperCAmelCase       )
							lowerCAmelCase__       =						pipe.to(__UpperCAmelCase       )

							pipe.set_progress_bar_config(disable=__UpperCAmelCase       )

							lowerCAmelCase__       =						self.get_dummy_inputs(__UpperCAmelCase		,    pil_image=__UpperCAmelCase       )
							lowerCAmelCase__       =						[
							    pipeline_inputs["image"],
							    pipeline_inputs["image"],
							]

							lowerCAmelCase__       =						pipe(**__UpperCAmelCase       )
							lowerCAmelCase__       =						output.images

							lowerCAmelCase__       =						self.get_dummy_inputs(__UpperCAmelCase		,    pil_image=__UpperCAmelCase       )
							lowerCAmelCase__       =						[
							    tuple_pipeline_inputs["image"],
							    tuple_pipeline_inputs["image"],
							]

							lowerCAmelCase__       =						pipe(
							    **__UpperCAmelCase		,    return_dict=__UpperCAmelCase		,    )[0]

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

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

							lowerCAmelCase__       =						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__       =						torch.device("cpu"       )


							class      lowercase__     :
								a_         =1





							lowerCAmelCase__       =						self.get_dummy_components()

							lowerCAmelCase__       =						self.pipeline_class(**__UpperCAmelCase       )
							lowerCAmelCase__       =						pipe.to(__UpperCAmelCase       )

							pipe.set_progress_bar_config(disable=__UpperCAmelCase       )

							lowerCAmelCase__       =						torch.Generator(device=__UpperCAmelCase       ).manual_seed(0       )
							lowerCAmelCase__       =						pipe.decoder.dtype
							lowerCAmelCase__       =						1

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

							lowerCAmelCase__       =						(
							    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__       =						pipe.prepare_latents(
							    __UpperCAmelCase		,    dtype=__UpperCAmelCase		,    device=__UpperCAmelCase		,    generator=__UpperCAmelCase		,    latents=__UpperCAmelCase		,    scheduler=DummyScheduler()       )

							lowerCAmelCase__       =						self.get_dummy_inputs(__UpperCAmelCase		,    pil_image=__UpperCAmelCase       )

							lowerCAmelCase__       =						pipe(
							    **__UpperCAmelCase		,    decoder_latents=__UpperCAmelCase		,    super_res_latents=__UpperCAmelCase       ).images

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

							lowerCAmelCase__       =						pipe(
							    **__UpperCAmelCase		,    decoder_latents=__UpperCAmelCase		,    super_res_latents=__UpperCAmelCase		,    image_embeddings=__UpperCAmelCase		,    ).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[str]:







							'''simple docstring'''





							lowerCAmelCase__       =						torch_device == "cpu"

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

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

	@skip_mps
	def 		UpperCAmelCase  (							self       )->					List[str]:







							'''simple docstring'''





							lowerCAmelCase__       =						torch_device == "cpu"
							lowerCAmelCase__       =						True
							lowerCAmelCase__       =						[
							    "decoder_num_inference_steps",
							    "super_res_num_inference_steps",
							]

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

	def 		UpperCAmelCase  (							self       )->					Any:







							'''simple docstring'''





							lowerCAmelCase__       =						[
							    "decoder_num_inference_steps",
							    "super_res_num_inference_steps",
							]

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

	@skip_mps
	def 		UpperCAmelCase  (							self       )->					int:







							'''simple docstring'''





							return super().test_dict_tuple_outputs_equivalent()

	@skip_mps
	def 		UpperCAmelCase  (							self       )->					List[str]:







							'''simple docstring'''





							return super().test_save_load_local()


	@skip_mps
	def 		UpperCAmelCase  (							self       )->					Optional[int]:







							'''simple docstring'''





							return super().test_save_load_optional_components()





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

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







							'''simple docstring'''





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


	def 		UpperCAmelCase  (							self       )->					int:







							'''simple docstring'''





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

							lowerCAmelCase__       =						UnCLIPImageVariationPipeline.from_pretrained(
							    "kakaobrain/karlo-v1-alpha-image-variations"		,    torch_dtype=torch.floataa       )
							lowerCAmelCase__       =						pipeline.to(__UpperCAmelCase       )
							pipeline.set_progress_bar_config(disable=__UpperCAmelCase       )

							lowerCAmelCase__       =						torch.Generator(device="cpu"       ).manual_seed(0       )
							lowerCAmelCase__       =						pipeline(
							    __UpperCAmelCase		,    generator=__UpperCAmelCase		,    output_type="np"		,    )

							lowerCAmelCase__       =						output.images[0]

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

							assert_mean_pixel_difference(__UpperCAmelCase		,    __UpperCAmelCase		,    15       )














from collections import defaultdict
from math import ceil, sqrt


def        _a  (     UpperCamelCase_					:				int = 1_000_000    ,   UpperCamelCase_					:				int = 10   ) ->   int:
						"""simple docstring"""







						lowerCAmelCase__       =						defaultdict(UpperCamelCase_   )

						for outer_width in range(3    ,   (t_limit // 4) + 2   ):
												if outer_width * outer_width > t_limit:
																		lowerCAmelCase__       =						max(
																		    ceil(sqrt(outer_width * outer_width - t_limit   )   )    ,   1   )
												else:
																		lowerCAmelCase__       =						1

												hole_width_lower_bound += (outer_width - hole_width_lower_bound) % 2

												for hole_width in range(UpperCamelCase_    ,   outer_width - 1    ,   2   ):
																		count[outer_width * outer_width - hole_width * hole_width] += 1

						return sum(1 for n in count.values() if 1 <= n <= 10   )


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






