





import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMInverseScheduler,
    DDIMScheduler,
    DPMSolverMultistepInverseScheduler,
    DPMSolverMultistepScheduler,
    StableDiffusionDiffEditPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_image, slow
from diffusers.utils.testing_utils import enable_full_determinism, floats_tensor, require_torch_gpu, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class      lowerCAmelCase							(      __UpperCamelCase,						__UpperCamelCase,						unittest.TestCase    ):
 UpperCAmelCase__									=	StableDiffusionDiffEditPipeline
 UpperCAmelCase__									=	TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"""height""", """width""", """image"""} | {"""image_latents"""}
 UpperCAmelCase__									=	TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {"""image"""} | {"""image_latents"""}
 UpperCAmelCase__									=	frozenset(
     []    )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
 UpperCAmelCase__									=	frozenset([]    )
 def   A_			(   self						:   Dict			)    ->			Union[str, Any]:
     torch.manual_seed(0			)
     lowerCamelCase__					:   List[Any]      				=	UNetaDConditionModel(
         block_out_channels=(32, 64)    ,  layers_per_block=2    ,  sample_size=32    ,  in_channels=4    ,  out_channels=4    ,  down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')    ,  up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')    ,  cross_attention_dim=32    ,  attention_head_dim=(2, 4)    ,  use_linear_projection=UpperCAmelCase    ,  )
     lowerCamelCase__					:   Optional[Any]      				=	DDIMScheduler(
         beta_start=0.0_0_0_8_5    ,  beta_end=0.0_1_2    ,  beta_schedule='scaled_linear'    ,  clip_sample=UpperCAmelCase    ,  set_alpha_to_one=UpperCAmelCase    ,  )
     lowerCamelCase__					:   str      				=	DDIMInverseScheduler(
         beta_start=0.0_0_0_8_5    ,  beta_end=0.0_1_2    ,  beta_schedule='scaled_linear'    ,  clip_sample=UpperCAmelCase    ,  set_alpha_to_zero=UpperCAmelCase    ,  )
     torch.manual_seed(0			)
     lowerCamelCase__					:   Union[str, Any]      				=	AutoencoderKL(
         block_out_channels=[32, 64]    ,  in_channels=3    ,  out_channels=3    ,  down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']    ,  up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']    ,  latent_channels=4    ,  sample_size=128    ,  )
     torch.manual_seed(0			)
     lowerCamelCase__					:   int      				=	CLIPTextConfig(
         bos_token_id=0    ,  eos_token_id=2    ,  hidden_size=32    ,  intermediate_size=37    ,  layer_norm_eps=1e-05    ,  num_attention_heads=4    ,  num_hidden_layers=5    ,  pad_token_id=1    ,  vocab_size=1000    ,  hidden_act='gelu'    ,  projection_dim=512    ,  )
     lowerCamelCase__					:   int      				=	CLIPTextModel(UpperCAmelCase			)
     lowerCamelCase__					:   int      				=	CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'			)

     lowerCamelCase__					:   Tuple      				=	{
         'unet': unet,
         'scheduler': scheduler,
         'inverse_scheduler': inverse_scheduler,
         'vae': vae,
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'safety_checker': None,
         'feature_extractor': None,
     }

     return components
 def   A_			(   self						:   int    ,  UpperCAmelCase						:   int    ,  UpperCAmelCase						:   Dict=0			)    ->			Optional[Any]:
     lowerCamelCase__					:   Any      				=	floats_tensor((1, 16, 16)    ,  rng=random.Random(UpperCAmelCase			)			).to(UpperCAmelCase			)
     lowerCamelCase__					:   int      				=	floats_tensor((1, 2, 4, 16, 16)    ,  rng=random.Random(UpperCAmelCase			)			).to(UpperCAmelCase			)
     if str(UpperCAmelCase			).startswith('mps'			):
         lowerCamelCase__					:   List[Any]      				=	torch.manual_seed(UpperCAmelCase			)
     else:
         lowerCamelCase__					:   Any      				=	torch.Generator(device=UpperCAmelCase			).manual_seed(UpperCAmelCase			)
     lowerCamelCase__					:   int      				=	{
         'prompt': 'a dog and a newt',
         'mask_image': mask,
         'image_latents': latents,
         'generator': generator,
         'num_inference_steps': 2,
         'inpaint_strength': 1.0,
         'guidance_scale': 6.0,
         'output_type': 'numpy',
     }

     return inputs
 def   A_			(   self						:   Optional[int]    ,  UpperCAmelCase						:   Any    ,  UpperCAmelCase						:   str=0			)    ->			Tuple:
     lowerCamelCase__					:   List[Any]      				=	floats_tensor((1, 3, 32, 32)    ,  rng=random.Random(UpperCAmelCase			)			).to(UpperCAmelCase			)
     lowerCamelCase__					:   Dict      				=	image.cpu().permute(0    ,  2    ,  3    ,  1			)[0]
     lowerCamelCase__					:   int      				=	Image.fromarray(np.uinta(UpperCAmelCase			)			).convert('RGB'			)
     if str(UpperCAmelCase			).startswith('mps'			):
         lowerCamelCase__					:   Union[str, Any]      				=	torch.manual_seed(UpperCAmelCase			)
     else:
         lowerCamelCase__					:   List[str]      				=	torch.Generator(device=UpperCAmelCase			).manual_seed(UpperCAmelCase			)
     lowerCamelCase__					:   List[Any]      				=	{
         'image': image,
         'source_prompt': 'a cat and a frog',
         'target_prompt': 'a dog and a newt',
         'generator': generator,
         'num_inference_steps': 2,
         'num_maps_per_mask': 2,
         'mask_encode_strength': 1.0,
         'guidance_scale': 6.0,
         'output_type': 'numpy',
     }

     return inputs
 def   A_			(   self						:   Optional[int]    ,  UpperCAmelCase						:   Union[str, Any]    ,  UpperCAmelCase						:   int=0			)    ->			List[str]:
     lowerCamelCase__					:   str      				=	floats_tensor((1, 3, 32, 32)    ,  rng=random.Random(UpperCAmelCase			)			).to(UpperCAmelCase			)
     lowerCamelCase__					:   str      				=	image.cpu().permute(0    ,  2    ,  3    ,  1			)[0]
     lowerCamelCase__					:   Tuple      				=	Image.fromarray(np.uinta(UpperCAmelCase			)			).convert('RGB'			)
     if str(UpperCAmelCase			).startswith('mps'			):
         lowerCamelCase__					:   List[str]      				=	torch.manual_seed(UpperCAmelCase			)
     else:
         lowerCamelCase__					:   str      				=	torch.Generator(device=UpperCAmelCase			).manual_seed(UpperCAmelCase			)
     lowerCamelCase__					:   Any      				=	{
         'image': image,
         'prompt': 'a cat and a frog',
         'generator': generator,
         'num_inference_steps': 2,
         'inpaint_strength': 1.0,
         'guidance_scale': 6.0,
         'decode_latents': True,
         'output_type': 'numpy',
     }
     return inputs
 def   A_			(   self						:   Union[str, Any]			)    ->			Optional[int]:
     if not hasattr(self.pipeline_class    ,  '_optional_components'			):
         return

     lowerCamelCase__					:   List[Any]      				=	self.get_dummy_components()
     lowerCamelCase__					:   Dict      				=	self.pipeline_class(**UpperCAmelCase			)
     pipe.to(UpperCAmelCase			)
     pipe.set_progress_bar_config(disable=UpperCAmelCase			)

     # set all optional components to None and update pipeline config accordingly
     for optional_component in pipe._optional_components:
         setattr(UpperCAmelCase    ,  UpperCAmelCase    ,  UpperCAmelCase			)
     pipe.register_modules(**{optional_component: None for optional_component in pipe._optional_components}			)

     lowerCamelCase__					:   Tuple      				=	self.get_dummy_inputs(UpperCAmelCase			)
     lowerCamelCase__					:   str      				=	pipe(**UpperCAmelCase			)[0]

     with tempfile.TemporaryDirectory() as tmpdir:
         pipe.save_pretrained(UpperCAmelCase			)
         lowerCamelCase__					:   Dict      				=	self.pipeline_class.from_pretrained(UpperCAmelCase			)
         pipe_loaded.to(UpperCAmelCase			)
         pipe_loaded.set_progress_bar_config(disable=UpperCAmelCase			)

     for optional_component in pipe._optional_components:
         self.assertTrue(
             getattr(UpperCAmelCase    ,  UpperCAmelCase			) is None    ,  F"""`{optional_component}` did not stay set to None after loading."""    ,  )

     lowerCamelCase__					:   Dict      				=	self.get_dummy_inputs(UpperCAmelCase			)
     lowerCamelCase__					:   List[Any]      				=	pipe_loaded(**UpperCAmelCase			)[0]

     lowerCamelCase__					:   Optional[int]      				=	np.abs(output - output_loaded			).max()
     self.assertLess(UpperCAmelCase    ,  1e-4			)
 def   A_			(   self						:   Union[str, Any]			)    ->			Union[str, Any]:
     lowerCamelCase__					:   Dict      				=	'cpu'

     lowerCamelCase__					:   List[Any]      				=	self.get_dummy_components()
     lowerCamelCase__					:   List[str]      				=	self.pipeline_class(**UpperCAmelCase			)
     pipe.to(UpperCAmelCase			)
     pipe.set_progress_bar_config(disable=UpperCAmelCase			)

     lowerCamelCase__					:   List[Any]      				=	self.get_dummy_mask_inputs(UpperCAmelCase			)
     lowerCamelCase__					:   int      				=	pipe.generate_mask(**UpperCAmelCase			)
     lowerCamelCase__					:   Dict      				=	mask[0, -3:, -3:]

     self.assertEqual(mask.shape    ,  (1, 16, 16)			)
     lowerCamelCase__					:   int      				=	np.array([0] * 9			)
     lowerCamelCase__					:   List[str]      				=	np.abs(mask_slice.flatten() - expected_slice			).max()
     self.assertLessEqual(UpperCAmelCase    ,  1e-3			)
     self.assertEqual(mask[0, -3, -4]    ,  0			)
 def   A_			(   self						:   Union[str, Any]			)    ->			Optional[int]:
     lowerCamelCase__					:   Union[str, Any]      				=	'cpu'

     lowerCamelCase__					:   List[Any]      				=	self.get_dummy_components()
     lowerCamelCase__					:   Optional[Any]      				=	self.pipeline_class(**UpperCAmelCase			)
     pipe.to(UpperCAmelCase			)
     pipe.set_progress_bar_config(disable=UpperCAmelCase			)

     lowerCamelCase__					:   Tuple      				=	self.get_dummy_inversion_inputs(UpperCAmelCase			)
     lowerCamelCase__					:   List[Any]      				=	pipe.invert(**UpperCAmelCase			).images
     lowerCamelCase__					:   int      				=	image[0, -1, -3:, -3:]

     self.assertEqual(image.shape    ,  (2, 32, 32, 3)			)
     lowerCamelCase__					:   List[Any]      				=	np.array(
         [0.5_1_5_0, 0.5_1_3_4, 0.5_0_4_3, 0.5_3_7_6, 0.4_6_9_4, 0.5_1_0_5_0, 0.5_0_1_5, 0.4_4_0_7, 0.4_7_9_9]    ,  )
     lowerCamelCase__					:   Dict      				=	np.abs(image_slice.flatten() - expected_slice			).max()
     self.assertLessEqual(UpperCAmelCase    ,  1e-3			)
 def   A_			(   self						:   Union[str, Any]			)    ->			Tuple:
     super().test_inference_batch_single_identical(expected_max_diff=5e-3			)







 def   A_			(   self						:   str			)    ->			List[str]:
     lowerCamelCase__					:   Any      				=	'cpu'

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

     lowerCamelCase__					:   List[str]      				=	{'beta_start': 0.0_0_0_8_5, 'beta_end': 0.0_1_2, 'beta_schedule': 'scaled_linear'}
     lowerCamelCase__					:   List[str]      				=	DPMSolverMultistepScheduler(**UpperCAmelCase			)
     lowerCamelCase__					:   int      				=	DPMSolverMultistepInverseScheduler(**UpperCAmelCase			)

     lowerCamelCase__					:   Any      				=	self.pipeline_class(**UpperCAmelCase			)
     pipe.to(UpperCAmelCase			)
     pipe.set_progress_bar_config(disable=UpperCAmelCase			)

     lowerCamelCase__					:   int      				=	self.get_dummy_inversion_inputs(UpperCAmelCase			)
     lowerCamelCase__					:   List[Any]      				=	pipe.invert(**UpperCAmelCase			).images
     lowerCamelCase__					:   str      				=	image[0, -1, -3:, -3:]

     self.assertEqual(image.shape    ,  (2, 32, 32, 3)			)
     lowerCamelCase__					:   Optional[Any]      				=	np.array(
         [0.5_1_5_0, 0.5_1_3_4, 0.5_0_4_3, 0.5_3_7_6, 0.4_6_9_4, 0.5_1_0_5_0, 0.5_0_1_5, 0.4_4_0_7, 0.4_7_9_9]    ,  )
     lowerCamelCase__					:   List[str]      				=	np.abs(image_slice.flatten() - expected_slice			).max()
     self.assertLessEqual(UpperCAmelCase    ,  1e-3			)


