




'''simple docstring'''




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   __UpperCamelCase       ( lowerCamelCase__			,					lowerCamelCase__			,					unittest.TestCase				):
       lowercase	:    Tuple      						=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	:    Tuple      						=frozenset(
           []				)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
       lowercase	:    Union[str, Any]      						=frozenset([]				)






       def      lowercase__     (							self  ):



              """simple docstring"""







              torch.manual_seed(0  )
              lowerCamelCase_							=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=lowerCAmelCase,		)
              lowerCamelCase_							=DDIMScheduler(
                  beta_start=0.0_0_0_8_5,		beta_end=0.0_1_2,		beta_schedule='''scaled_linear''',		clip_sample=lowerCAmelCase,		set_alpha_to_one=lowerCAmelCase,		)
              lowerCamelCase_							=DDIMInverseScheduler(
                  beta_start=0.0_0_0_8_5,		beta_end=0.0_1_2,		beta_schedule='''scaled_linear''',		clip_sample=lowerCAmelCase,		set_alpha_to_zero=lowerCAmelCase,		)
              torch.manual_seed(0  )
              lowerCamelCase_							=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_							=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=1_000,		hidden_act='''gelu''',		projection_dim=512,		)
              lowerCamelCase_							=CLIPTextModel(lowerCAmelCase  )
              lowerCamelCase_							=CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )

              lowerCamelCase_							={
                  '''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      lowercase__     (							self,		lowerCAmelCase,		lowerCAmelCase=0  ):



              """simple docstring"""







              lowerCamelCase_							=floats_tensor((1, 16, 16),		rng=random.Random(lowerCAmelCase  )  ).to(lowerCAmelCase  )
              lowerCamelCase_							=floats_tensor((1, 2, 4, 16, 16),		rng=random.Random(lowerCAmelCase  )  ).to(lowerCAmelCase  )
              if str(lowerCAmelCase  ).startswith('''mps'''  ):
                     lowerCamelCase_							=torch.manual_seed(lowerCAmelCase  )
              else:
                     lowerCamelCase_							=torch.Generator(device=lowerCAmelCase  ).manual_seed(lowerCAmelCase  )
              lowerCamelCase_							={
                  '''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      lowercase__     (							self,		lowerCAmelCase,		lowerCAmelCase=0  ):



              """simple docstring"""







              lowerCamelCase_							=floats_tensor((1, 3, 32, 32),		rng=random.Random(lowerCAmelCase  )  ).to(lowerCAmelCase  )
              lowerCamelCase_							=image.cpu().permute(0,		2,		3,		1  )[0]
              lowerCamelCase_							=Image.fromarray(np.uinta(lowerCAmelCase  )  ).convert('''RGB'''  )
              if str(lowerCAmelCase  ).startswith('''mps'''  ):
                     lowerCamelCase_							=torch.manual_seed(lowerCAmelCase  )
              else:
                     lowerCamelCase_							=torch.Generator(device=lowerCAmelCase  ).manual_seed(lowerCAmelCase  )
              lowerCamelCase_							={
                  '''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      lowercase__     (							self,		lowerCAmelCase,		lowerCAmelCase=0  ):



              """simple docstring"""







              lowerCamelCase_							=floats_tensor((1, 3, 32, 32),		rng=random.Random(lowerCAmelCase  )  ).to(lowerCAmelCase  )
              lowerCamelCase_							=image.cpu().permute(0,		2,		3,		1  )[0]
              lowerCamelCase_							=Image.fromarray(np.uinta(lowerCAmelCase  )  ).convert('''RGB'''  )
              if str(lowerCAmelCase  ).startswith('''mps'''  ):
                     lowerCamelCase_							=torch.manual_seed(lowerCAmelCase  )
              else:
                     lowerCamelCase_							=torch.Generator(device=lowerCAmelCase  ).manual_seed(lowerCAmelCase  )
              lowerCamelCase_							={
                  '''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      lowercase__     (							self  ):



              """simple docstring"""







              if not hasattr(self.pipeline_class,		'''_optional_components'''  ):
                     return

              lowerCamelCase_							=self.get_dummy_components()
              lowerCamelCase_							=self.pipeline_class(**lowerCAmelCase  )
              pipe.to(lowerCAmelCase  )
              pipe.set_progress_bar_config(disable=lowerCAmelCase  )

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

              lowerCamelCase_							=self.get_dummy_inputs(lowerCAmelCase  )
              lowerCamelCase_							=pipe(**lowerCAmelCase  )[0]

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

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

              lowerCamelCase_							=self.get_dummy_inputs(lowerCAmelCase  )
              lowerCamelCase_							=pipe_loaded(**lowerCAmelCase  )[0]

              lowerCamelCase_							=np.abs(output - output_loaded  ).max()
              self.assertLess(lowerCAmelCase,		1e-4  )






       def      lowercase__     (							self  ):



              """simple docstring"""







              lowerCamelCase_							='''cpu'''

              lowerCamelCase_							=self.get_dummy_components()
              lowerCamelCase_							=self.pipeline_class(**lowerCAmelCase  )
              pipe.to(lowerCAmelCase  )
              pipe.set_progress_bar_config(disable=lowerCAmelCase  )

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

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






       def      lowercase__     (							self  ):



              """simple docstring"""







              lowerCamelCase_							='''cpu'''

              lowerCamelCase_							=self.get_dummy_components()
              lowerCamelCase_							=self.pipeline_class(**lowerCAmelCase  )
              pipe.to(lowerCAmelCase  )
              pipe.set_progress_bar_config(disable=lowerCAmelCase  )

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

              self.assertEqual(image.shape,		(2, 32, 32, 3)  )
              lowerCamelCase_							=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_							=np.abs(image_slice.flatten() - expected_slice  ).max()
              self.assertLessEqual(lowerCAmelCase,		1e-3  )






       def      lowercase__     (							self  ):



              """simple docstring"""







              super().test_inference_batch_single_identical(expected_max_diff=5e-3  )






       def      lowercase__     (							self  ):



              """simple docstring"""







              lowerCamelCase_							='''cpu'''

              lowerCamelCase_							=self.get_dummy_components()

              lowerCamelCase_							={'''beta_start''': 0.0_0_0_8_5, '''beta_end''': 0.0_1_2, '''beta_schedule''': '''scaled_linear'''}
              lowerCamelCase_							=DPMSolverMultistepScheduler(**lowerCAmelCase  )
              lowerCamelCase_							=DPMSolverMultistepInverseScheduler(**lowerCAmelCase  )

              lowerCamelCase_							=self.pipeline_class(**lowerCAmelCase  )
              pipe.to(lowerCAmelCase  )
              pipe.set_progress_bar_config(disable=lowerCAmelCase  )

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

              self.assertEqual(image.shape,		(2, 32, 32, 3)  )
              lowerCamelCase_							=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_							=np.abs(image_slice.flatten() - expected_slice  ).max()
              self.assertLessEqual(lowerCAmelCase,		1e-3  )





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






       def      lowercase__     (							self  ):



              """simple docstring"""







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






       @classmethod
       def      lowercase__     (							cls  ):



              """simple docstring"""







              lowerCamelCase_							=load_image(
                  '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png'''  )

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

              lowerCamelCase_							=raw_image






       def      lowercase__     (							self  ):



              """simple docstring"""







              lowerCamelCase_							=torch.manual_seed(0  )

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

              lowerCamelCase_							='''a bowl of fruit'''
              lowerCamelCase_							='''a bowl of pears'''

              lowerCamelCase_							=pipe.generate_mask(
                  image=self.raw_image,		source_prompt=lowerCAmelCase,		target_prompt=lowerCAmelCase,		generator=lowerCAmelCase,		)

              lowerCamelCase_							=pipe.invert(
                  prompt=lowerCAmelCase,		image=self.raw_image,		inpaint_strength=0.7,		generator=lowerCAmelCase  ).latents

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

              lowerCamelCase_							=(
                  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      lowercase__     (							self  ):



              """simple docstring"""







              lowerCamelCase_							=torch.manual_seed(0  )

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

              lowerCamelCase_							='''a bowl of fruit'''
              lowerCamelCase_							='''a bowl of pears'''

              lowerCamelCase_							=pipe.generate_mask(
                  image=self.raw_image,		source_prompt=lowerCAmelCase,		target_prompt=lowerCAmelCase,		generator=lowerCAmelCase,		)

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

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

              lowerCamelCase_							=(
                  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'''

import warnings


warnings.warn(
    "memory_utils has been reorganized to utils.memory. Import `find_executable_batchsize` from the main `__init__`: "
    "`from accelerate import find_executable_batch_size` to avoid this warning.",
    FutureWarning,
)



