



"""simple docstring"""







import math
from collections import OrderedDict

import torch
from packaging import version
from torch import Tensor, nn

from .utils import logging


UpperCAmelCase:   Optional[Any]       			=							logging.get_logger(__name__)





class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def __init__(			self					):
			super().__init__()
			if version.parse(torch.__version__					) < version.parse("""1.12.0"""					):
					raise ImportError(
					    f"""You are using torch=={torch.__version__}, but torch>=1.12.0 is required to use """
					    """PytorchGELUTanh. Please upgrade torch."""					)






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return nn.functional.gelu(UpperCAmelCase_							,approximate="""tanh"""					)







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi					) * (input + 0.044715 * torch.pow(UpperCAmelCase_							,3.0					))					))







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def __init__(			self							,UpperCAmelCase_ = False					):
			super().__init__()
			if use_gelu_python:
					_lowercase							:      str			=					self._gelu_python
			else:
					_lowercase							:      List[str]			=					nn.functional.gelu






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return input * 0.5 * (1.0 + torch.erf(input / math.sqrt(2.0					)					))






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return self.act(UpperCAmelCase_					)







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return 0.5 * input * (1.0 + torch.tanh(input * 0.7978845608 * (1.0 + 0.044715 * input * input)					))







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return input * torch.sigmoid(1.702 * input					)







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def __init__(			self							,UpperCAmelCase_							,UpperCAmelCase_					):
			if min > max:
					raise ValueError(f"""min should be < max (got min: {min}, max: {max})"""					)

			super().__init__()
			_lowercase							:      Tuple			=					min
			_lowercase							:      List[Any]			=					max






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return torch.clip(gelu(UpperCAmelCase_					)							,self.min							,self.max					)







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def __init__(			self					):
			super().__init__()
			_lowercase							:      Dict			=					math.sqrt(2 / math.pi					)






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return 0.5 * input * (1 + torch.tanh(self.precomputed_constant * (input + 0.044715 * torch.pow(UpperCAmelCase_							,3					))					))







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return nn.functional.silu(UpperCAmelCase_					)







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def __init__(			self					):
			super().__init__()
			if version.parse(torch.__version__					) < version.parse("""1.9.0"""					):
					_lowercase							:      str			=					self._mish_python
			else:
					_lowercase							:      str			=					nn.functional.mish






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return input * torch.tanh(nn.functional.softplus(UpperCAmelCase_					)					)






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return self.act(UpperCAmelCase_					)







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return input







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_							,UpperCAmelCase_=0.707107							,UpperCAmelCase_=0.282095					):
			_lowercase							:      List[Any]			=					(input - mu).div(sigma * math.sqrt(2.0					)					)
			return 0.5 * (1.0 + torch.erf(UpperCAmelCase_					))







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			_lowercase							:      Any			=					nn.functional.relu(UpperCAmelCase_					)
			_lowercase							:      Tuple			=					torch.square(UpperCAmelCase_					)
			return squared







class    UpperCamelCase     (			snake_case    ):




	"""simple docstring"""






	def __getitem__(			self							,UpperCAmelCase_					):
			_lowercase							:      Any			=					super().__getitem__(UpperCAmelCase_					)
			_lowercase			,   _lowercase							:      Optional[int]			=					content if isinstance(UpperCAmelCase_							,UpperCAmelCase_					) else (content, {})
			return cls(**UpperCAmelCase_					)







UpperCAmelCase:   List[str]       			=							{
    """gelu""": GELUActivation,
    """gelu_10""": (ClippedGELUActivation, {"""min""": -10, """max""": 10}),
    """gelu_fast""": FastGELUActivation,
    """gelu_new""": NewGELUActivation,
    """gelu_python""": (GELUActivation, {"""use_gelu_python""": True}),
    """gelu_pytorch_tanh""": PytorchGELUTanh,
    """gelu_accurate""": AccurateGELUActivation,
    """laplace""": LaplaceActivation,
    """linear""": LinearActivation,
    """mish""": MishActivation,
    """quick_gelu""": QuickGELUActivation,
    """relu""": nn.ReLU,
    """relu2""": ReLUSquaredActivation,
    """relu6""": nn.ReLUa,
    """sigmoid""": nn.Sigmoid,
    """silu""": SiLUActivation,
    """swish""": SiLUActivation,
    """tanh""": nn.Tanh,
}
UpperCAmelCase:   Any       			=							ClassInstantier(ACTaCLS)