@require_torch_gpu
@slow
class      lowerCAmelCase							(      unittest.TestCase    ):
 def   A_			(   self						:   Optional[int]			)    ->			Dict:
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()
 @classmethod
 def   A_			(   cls						:   Optional[Any]			)    ->			Any:
     lowerCamelCase__					:   Dict      				=	load_image(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png'			)

     lowerCamelCase__					:   int      				=	raw_image.convert('RGB'			).resize((768, 768)			)

     lowerCamelCase__					:   Union[str, Any]      				=	raw_image
 def   A_			(   self						:   Optional[Any]			)    ->			int:
     lowerCamelCase__					:   Dict      				=	torch.manual_seed(0			)

     lowerCamelCase__					:   Dict      				=	StableDiffusionDiffEditPipeline.from_pretrained(
         'stabilityai/stable-diffusion-2-1'    ,  safety_checker=UpperCAmelCase    ,  torch_dtype=torch.floataa			)
     lowerCamelCase__					:   Tuple      				=	DDIMScheduler.from_config(pipe.scheduler.config			)
     lowerCamelCase__					:   int      				=	DDIMInverseScheduler.from_config(pipe.scheduler.config			)
     pipe.enable_model_cpu_offload()
     pipe.set_progress_bar_config(disable=UpperCAmelCase			)

     lowerCamelCase__					:   Union[str, Any]      				=	'a bowl of fruit'
     lowerCamelCase__					:   List[str]      				=	'a bowl of pears'

     lowerCamelCase__					:   Optional[int]      				=	pipe.generate_mask(
         image=self.raw_image    ,  source_prompt=UpperCAmelCase    ,  target_prompt=UpperCAmelCase    ,  generator=UpperCAmelCase    ,  )

     lowerCamelCase__					:   int      				=	pipe.invert(
         prompt=UpperCAmelCase    ,  image=self.raw_image    ,  inpaint_strength=0.7    ,  generator=UpperCAmelCase			).latents

     lowerCamelCase__					:   List[Any]      				=	pipe(
         prompt=UpperCAmelCase    ,  mask_image=UpperCAmelCase    ,  image_latents=UpperCAmelCase    ,  generator=UpperCAmelCase    ,  negative_prompt=UpperCAmelCase    ,  inpaint_strength=0.7    ,  output_type='numpy'    ,  ).images[0]

     lowerCamelCase__					:   List[Any]      				=	(
         np.array(
             load_image(
                 'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                 '/diffedit/pears.png'			).resize((768, 768)			)			)
         / 255
     )
     assert np.abs((expected_image - image).max()			) < 5e-1







 def   A_			(   self						:   str			)    ->			Optional[Any]:
     lowerCamelCase__					:   Any      				=	torch.manual_seed(0			)

     lowerCamelCase__					:   Union[str, Any]      				=	StableDiffusionDiffEditPipeline.from_pretrained(
         'stabilityai/stable-diffusion-2-1'    ,  safety_checker=UpperCAmelCase    ,  torch_dtype=torch.floataa			)
     lowerCamelCase__					:   Union[str, Any]      				=	DPMSolverMultistepScheduler.from_config(pipe.scheduler.config			)
     lowerCamelCase__					:   Optional[int]      				=	DPMSolverMultistepInverseScheduler.from_config(pipe.scheduler.config			)
     pipe.enable_model_cpu_offload()
     pipe.set_progress_bar_config(disable=UpperCAmelCase			)

     lowerCamelCase__					:   Optional[Any]      				=	'a bowl of fruit'
     lowerCamelCase__					:   Optional[Any]      				=	'a bowl of pears'

     lowerCamelCase__					:   Tuple      				=	pipe.generate_mask(
         image=self.raw_image    ,  source_prompt=UpperCAmelCase    ,  target_prompt=UpperCAmelCase    ,  generator=UpperCAmelCase    ,  )

     lowerCamelCase__					:   Optional[int]      				=	pipe.invert(
         prompt=UpperCAmelCase    ,  image=self.raw_image    ,  inpaint_strength=0.7    ,  generator=UpperCAmelCase    ,  num_inference_steps=25    ,  ).latents

     lowerCamelCase__					:   Optional[Any]      				=	pipe(
         prompt=UpperCAmelCase    ,  mask_image=UpperCAmelCase    ,  image_latents=UpperCAmelCase    ,  generator=UpperCAmelCase    ,  negative_prompt=UpperCAmelCase    ,  inpaint_strength=0.7    ,  num_inference_steps=25    ,  output_type='numpy'    ,  ).images[0]

     lowerCamelCase__					:   List[str]      				=	(
         np.array(
             load_image(
                 'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                 '/diffedit/pears.png'			).resize((768, 768)			)			)
         / 255
     )
     assert np.abs((expected_image - image).max()			) < 5e-1












from collections import defaultdict
from math import ceil, sqrt


def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase = 100_0000		,			_UpperCAmelCase = 10				)    ->  int:
    lowerCamelCase__					:   defaultdict      				=	defaultdict(_UpperCAmelCase				)

    for outer_width in range(3		,			(t_limit // 4) + 2				):
        if outer_width * outer_width > t_limit:
            lowerCamelCase__					:   List[str]      				=	max(
                ceil(sqrt(outer_width * outer_width - t_limit				)				)		,			1				)
        else:
            lowerCamelCase__					:   int      				=	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() = }""")





