





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        __magic_name__	(				A_     ,					A_     ,					unittest.TestCase				):





     """simple docstring"""



     __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     SCREAMING_SNAKE_CASE		(				self							:Tuple	):





        '''simple docstring'''




        torch.manual_seed(0	)
        A_       :					Dict        =      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=snake_case__      ,   )
        A_       :					str        =      DDIMScheduler(
            beta_start=0.00085      ,   beta_end=0.012      ,   beta_schedule="scaled_linear"      ,   clip_sample=snake_case__      ,   set_alpha_to_one=snake_case__      ,   )
        A_       :					List[Any]        =      DDIMInverseScheduler(
            beta_start=0.00085      ,   beta_end=0.012      ,   beta_schedule="scaled_linear"      ,   clip_sample=snake_case__      ,   set_alpha_to_zero=snake_case__      ,   )
        torch.manual_seed(0	)
        A_       :					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	)
        A_       :					Union[str, Any]        =      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      ,   )
        A_       :					str        =      CLIPTextModel(snake_case__	)
        A_       :					Optional[Any]        =      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"	)

        A_       :					List[str]        =      {
            "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							:List[str]      ,   snake_case							:Dict      ,   snake_case							:Union[str, Any]=0	):





        '''simple docstring'''




        A_       :					List[Any]        =      floats_tensor((1, 16, 16)      ,   rng=random.Random(snake_case__	)	).to(snake_case__	)
        A_       :					List[Any]        =      floats_tensor((1, 2, 4, 16, 16)      ,   rng=random.Random(snake_case__	)	).to(snake_case__	)
        if str(snake_case__	).startswith("mps"	):
           A_       :					Any        =      torch.manual_seed(snake_case__	)
        else:
           A_       :					Optional[Any]        =      torch.Generator(device=snake_case__	).manual_seed(snake_case__	)
        A_       :					Union[str, Any]        =      {
            "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							:Tuple      ,   snake_case							:Optional[int]      ,   snake_case							:Optional[int]=0	):





        '''simple docstring'''




        A_       :					Tuple        =      floats_tensor((1, 3, 32, 32)      ,   rng=random.Random(snake_case__	)	).to(snake_case__	)
        A_       :					Optional[Any]        =      image.cpu().permute(0      ,   2      ,   3      ,   1	)[0]
        A_       :					Any        =      Image.fromarray(np.uinta(snake_case__	)	).convert("RGB"	)
        if str(snake_case__	).startswith("mps"	):
           A_       :					int        =      torch.manual_seed(snake_case__	)
        else:
           A_       :					Optional[int]        =      torch.Generator(device=snake_case__	).manual_seed(snake_case__	)
        A_       :					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     SCREAMING_SNAKE_CASE		(				self							:Any      ,   snake_case							:Optional[Any]      ,   snake_case							:Tuple=0	):





        '''simple docstring'''




        A_       :					Optional[Any]        =      floats_tensor((1, 3, 32, 32)      ,   rng=random.Random(snake_case__	)	).to(snake_case__	)
        A_       :					Tuple        =      image.cpu().permute(0      ,   2      ,   3      ,   1	)[0]
        A_       :					List[Any]        =      Image.fromarray(np.uinta(snake_case__	)	).convert("RGB"	)
        if str(snake_case__	).startswith("mps"	):
           A_       :					str        =      torch.manual_seed(snake_case__	)
        else:
           A_       :					int        =      torch.Generator(device=snake_case__	).manual_seed(snake_case__	)
        A_       :					List[str]        =      {
            "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							:List[str]	):





        '''simple docstring'''




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

        A_       :					int        =      self.get_dummy_components()
        A_       :					Optional[int]        =      self.pipeline_class(**snake_case__	)
        pipe.to(snake_case__	)
        pipe.set_progress_bar_config(disable=snake_case__	)

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

        A_       :					Dict        =      self.get_dummy_inputs(snake_case__	)
        A_       :					int        =      pipe(**snake_case__	)[0]

        with tempfile.TemporaryDirectory() as tmpdir:
           pipe.save_pretrained(snake_case__	)
           A_       :					Any        =      self.pipeline_class.from_pretrained(snake_case__	)
           pipe_loaded.to(snake_case__	)
           pipe_loaded.set_progress_bar_config(disable=snake_case__	)

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

        A_       :					Optional[int]        =      self.get_dummy_inputs(snake_case__	)
        A_       :					int        =      pipe_loaded(**snake_case__	)[0]

        A_       :					int        =      np.abs(output - output_loaded	).max()
        self.assertLess(snake_case__      ,   1e-4	)



     def     SCREAMING_SNAKE_CASE		(				self							:str	):





        '''simple docstring'''




        A_       :					Union[str, Any]        =      "cpu"

        A_       :					int        =      self.get_dummy_components()
        A_       :					Any        =      self.pipeline_class(**snake_case__	)
        pipe.to(snake_case__	)
        pipe.set_progress_bar_config(disable=snake_case__	)

        A_       :					Any        =      self.get_dummy_mask_inputs(snake_case__	)
        A_       :					Union[str, Any]        =      pipe.generate_mask(**snake_case__	)
        A_       :					List[Any]        =      mask[0, -3:, -3:]

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



     def     SCREAMING_SNAKE_CASE		(				self							:str	):





        '''simple docstring'''




        A_       :					Any        =      "cpu"

        A_       :					List[Any]        =      self.get_dummy_components()
        A_       :					Optional[int]        =      self.pipeline_class(**snake_case__	)
        pipe.to(snake_case__	)
        pipe.set_progress_bar_config(disable=snake_case__	)

        A_       :					List[str]        =      self.get_dummy_inversion_inputs(snake_case__	)
        A_       :					List[str]        =      pipe.invert(**snake_case__	).images
        A_       :					Optional[Any]        =      image[0, -1, -3:, -3:]

        self.assertEqual(image.shape      ,   (2, 32, 32, 3)	)
        A_       :					int        =      np.array(
            [0.5150, 0.5134, 0.5043, 0.5376, 0.4694, 0.51050, 0.5015, 0.4407, 0.4799]      ,   )
        A_       :					Optional[Any]        =      np.abs(image_slice.flatten() - expected_slice	).max()
        self.assertLessEqual(snake_case__      ,   1e-3	)



     def     SCREAMING_SNAKE_CASE		(				self							:int	):





        '''simple docstring'''




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





     def     SCREAMING_SNAKE_CASE		(				self							:str	):





        '''simple docstring'''




        A_       :					Any        =      "cpu"

        A_       :					int        =      self.get_dummy_components()

        A_       :					Tuple        =      {"beta_start": 0.00085, "beta_end": 0.012, "beta_schedule": "scaled_linear"}
        A_       :					int        =      DPMSolverMultistepScheduler(**snake_case__	)
        A_       :					Optional[Any]        =      DPMSolverMultistepInverseScheduler(**snake_case__	)

        A_       :					Any        =      self.pipeline_class(**snake_case__	)
        pipe.to(snake_case__	)
        pipe.set_progress_bar_config(disable=snake_case__	)

        A_       :					Tuple        =      self.get_dummy_inversion_inputs(snake_case__	)
        A_       :					Optional[int]        =      pipe.invert(**snake_case__	).images
        A_       :					Optional[Any]        =      image[0, -1, -3:, -3:]

        self.assertEqual(image.shape      ,   (2, 32, 32, 3)	)
        A_       :					int        =      np.array(
            [0.5150, 0.5134, 0.5043, 0.5376, 0.4694, 0.51050, 0.5015, 0.4407, 0.4799]      ,   )
        A_       :					Any        =      np.abs(image_slice.flatten() - expected_slice	).max()
        self.assertLessEqual(snake_case__      ,   1e-3	)




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





     """simple docstring"""



     def     SCREAMING_SNAKE_CASE		(				self							:List[Any]	):





        '''simple docstring'''




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



     @classmethod
     def     SCREAMING_SNAKE_CASE		(				cls							:Optional[int]	):





        '''simple docstring'''




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

        A_       :					Tuple        =      raw_image.convert("RGB"	).resize((768, 768)	)

        A_       :					Tuple        =      raw_image



     def     SCREAMING_SNAKE_CASE		(				self							:Union[str, Any]	):





        '''simple docstring'''




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

        A_       :					List[str]        =      StableDiffusionDiffEditPipeline.from_pretrained(
            "stabilityai/stable-diffusion-2-1"      ,   safety_checker=snake_case__      ,   torch_dtype=torch.floataa	)
        A_       :					int        =      DDIMScheduler.from_config(pipe.scheduler.config	)
        A_       :					Dict        =      DDIMInverseScheduler.from_config(pipe.scheduler.config	)
        pipe.enable_model_cpu_offload()
        pipe.set_progress_bar_config(disable=snake_case__	)

        A_       :					Union[str, Any]        =      "a bowl of fruit"
        A_       :					int        =      "a bowl of pears"

        A_       :					List[Any]        =      pipe.generate_mask(
            image=self.raw_image      ,   source_prompt=snake_case__      ,   target_prompt=snake_case__      ,   generator=snake_case__      ,   )

        A_       :					Optional[int]        =      pipe.invert(
            prompt=snake_case__      ,   image=self.raw_image      ,   inpaint_strength=0.7      ,   generator=snake_case__	).latents

        A_       :					List[str]        =      pipe(
            prompt=snake_case__      ,   mask_image=snake_case__      ,   image_latents=snake_case__      ,   generator=snake_case__      ,   negative_prompt=snake_case__      ,   inpaint_strength=0.7      ,   output_type="numpy"      ,   ).images[0]

        A_       :					Dict        =      (
            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							:Union[str, Any]	):





        '''simple docstring'''




        A_       :					int        =      torch.manual_seed(0	)

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

        A_       :					str        =      "a bowl of fruit"
        A_       :					Optional[int]        =      "a bowl of pears"

        A_       :					Optional[Any]        =      pipe.generate_mask(
            image=self.raw_image      ,   source_prompt=snake_case__      ,   target_prompt=snake_case__      ,   generator=snake_case__      ,   )

        A_       :					List[str]        =      pipe.invert(
            prompt=snake_case__      ,   image=self.raw_image      ,   inpaint_strength=0.7      ,   generator=snake_case__      ,   num_inference_steps=25      ,   ).latents

        A_       :					Dict        =      pipe(
            prompt=snake_case__      ,   mask_image=snake_case__      ,   image_latents=snake_case__      ,   generator=snake_case__      ,   negative_prompt=snake_case__      ,   inpaint_strength=0.7      ,   num_inference_steps=25      ,   output_type="numpy"      ,   ).images[0]

        A_       :					int        =      (
            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 ..utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_scipy_available,
    is_torch_available,
    is_torchsde_available,
)


try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ..utils.dummy_pt_objects import *  # noqa F403
else:
   from .scheduling_consistency_models import CMStochasticIterativeScheduler
   from .scheduling_ddim import DDIMScheduler
   from .scheduling_ddim_inverse import DDIMInverseScheduler
   from .scheduling_ddim_parallel import DDIMParallelScheduler
   from .scheduling_ddpm import DDPMScheduler
   from .scheduling_ddpm_parallel import DDPMParallelScheduler
   from .scheduling_deis_multistep import DEISMultistepScheduler
   from .scheduling_dpmsolver_multistep import DPMSolverMultistepScheduler
   from .scheduling_dpmsolver_multistep_inverse import DPMSolverMultistepInverseScheduler
   from .scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
   from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler
   from .scheduling_euler_discrete import EulerDiscreteScheduler
   from .scheduling_heun_discrete import HeunDiscreteScheduler
   from .scheduling_ipndm import IPNDMScheduler
   from .scheduling_k_dpm_2_ancestral_discrete import KDPMaAncestralDiscreteScheduler
   from .scheduling_k_dpm_2_discrete import KDPMaDiscreteScheduler
   from .scheduling_karras_ve import KarrasVeScheduler
   from .scheduling_pndm import PNDMScheduler
   from .scheduling_repaint import RePaintScheduler
   from .scheduling_sde_ve import ScoreSdeVeScheduler
   from .scheduling_sde_vp import ScoreSdeVpScheduler
   from .scheduling_unclip import UnCLIPScheduler
   from .scheduling_unipc_multistep import UniPCMultistepScheduler
   from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
   from .scheduling_vq_diffusion import VQDiffusionScheduler

try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ..utils.dummy_flax_objects import *  # noqa F403
else:
   from .scheduling_ddim_flax import FlaxDDIMScheduler
   from .scheduling_ddpm_flax import FlaxDDPMScheduler
   from .scheduling_dpmsolver_multistep_flax import FlaxDPMSolverMultistepScheduler
   from .scheduling_karras_ve_flax import FlaxKarrasVeScheduler
   from .scheduling_lms_discrete_flax import FlaxLMSDiscreteScheduler
   from .scheduling_pndm_flax import FlaxPNDMScheduler
   from .scheduling_sde_ve_flax import FlaxScoreSdeVeScheduler
   from .scheduling_utils_flax import (
       FlaxKarrasDiffusionSchedulers,
       FlaxSchedulerMixin,
       FlaxSchedulerOutput,
       broadcast_to_shape_from_left,
   )


try:
   if not (is_torch_available() and is_scipy_available()):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ..utils.dummy_torch_and_scipy_objects import *  # noqa F403
else:
   from .scheduling_lms_discrete import LMSDiscreteScheduler

try:
   if not (is_torch_available() and is_torchsde_available()):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ..utils.dummy_torch_and_torchsde_objects import *  # noqa F403
else:
   from .scheduling_dpmsolver_sde import DPMSolverSDEScheduler





