


'''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					):
	A_				=    StableDiffusionDiffEditPipeline
	A_				=    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"height", "width", "image"} | {"image_latents"}
	A_				=    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {"image"} | {"image_latents"}
	A_				=    frozenset(
	    []					)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
	A_				=    frozenset([]					)




	def       __UpperCAmelCase							(   self							):

						'''simple docstring'''





						torch.manual_seed(0							)
						__a			:       List[str]					    =  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			:       Tuple					    =  DDIMScheduler(
						    beta_start=0.00085						,				beta_end=0.012						,				beta_schedule='scaled_linear'						,				clip_sample=__a						,				set_alpha_to_one=__a						,				)
						__a			:       Any					    =  DDIMInverseScheduler(
						    beta_start=0.00085						,				beta_end=0.012						,				beta_schedule='scaled_linear'						,				clip_sample=__a						,				set_alpha_to_zero=__a						,				)
						torch.manual_seed(0							)
						__a			:       Union[str, 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-0_5						,				num_attention_heads=4						,				num_hidden_layers=5						,				pad_token_id=1						,				vocab_size=1000						,				hidden_act='gelu'						,				projection_dim=512						,				)
						__a			:       List[Any]					    =  CLIPTextModel(__a							)
						__a			:       Any					    =  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'							)

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

						'''simple docstring'''





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

						'''simple docstring'''





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

						__a			:       Tuple					    =  self.get_dummy_components()
						__a			:       List[Any]					    =  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			:       Optional[Any]					    =  self.get_dummy_inputs(__a							)
						__a			:       str					    =  pipe(**__a							)[0]

						with tempfile.TemporaryDirectory() as tmpdir:
											pipe.save_pretrained(__a							)
											__a			:       Any					    =  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			:       str					    =  self.get_dummy_inputs(__a							)
						__a			:       Tuple					    =  pipe_loaded(**__a							)[0]

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




	def       __UpperCAmelCase							(   self							):

						'''simple docstring'''





						__a			:       Tuple					    =  'cpu'

						__a			:       Union[str, Any]					    =  self.get_dummy_components()
						__a			:       Optional[Any]					    =  self.pipeline_class(**__a							)
						pipe.to(__a							)
						pipe.set_progress_bar_config(disable=__a							)

						__a			:       Any					    =  self.get_dummy_mask_inputs(__a							)
						__a			:       Dict					    =  pipe.generate_mask(**__a							)
						__a			:       Optional[int]					    =  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(__a						,				1E-3							)
						self.assertEqual(mask[0, -3, -4]						,				0							)




	def       __UpperCAmelCase							(   self							):

						'''simple docstring'''





						__a			:       Any					    =  'cpu'

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

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

						self.assertEqual(image.shape						,				(2, 32, 32, 3)							)
						__a			:       Optional[Any]					    =  np.array(
						    [0.5150, 0.5134, 0.5043, 0.5376, 0.4694, 0.51050, 0.5015, 0.4407, 0.4799]						,				)
						__a			:       str					    =  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			:       Any					    =  'cpu'

						__a			:       str					    =  self.get_dummy_components()

						__a			:       str					    =  {'beta_start': 0.00085, 'beta_end': 0.012, 'beta_schedule': 'scaled_linear'}
						__a			:       Any					    =  DPMSolverMultistepScheduler(**__a							)
						__a			:       Any					    =  DPMSolverMultistepInverseScheduler(**__a							)

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

						__a			:       str					    =  self.get_dummy_inversion_inputs(__a							)
						__a			:       List[Any]					    =  pipe.invert(**__a							).images
						__a			:       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			:       Union[str, Any]					    =  np.abs(image_slice.flatten() - expected_slice							).max()
						self.assertLessEqual(__a						,				1E-3							)





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




	def       __UpperCAmelCase							(   self							):

						'''simple docstring'''





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




	@classmethod
	def       __UpperCAmelCase							(   cls							):

						'''simple docstring'''





						__a			:       Optional[int]					    =  load_image(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png'							)

						__a			:       str					    =  raw_image.convert('RGB'							).resize((768, 768)							)

						__a			:       str					    =  raw_image




	def       __UpperCAmelCase							(   self							):

						'''simple docstring'''





						__a			:       Dict					    =  torch.manual_seed(0							)

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

						__a			:       int					    =  'a bowl of fruit'
						__a			:       Dict					    =  'a bowl of pears'

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

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

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

						__a			:       str					    =  (
						    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			:       Any					    =  torch.manual_seed(0							)

						__a			:       List[str]					    =  StableDiffusionDiffEditPipeline.from_pretrained(
						    'stabilityai/stable-diffusion-2-1'						,				safety_checker=__a						,				torch_dtype=torch.floataa							)
						__a			:       Union[str, Any]					    =  DPMSolverMultistepScheduler.from_config(pipe.scheduler.config							)
						__a			:       List[str]					    =  DPMSolverMultistepInverseScheduler.from_config(pipe.scheduler.config							)
						pipe.enable_model_cpu_offload()
						pipe.set_progress_bar_config(disable=__a							)

						__a			:       Dict					    =  'a bowl of fruit'
						__a			:       Optional[Any]					    =  'a bowl of pears'

						__a			:       Optional[int]					    =  pipe.generate_mask(
						    image=self.raw_image						,				source_prompt=__a						,				target_prompt=__a						,				generator=__a						,				)

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

						__a			:       List[str]					    =  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			:       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







'''simple docstring'''


import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.text import TextDatasetReader

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase      ):
		assert isinstance(__lowerCAmelCase					,       __lowerCAmelCase      )
		assert dataset.num_rows == 4
		assert dataset.num_columns == 1
		assert dataset.column_names == ["text"]
		for feature, expected_dtype in expected_features.items():
				assert dataset.features[feature].dtype == expected_dtype






@pytest.mark.parametrize("keep_in_memory"					,       [False, True]      )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				Union[str, Any]	    =					tmp_path / "cache"
		_UpperCAmelCase						:				int	    =					{"text": "string"}
		with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
				_UpperCAmelCase						:				Union[str, Any]	    =					TextDatasetReader(__lowerCAmelCase					,       cache_dir=__lowerCAmelCase					,       keep_in_memory=__lowerCAmelCase      ).read()
		_check_text_dataset(__lowerCAmelCase					,       __lowerCAmelCase      )






@pytest.mark.parametrize(
    "features"					,       [
        None,
        {"text": "string"},
        {"text": "int32"},
        {"text": "float32"},
    ]					,       )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				List[Any]	    =					tmp_path / "cache"
		_UpperCAmelCase						:				Any	    =					{"text": "string"}
		_UpperCAmelCase						:				Optional[Any]	    =					features.copy() if features else default_expected_features
		_UpperCAmelCase						:				Union[str, Any]	    =					(
		    Features({feature: Value(__lowerCAmelCase      ) for feature, dtype in features.items()}      ) if features is not None else None
		)
		_UpperCAmelCase						:				Union[str, Any]	    =					TextDatasetReader(__lowerCAmelCase					,       features=__lowerCAmelCase					,       cache_dir=__lowerCAmelCase      ).read()
		_check_text_dataset(__lowerCAmelCase					,       __lowerCAmelCase      )






@pytest.mark.parametrize("split"					,       [None, NamedSplit("train"      ), "train", "test"]      )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				Any	    =					tmp_path / "cache"
		_UpperCAmelCase						:				Dict	    =					{"text": "string"}
		_UpperCAmelCase						:				Union[str, Any]	    =					TextDatasetReader(__lowerCAmelCase					,       cache_dir=__lowerCAmelCase					,       split=__lowerCAmelCase      ).read()
		_check_text_dataset(__lowerCAmelCase					,       __lowerCAmelCase      )
		assert dataset.split == split if split else "train"






@pytest.mark.parametrize("path_type"					,       [str, list]      )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		if issubclass(__lowerCAmelCase					,       __lowerCAmelCase      ):
				_UpperCAmelCase						:				Union[str, Any]	    =					text_path
		elif issubclass(__lowerCAmelCase					,       __lowerCAmelCase      ):
				_UpperCAmelCase						:				int	    =					[text_path]
		_UpperCAmelCase						:				List[Any]	    =					tmp_path / "cache"
		_UpperCAmelCase						:				Union[str, Any]	    =					{"text": "string"}
		_UpperCAmelCase						:				Optional[int]	    =					TextDatasetReader(__lowerCAmelCase					,       cache_dir=__lowerCAmelCase      ).read()
		_check_text_dataset(__lowerCAmelCase					,       __lowerCAmelCase      )






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase=("train",)      ):
		assert isinstance(__lowerCAmelCase					,       __lowerCAmelCase      )
		for split in splits:
				_UpperCAmelCase						:				List[str]	    =					dataset_dict[split]
				assert dataset.num_rows == 4
				assert dataset.num_columns == 1
				assert dataset.column_names == ["text"]
				for feature, expected_dtype in expected_features.items():
						assert dataset.features[feature].dtype == expected_dtype






@pytest.mark.parametrize("keep_in_memory"					,       [False, True]      )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				List[Any]	    =					tmp_path / "cache"
		_UpperCAmelCase						:				Tuple	    =					{"text": "string"}
		with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
				_UpperCAmelCase						:				Any	    =					TextDatasetReader({"train": text_path}					,       cache_dir=__lowerCAmelCase					,       keep_in_memory=__lowerCAmelCase      ).read()
		_check_text_datasetdict(__lowerCAmelCase					,       __lowerCAmelCase      )






@pytest.mark.parametrize(
    "features"					,       [
        None,
        {"text": "string"},
        {"text": "int32"},
        {"text": "float32"},
    ]					,       )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				List[Any]	    =					tmp_path / "cache"
		# CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
		_UpperCAmelCase						:				List[Any]	    =					{"text": "string"}
		_UpperCAmelCase						:				List[str]	    =					features.copy() if features else default_expected_features
		_UpperCAmelCase						:				Optional[int]	    =					(
		    Features({feature: Value(__lowerCAmelCase      ) for feature, dtype in features.items()}      ) if features is not None else None
		)
		_UpperCAmelCase						:				Tuple	    =					TextDatasetReader({"train": text_path}					,       features=__lowerCAmelCase					,       cache_dir=__lowerCAmelCase      ).read()
		_check_text_datasetdict(__lowerCAmelCase					,       __lowerCAmelCase      )






@pytest.mark.parametrize("split"					,       [None, NamedSplit("train"      ), "train", "test"]      )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		if split:
				_UpperCAmelCase						:				int	    =					{split: text_path}
		else:
				_UpperCAmelCase						:				Tuple	    =					"train"
				_UpperCAmelCase						:				List[str]	    =					{"train": text_path, "test": text_path}
		_UpperCAmelCase						:				Optional[Any]	    =					tmp_path / "cache"
		_UpperCAmelCase						:				Optional[int]	    =					{"text": "string"}
		_UpperCAmelCase						:				int	    =					TextDatasetReader(__lowerCAmelCase					,       cache_dir=__lowerCAmelCase      ).read()
		_check_text_datasetdict(__lowerCAmelCase					,       __lowerCAmelCase					,       splits=list(path.keys()      )      )
		assert all(dataset[split].split == split for split in path.keys()      )



