






"""simple docstring"""






from ...configuration_utils import PretrainedConfig
from ...utils import logging


SCREAMING_SNAKE_CASE						:						Any					   =       logging.get_logger(__name__)

SCREAMING_SNAKE_CASE						:						Any					   =       {
    '''funnel-transformer/small''': '''https://huggingface.co/funnel-transformer/small/resolve/main/config.json''',
    '''funnel-transformer/small-base''': '''https://huggingface.co/funnel-transformer/small-base/resolve/main/config.json''',
    '''funnel-transformer/medium''': '''https://huggingface.co/funnel-transformer/medium/resolve/main/config.json''',
    '''funnel-transformer/medium-base''': '''https://huggingface.co/funnel-transformer/medium-base/resolve/main/config.json''',
    '''funnel-transformer/intermediate''': (
        '''https://huggingface.co/funnel-transformer/intermediate/resolve/main/config.json'''
    ),
    '''funnel-transformer/intermediate-base''': (
        '''https://huggingface.co/funnel-transformer/intermediate-base/resolve/main/config.json'''
    ),
    '''funnel-transformer/large''': '''https://huggingface.co/funnel-transformer/large/resolve/main/config.json''',
    '''funnel-transformer/large-base''': '''https://huggingface.co/funnel-transformer/large-base/resolve/main/config.json''',
    '''funnel-transformer/xlarge''': '''https://huggingface.co/funnel-transformer/xlarge/resolve/main/config.json''',
    '''funnel-transformer/xlarge-base''': '''https://huggingface.co/funnel-transformer/xlarge-base/resolve/main/config.json''',
}





class       snake_case_       (    _lowerCamelCase       ):



	"""simple docstring"""

	SCREAMING_SNAKE_CASE_:							int   			=					"""funnel"""
	SCREAMING_SNAKE_CASE_:							Any   			=					{
	    """hidden_size""": """d_model""",
	    """num_attention_heads""": """n_head""",
	}


	def __init__(      self  ,					__a=3_0522  ,					__a=[4, 4, 4]  ,					__a=None  ,					__a=2  ,					__a=768  ,					__a=12  ,					__a=64  ,					__a=3072  ,					__a="gelu_new"  ,					__a=0.1  ,					__a=0.1  ,					__a=0.0  ,					__a=0.1  ,					__a=None  ,					__a=1E-9  ,					__a="mean"  ,					__a="relative_shift"  ,					__a=True  ,					__a=True  ,					__a=True  ,					**__a  ,					):


					"""simple docstring"""

					A__							 =       vocab_size
					A__							 =       block_sizes
					A__							 =       [1] * len(__a						) if block_repeats is None else block_repeats
					assert len(__a						) == len(
					    self.block_repeats						), "`block_sizes` and `block_repeats` should have the same length."
					A__							 =       num_decoder_layers
					A__							 =       d_model
					A__							 =       n_head
					A__							 =       d_head
					A__							 =       d_inner
					A__							 =       hidden_act
					A__							 =       hidden_dropout
					A__							 =       attention_dropout
					A__							 =       activation_dropout
					A__							 =       initializer_range
					A__							 =       initializer_std
					A__							 =       layer_norm_eps
					assert pooling_type in [
					    "mean",
					    "max",
					], f'''Got {pooling_type} for `pooling_type` but only \'mean\' and \'max\' are supported.'''
					A__							 =       pooling_type
					assert attention_type in [
					    "relative_shift",
					    "factorized",
					], f'''Got {attention_type} for `attention_type` but only \'relative_shift\' and \'factorized\' are supported.'''
					A__							 =       attention_type
					A__							 =       separate_cls
					A__							 =       truncate_seq
					A__							 =       pool_q_only

					super().__init__(**__a						)


	@property
	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					return sum(self.block_sizes						)


	@num_hidden_layers.setter
	def       _UpperCAmelCase						(      self  ,					__a						):


					"""simple docstring"""

					raise NotImplementedError(
					    'This model does not support the setting of `num_hidden_layers`. Please set `block_sizes`.'						)


	@property
	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					return len(self.block_sizes						)






	@num_blocks.setter
	def       _UpperCAmelCase						(      self  ,					__a						):


					"""simple docstring"""

					raise NotImplementedError('This model does not support the setting of `num_blocks`. Please set `block_sizes`.'						)