def       __SCREAMING_SNAKE_CASE      (   __UpperCAmelCase	):
		if activation_string in ACTaFN:
				return ACTaFN[activation_string]
		else:
				raise KeyError(F"""function {activation_string} not found in ACT2FN mapping {list(ACTaFN.keys()	)}"""	)


# For backwards compatibility with: from activations import gelu_python
UpperCAmelCase:   Any       			=							get_activation("""gelu_python""")
UpperCAmelCase:   int       			=							get_activation("""gelu_new""")
UpperCAmelCase:   List[Any]       			=							get_activation("""gelu""")
UpperCAmelCase:   Union[str, Any]       			=							get_activation("""gelu_fast""")
UpperCAmelCase:   int       			=							get_activation("""quick_gelu""")
UpperCAmelCase:   List[str]       			=							get_activation("""silu""")
UpperCAmelCase:   Tuple       			=							get_activation("""mish""")
UpperCAmelCase:   Optional[int]       			=							get_activation("""linear""")





"""simple docstring"""

import gc
import unittest

import numpy as np
import torch

from diffusers import RePaintPipeline, RePaintScheduler, UNetaDModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_image,
    load_numpy,
    nightly,
    require_torch_gpu,
    skip_mps,
    torch_device,
)

from ..pipeline_params import IMAGE_INPAINTING_BATCH_PARAMS, IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()






class        _lowercase   (				__UpperCAmelCase				,       unittest.TestCase		):
	lowercase_        =	RePaintPipeline
	lowercase_        =	IMAGE_INPAINTING_PARAMS - {'width', 'height', 'guidance_scale'}
	lowercase_        =	PipelineTesterMixin.required_optional_params - {
	    'latents',
	    'num_images_per_prompt',
	    'callback',
	    'callback_steps',
	}
	lowercase_        =	IMAGE_INPAINTING_BATCH_PARAMS
	lowercase_        =	False




	def 			_UpperCamelCase    ( self					)		->							Optional[Any]:
			torch.manual_seed(0					)
			torch.manual_seed(0					)
			lowerCamelCase     :			Optional[Any]		=    UNetaDModel(
			    block_out_channels=(32, 64)		, layers_per_block=2		, sample_size=32		, in_channels=3		, out_channels=3		, down_block_types=('DownBlock2D', 'AttnDownBlock2D')		, up_block_types=('AttnUpBlock2D', 'UpBlock2D')		, )
			lowerCamelCase     :			Union[str, Any]		=    RePaintScheduler()
			lowerCamelCase     :			Dict		=    {'unet': unet, 'scheduler': scheduler}
			return components




	def 			_UpperCamelCase    ( self		, UpperCAmelCase_		, UpperCAmelCase_=0					)		->							Optional[int]:
			if str(UpperCAmelCase_					).startswith('mps'					):
					lowerCamelCase     :			str		=    torch.manual_seed(UpperCAmelCase_					)
			else:
					lowerCamelCase     :			Tuple		=    torch.Generator(device=UpperCAmelCase_					).manual_seed(UpperCAmelCase_					)
			lowerCamelCase     :			str		=    np.random.RandomState(UpperCAmelCase_					).standard_normal((1, 3, 32, 32)					)
			lowerCamelCase     :			Tuple		=    torch.from_numpy(UpperCAmelCase_					).to(device=UpperCAmelCase_		, dtype=torch.floataa					)
			lowerCamelCase     :			Optional[int]		=    (image > 0).to(device=UpperCAmelCase_		, dtype=torch.floataa					)
			lowerCamelCase     :			Optional[Any]		=    {
			    'image': image,
			    'mask_image': mask,
			    'generator': generator,
			    'num_inference_steps': 5,
			    'eta': 0.0,
			    'jump_length': 2,
			    'jump_n_sample': 2,
			    'output_type': 'numpy',
			}
			return inputs




	def 			_UpperCamelCase    ( self					)		->							str:
			lowerCamelCase     :			Any		=    'cpu'  # ensure determinism for the device-dependent torch.Generator
			lowerCamelCase     :			Union[str, Any]		=    self.get_dummy_components()
			lowerCamelCase     :			List[Any]		=    RePaintPipeline(**UpperCAmelCase_					)
			lowerCamelCase     :			int		=    sd_pipe.to(UpperCAmelCase_					)
			sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_					)

			lowerCamelCase     :			Any		=    self.get_dummy_inputs(UpperCAmelCase_					)
			lowerCamelCase     :			Dict		=    sd_pipe(**UpperCAmelCase_					).images
			lowerCamelCase     :			List[Any]		=    image[0, -3:, -3:, -1]

			assert image.shape == (1, 32, 32, 3)
			lowerCamelCase     :			Optional[Any]		=    np.array([1.0000, 0.5426, 0.5497, 0.2200, 1.0000, 1.0000, 0.5623, 1.0000, 0.6274]					)

			assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-3




	@skip_mps
	def 			_UpperCamelCase    ( self					)		->							Union[str, Any]:
			return super().test_save_load_local()




	@unittest.skip('non-deterministic pipeline'					)
	def 			_UpperCamelCase    ( self					)		->							int:
			return super().test_inference_batch_single_identical()




	@skip_mps
	def 			_UpperCamelCase    ( self					)		->							int:
			return super().test_dict_tuple_outputs_equivalent()




	@skip_mps
	def 			_UpperCamelCase    ( self					)		->							Optional[int]:
			return super().test_save_load_optional_components()




	@skip_mps
	def 			_UpperCamelCase    ( self					)		->							Tuple:
			return super().test_attention_slicing_forward_pass()






