






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 		_snake_case	(	A__ ,     A__ ,     unittest.TestCase       ):
			_lowercase		:     List[Any]        =				StableDiffusionDiffEditPipeline
			_lowercase		:     List[str]        =				TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {'''height''', '''width''', '''image'''} | {'''image_latents'''}
			_lowercase		:     Optional[int]        =				TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {'''image'''} | {'''image_latents'''}
			_lowercase		:     List[str]        =				frozenset(
			    []       )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
			_lowercase		:     List[str]        =				frozenset([]       )






			def  SCREAMING_SNAKE_CASE__     (							self)   ->							Optional[Any]:
						torch.manual_seed(0)
						SCREAMING_SNAKE_CASE							=			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=a		,       )
						SCREAMING_SNAKE_CASE							=			DDIMScheduler(
						    beta_start=0.0_00_85		,       beta_end=0.0_12		,       beta_schedule='scaled_linear'		,       clip_sample=a		,       set_alpha_to_one=a		,       )
						SCREAMING_SNAKE_CASE							=			DDIMInverseScheduler(
						    beta_start=0.0_00_85		,       beta_end=0.0_12		,       beta_schedule='scaled_linear'		,       clip_sample=a		,       set_alpha_to_zero=a		,       )
						torch.manual_seed(0)
						SCREAMING_SNAKE_CASE							=			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)
						SCREAMING_SNAKE_CASE							=			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		,       )
						SCREAMING_SNAKE_CASE							=			CLIPTextModel(a)
						SCREAMING_SNAKE_CASE							=			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

						SCREAMING_SNAKE_CASE							=			{
						    '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  SCREAMING_SNAKE_CASE__     (							self		,       a		,       a=0)   ->							List[Any]:
						SCREAMING_SNAKE_CASE							=			floats_tensor((1, 16, 16)		,       rng=random.Random(a)).to(a)
						SCREAMING_SNAKE_CASE							=			floats_tensor((1, 2, 4, 16, 16)		,       rng=random.Random(a)).to(a)
						if str(a).startswith('mps'):
									SCREAMING_SNAKE_CASE							=			torch.manual_seed(a)
						else:
									SCREAMING_SNAKE_CASE							=			torch.Generator(device=a).manual_seed(a)
						SCREAMING_SNAKE_CASE							=			{
						    '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  SCREAMING_SNAKE_CASE__     (							self		,       a		,       a=0)   ->							List[Any]:
						SCREAMING_SNAKE_CASE							=			floats_tensor((1, 3, 32, 32)		,       rng=random.Random(a)).to(a)
						SCREAMING_SNAKE_CASE							=			image.cpu().permute(0		,       2		,       3		,       1)[0]
						SCREAMING_SNAKE_CASE							=			Image.fromarray(np.uinta(a)).convert('RGB')
						if str(a).startswith('mps'):
									SCREAMING_SNAKE_CASE							=			torch.manual_seed(a)
						else:
									SCREAMING_SNAKE_CASE							=			torch.Generator(device=a).manual_seed(a)
						SCREAMING_SNAKE_CASE							=			{
						    '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  SCREAMING_SNAKE_CASE__     (							self		,       a		,       a=0)   ->							Optional[int]:
						SCREAMING_SNAKE_CASE							=			floats_tensor((1, 3, 32, 32)		,       rng=random.Random(a)).to(a)
						SCREAMING_SNAKE_CASE							=			image.cpu().permute(0		,       2		,       3		,       1)[0]
						SCREAMING_SNAKE_CASE							=			Image.fromarray(np.uinta(a)).convert('RGB')
						if str(a).startswith('mps'):
									SCREAMING_SNAKE_CASE							=			torch.manual_seed(a)
						else:
									SCREAMING_SNAKE_CASE							=			torch.Generator(device=a).manual_seed(a)
						SCREAMING_SNAKE_CASE							=			{
						    '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  SCREAMING_SNAKE_CASE__     (							self)   ->							Optional[int]:
						if not hasattr(self.pipeline_class		,       '_optional_components'):
									return

						SCREAMING_SNAKE_CASE							=			self.get_dummy_components()
						SCREAMING_SNAKE_CASE							=			self.pipeline_class(**a)
						pipe.to(a)
						pipe.set_progress_bar_config(disable=a)

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

						SCREAMING_SNAKE_CASE							=			self.get_dummy_inputs(a)
						SCREAMING_SNAKE_CASE							=			pipe(**a)[0]

						with tempfile.TemporaryDirectory() as tmpdir:
									pipe.save_pretrained(a)
									SCREAMING_SNAKE_CASE							=			self.pipeline_class.from_pretrained(a)
									pipe_loaded.to(a)
									pipe_loaded.set_progress_bar_config(disable=a)

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

						SCREAMING_SNAKE_CASE							=			self.get_dummy_inputs(a)
						SCREAMING_SNAKE_CASE							=			pipe_loaded(**a)[0]

						SCREAMING_SNAKE_CASE							=			np.abs(output - output_loaded).max()
						self.assertLess(a		,       1E-4)






			def  SCREAMING_SNAKE_CASE__     (							self)   ->							str:
						SCREAMING_SNAKE_CASE							=			'cpu'

						SCREAMING_SNAKE_CASE							=			self.get_dummy_components()
						SCREAMING_SNAKE_CASE							=			self.pipeline_class(**a)
						pipe.to(a)
						pipe.set_progress_bar_config(disable=a)

						SCREAMING_SNAKE_CASE							=			self.get_dummy_mask_inputs(a)
						SCREAMING_SNAKE_CASE							=			pipe.generate_mask(**a)
						SCREAMING_SNAKE_CASE							=			mask[0, -3:, -3:]

						self.assertEqual(mask.shape		,       (1, 16, 16))
						SCREAMING_SNAKE_CASE							=			np.array([0] * 9)
						SCREAMING_SNAKE_CASE							=			np.abs(mask_slice.flatten() - expected_slice).max()
						self.assertLessEqual(a		,       1E-3)
						self.assertEqual(mask[0, -3, -4]		,       0)






			def  SCREAMING_SNAKE_CASE__     (							self)   ->							str:
						SCREAMING_SNAKE_CASE							=			'cpu'

						SCREAMING_SNAKE_CASE							=			self.get_dummy_components()
						SCREAMING_SNAKE_CASE							=			self.pipeline_class(**a)
						pipe.to(a)
						pipe.set_progress_bar_config(disable=a)

						SCREAMING_SNAKE_CASE							=			self.get_dummy_inversion_inputs(a)
						SCREAMING_SNAKE_CASE							=			pipe.invert(**a).images
						SCREAMING_SNAKE_CASE							=			image[0, -1, -3:, -3:]

						self.assertEqual(image.shape		,       (2, 32, 32, 3))
						SCREAMING_SNAKE_CASE							=			np.array(
						    [0.51_50, 0.51_34, 0.50_43, 0.53_76, 0.46_94, 0.5_10_50, 0.50_15, 0.44_07, 0.47_99]		,       )
						SCREAMING_SNAKE_CASE							=			np.abs(image_slice.flatten() - expected_slice).max()
						self.assertLessEqual(a		,       1E-3)






			def  SCREAMING_SNAKE_CASE__     (							self)   ->							Dict:
						super().test_inference_batch_single_identical(expected_max_diff=5E-3)







			def  SCREAMING_SNAKE_CASE__     (							self)   ->							List[Any]:
						SCREAMING_SNAKE_CASE							=			'cpu'

						SCREAMING_SNAKE_CASE							=			self.get_dummy_components()

						SCREAMING_SNAKE_CASE							=			{'beta_start': 0.0_00_85, 'beta_end': 0.0_12, 'beta_schedule': 'scaled_linear'}
						SCREAMING_SNAKE_CASE							=			DPMSolverMultistepScheduler(**a)
						SCREAMING_SNAKE_CASE							=			DPMSolverMultistepInverseScheduler(**a)

						SCREAMING_SNAKE_CASE							=			self.pipeline_class(**a)
						pipe.to(a)
						pipe.set_progress_bar_config(disable=a)

						SCREAMING_SNAKE_CASE							=			self.get_dummy_inversion_inputs(a)
						SCREAMING_SNAKE_CASE							=			pipe.invert(**a).images
						SCREAMING_SNAKE_CASE							=			image[0, -1, -3:, -3:]

						self.assertEqual(image.shape		,       (2, 32, 32, 3))
						SCREAMING_SNAKE_CASE							=			np.array(
						    [0.51_50, 0.51_34, 0.50_43, 0.53_76, 0.46_94, 0.5_10_50, 0.50_15, 0.44_07, 0.47_99]		,       )
						SCREAMING_SNAKE_CASE							=			np.abs(image_slice.flatten() - expected_slice).max()
						self.assertLessEqual(a		,       1E-3)



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






			def  SCREAMING_SNAKE_CASE__     (							self)   ->							Any:
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()






			@classmethod
			def  SCREAMING_SNAKE_CASE__     (							cls)   ->							List[Any]:
						SCREAMING_SNAKE_CASE							=			load_image(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png')

						SCREAMING_SNAKE_CASE							=			raw_image.convert('RGB').resize((768, 768))

						SCREAMING_SNAKE_CASE							=			raw_image






			def  SCREAMING_SNAKE_CASE__     (							self)   ->							List[Any]:
						SCREAMING_SNAKE_CASE							=			torch.manual_seed(0)

						SCREAMING_SNAKE_CASE							=			StableDiffusionDiffEditPipeline.from_pretrained(
						    'stabilityai/stable-diffusion-2-1'		,       safety_checker=a		,       torch_dtype=torch.floataa)
						SCREAMING_SNAKE_CASE							=			DDIMScheduler.from_config(pipe.scheduler.config)
						SCREAMING_SNAKE_CASE							=			DDIMInverseScheduler.from_config(pipe.scheduler.config)
						pipe.enable_model_cpu_offload()
						pipe.set_progress_bar_config(disable=a)

						SCREAMING_SNAKE_CASE							=			'a bowl of fruit'
						SCREAMING_SNAKE_CASE							=			'a bowl of pears'

						SCREAMING_SNAKE_CASE							=			pipe.generate_mask(
						    image=self.raw_image		,       source_prompt=a		,       target_prompt=a		,       generator=a		,       )

						SCREAMING_SNAKE_CASE							=			pipe.invert(
						    prompt=a		,       image=self.raw_image		,       inpaint_strength=0.7		,       generator=a).latents

						SCREAMING_SNAKE_CASE							=			pipe(
						    prompt=a		,       mask_image=a		,       image_latents=a		,       generator=a		,       negative_prompt=a		,       inpaint_strength=0.7		,       output_type='numpy'		,       ).images[0]

						SCREAMING_SNAKE_CASE							=			(
						    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  SCREAMING_SNAKE_CASE__     (							self)   ->							str:
						SCREAMING_SNAKE_CASE							=			torch.manual_seed(0)

						SCREAMING_SNAKE_CASE							=			StableDiffusionDiffEditPipeline.from_pretrained(
						    'stabilityai/stable-diffusion-2-1'		,       safety_checker=a		,       torch_dtype=torch.floataa)
						SCREAMING_SNAKE_CASE							=			DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
						SCREAMING_SNAKE_CASE							=			DPMSolverMultistepInverseScheduler.from_config(pipe.scheduler.config)
						pipe.enable_model_cpu_offload()
						pipe.set_progress_bar_config(disable=a)

						SCREAMING_SNAKE_CASE							=			'a bowl of fruit'
						SCREAMING_SNAKE_CASE							=			'a bowl of pears'

						SCREAMING_SNAKE_CASE							=			pipe.generate_mask(
						    image=self.raw_image		,       source_prompt=a		,       target_prompt=a		,       generator=a		,       )

						SCREAMING_SNAKE_CASE							=			pipe.invert(
						    prompt=a		,       image=self.raw_image		,       inpaint_strength=0.7		,       generator=a		,       num_inference_steps=25		,       ).latents

						SCREAMING_SNAKE_CASE							=			pipe(
						    prompt=a		,       mask_image=a		,       image_latents=a		,       generator=a		,       negative_prompt=a		,       inpaint_strength=0.7		,       num_inference_steps=25		,       output_type='numpy'		,       ).images[0]

						SCREAMING_SNAKE_CASE							=			(
						    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





'''simple docstring'''


def   A_      (       SCREAMING_SNAKE_CASE_						)			->bool:
  if num < 0:
    return False

  lowercase_     =       num
  lowercase_     =       0
  while num > 0:
    lowercase_     =       rev_num * 10 + (num % 10)
    num //= 10

  return num_copy == rev_num


if __name__ == "__main__":
       import doctest

       doctest.testmod()