"""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       snake_case_       (    _lowerCamelCase    ,							_lowerCamelCase    ,							unittest.TestCase       ):



	"""simple docstring"""

	SCREAMING_SNAKE_CASE_:							Union[str, Any]   			=					StableDiffusionDiffEditPipeline
	SCREAMING_SNAKE_CASE_:							Union[str, Any]   			=					TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"""height""", """width""", """image"""} | {"""image_latents"""}
	SCREAMING_SNAKE_CASE_:							int   			=					TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {"""image"""} | {"""image_latents"""}
	SCREAMING_SNAKE_CASE_:							str   			=					frozenset(
	    []       )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
	SCREAMING_SNAKE_CASE_:							int   			=					frozenset([]       )


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					torch.manual_seed(0						)
					A__							 =       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  ,					)
					A__							 =       DDIMScheduler(
					    beta_start=0.0_0085  ,					beta_end=0.012  ,					beta_schedule='scaled_linear'  ,					clip_sample=__a  ,					set_alpha_to_one=__a  ,					)
					A__							 =       DDIMInverseScheduler(
					    beta_start=0.0_0085  ,					beta_end=0.012  ,					beta_schedule='scaled_linear'  ,					clip_sample=__a  ,					set_alpha_to_zero=__a  ,					)
					torch.manual_seed(0						)
					A__							 =       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__							 =       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  ,					)
					A__							 =       CLIPTextModel(__a						)
					A__							 =       CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)

					A__							 =       {
					    '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       _UpperCAmelCase						(      self  ,					__a  ,					__a=0						):


					"""simple docstring"""

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


					"""simple docstring"""

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


					"""simple docstring"""

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


					"""simple docstring"""

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

					A__							 =       self.get_dummy_components()
					A__							 =       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}						)

					A__							 =       self.get_dummy_inputs(__a						)
					A__							 =       pipe(**__a						)[0]

					with tempfile.TemporaryDirectory() as tmpdir:
									pipe.save_pretrained(__a						)
									A__							 =       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.'''  ,					)

					A__							 =       self.get_dummy_inputs(__a						)
					A__							 =       pipe_loaded(**__a						)[0]

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


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       'cpu'

					A__							 =       self.get_dummy_components()
					A__							 =       self.pipeline_class(**__a						)
					pipe.to(__a						)
					pipe.set_progress_bar_config(disable=__a						)

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

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


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       'cpu'

					A__							 =       self.get_dummy_components()
					A__							 =       self.pipeline_class(**__a						)
					pipe.to(__a						)
					pipe.set_progress_bar_config(disable=__a						)

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

					self.assertEqual(image.shape  ,					(2, 32, 32, 3)						)
					A__							 =       np.array(
					    [0.5150, 0.5134, 0.5043, 0.5376, 0.4694, 0.5_1050, 0.5015, 0.4407, 0.4799]  ,					)
					A__							 =       np.abs(image_slice.flatten() - expected_slice						).max()
					self.assertLessEqual(__a  ,					1E-3						)


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

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






	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       'cpu'

					A__							 =       self.get_dummy_components()

					A__							 =       {'beta_start': 0.0_0085, 'beta_end': 0.012, 'beta_schedule': 'scaled_linear'}
					A__							 =       DPMSolverMultistepScheduler(**__a						)
					A__							 =       DPMSolverMultistepInverseScheduler(**__a						)

					A__							 =       self.pipeline_class(**__a						)
					pipe.to(__a						)
					pipe.set_progress_bar_config(disable=__a						)

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

					self.assertEqual(image.shape  ,					(2, 32, 32, 3)						)
					A__							 =       np.array(
					    [0.5150, 0.5134, 0.5043, 0.5376, 0.4694, 0.5_1050, 0.5015, 0.4407, 0.4799]  ,					)
					A__							 =       np.abs(image_slice.flatten() - expected_slice						).max()
					self.assertLessEqual(__a  ,					1E-3						)





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



	"""simple docstring"""

	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

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


	@classmethod
	def       _UpperCAmelCase						(      cls						):


					"""simple docstring"""

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

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

					A__							 =       raw_image


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       torch.manual_seed(0						)

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

					A__							 =       'a bowl of fruit'
					A__							 =       'a bowl of pears'

					A__							 =       pipe.generate_mask(
					    image=self.raw_image  ,					source_prompt=__a  ,					target_prompt=__a  ,					generator=__a  ,					)

					A__							 =       pipe.invert(
					    prompt=__a  ,					image=self.raw_image  ,					inpaint_strength=0.7  ,					generator=__a						).latents

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

					A__							 =       (
					    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       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       torch.manual_seed(0						)

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

					A__							 =       'a bowl of fruit'
					A__							 =       'a bowl of pears'

					A__							 =       pipe.generate_mask(
					    image=self.raw_image  ,					source_prompt=__a  ,					target_prompt=__a  ,					generator=__a  ,					)

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

					A__							 =       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]

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