@nightly
@require_torch_gpu
class        _lowercase   (				unittest.TestCase		):




	def 			_UpperCamelCase    ( self					)		->							Union[str, Any]:
			super().tearDown()
			gc.collect()
			torch.cuda.empty_cache()




	def 			_UpperCamelCase    ( self					)		->							int:
			lowerCamelCase     :			Optional[int]		=    load_image(
			    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/'
			    'repaint/celeba_hq_256.png'					)
			lowerCamelCase     :			str		=    load_image(
			    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/mask_256.png'					)
			lowerCamelCase     :			Optional[Any]		=    load_numpy(
			    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/'
			    'repaint/celeba_hq_256_result.npy'					)

			lowerCamelCase     :			Any		=    'google/ddpm-ema-celebahq-256'
			lowerCamelCase     :			List[Any]		=    UNetaDModel.from_pretrained(UpperCAmelCase_					)
			lowerCamelCase     :			Any		=    RePaintScheduler.from_pretrained(UpperCAmelCase_					)

			lowerCamelCase     :			Any		=    RePaintPipeline(unet=UpperCAmelCase_		, scheduler=UpperCAmelCase_					).to(UpperCAmelCase_					)
			repaint.set_progress_bar_config(disable=UpperCAmelCase_					)
			repaint.enable_attention_slicing()

			lowerCamelCase     :			Dict		=    torch.manual_seed(0					)
			lowerCamelCase     :			str		=    repaint(
			    UpperCAmelCase_		, UpperCAmelCase_		, num_inference_steps=250		, eta=0.0		, jump_length=10		, jump_n_sample=10		, generator=UpperCAmelCase_		, output_type='np'		, )
			lowerCamelCase     :			Optional[int]		=    output.images[0]

			assert image.shape == (256, 256, 3)
			assert np.abs(expected_image - image					).mean() < 1E-2






