


import io
import json

import fsspec
import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.json import JsonDatasetReader, JsonDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases


def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__  ):
					assert isinstance(UpperCAmelCase__						,							UpperCAmelCase__  )
					assert dataset.num_rows == 4
					assert dataset.num_columns == 3
					assert dataset.column_names == ["col_1", "col_2", "col_3"]
					for feature, expected_dtype in expected_features.items():
										assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("""keep_in_memory"""						,							[False, True]  )
def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__  ):
					lowercase_  							=			tmp_path / """cache"""
					lowercase_  							=			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
					with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
										lowercase_  							=			JsonDatasetReader(UpperCAmelCase__						,							cache_dir=UpperCAmelCase__						,							keep_in_memory=UpperCAmelCase__  ).read()
					_check_json_dataset(UpperCAmelCase__						,							UpperCAmelCase__  )


@pytest.mark.parametrize(
    """features"""						,							[
        None,
        {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""},
        {"""col_1""": """string""", """col_2""": """string""", """col_3""": """string"""},
        {"""col_1""": """int32""", """col_2""": """int32""", """col_3""": """int32"""},
        {"""col_1""": """float32""", """col_2""": """float32""", """col_3""": """float32"""},
    ]						,							)
def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__  ):
					lowercase_  							=			tmp_path / """cache"""
					lowercase_  							=			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
					lowercase_  							=			features.copy() if features else default_expected_features
					lowercase_  							=			(
					    Features({feature: Value(UpperCAmelCase__  ) for feature, dtype in features.items()}  ) if features is not None else None
					)
					lowercase_  							=			JsonDatasetReader(UpperCAmelCase__						,							features=UpperCAmelCase__						,							cache_dir=UpperCAmelCase__  ).read()
					_check_json_dataset(UpperCAmelCase__						,							UpperCAmelCase__  )


@pytest.mark.parametrize(
    """features"""						,							[
        None,
        {"""col_3""": """float64""", """col_1""": """string""", """col_2""": """int64"""},
    ]						,							)
def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__  ):
					lowercase_  							=			tmp_path / """cache"""
					lowercase_  							=			{"""col_3""": """float64""", """col_1""": """string""", """col_2""": """int64"""}
					lowercase_  							=			features.copy() if features else default_expected_features
					lowercase_  							=			(
					    Features({feature: Value(UpperCAmelCase__  ) for feature, dtype in features.items()}  ) if features is not None else None
					)
					lowercase_  							=			JsonDatasetReader(UpperCAmelCase__						,							features=UpperCAmelCase__						,							cache_dir=UpperCAmelCase__  ).read()
					assert isinstance(UpperCAmelCase__						,							UpperCAmelCase__  )
					assert dataset.num_rows == 2
					assert dataset.num_columns == 3
					assert dataset.column_names == ["col_3", "col_1", "col_2"]
					for feature, expected_dtype in expected_features.items():
										assert dataset.features[feature].dtype == expected_dtype


def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__  ):
					# jsonl_312_path features are {"col_3": "float64", "col_1": "string", "col_2": "int64"}
					lowercase_  							=			{"""col_2""": """int64""", """col_3""": """float64""", """col_1""": """string"""}
					lowercase_  							=			features.copy()
					lowercase_  							=			(
					    Features({feature: Value(UpperCAmelCase__  ) for feature, dtype in features.items()}  ) if features is not None else None
					)
					lowercase_  							=			tmp_path / """cache"""
					lowercase_  							=			JsonDatasetReader(UpperCAmelCase__						,							features=UpperCAmelCase__						,							cache_dir=UpperCAmelCase__  ).read()
					assert isinstance(UpperCAmelCase__						,							UpperCAmelCase__  )
					assert dataset.num_rows == 2
					assert dataset.num_columns == 3
					assert dataset.column_names == ["col_2", "col_3", "col_1"]
					for feature, expected_dtype in expected_features.items():
										assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("""split"""						,							[None, NamedSplit("""train"""  ), """train""", """test"""]  )
def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__  ):
					lowercase_  							=			tmp_path / """cache"""
					lowercase_  							=			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
					lowercase_  							=			JsonDatasetReader(UpperCAmelCase__						,							cache_dir=UpperCAmelCase__						,							split=UpperCAmelCase__  ).read()
					_check_json_dataset(UpperCAmelCase__						,							UpperCAmelCase__  )
					assert dataset.split == split if split else "train"


@pytest.mark.parametrize("""path_type"""						,							[str, list]  )
def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__  ):
					if issubclass(UpperCAmelCase__						,							UpperCAmelCase__  ):
										lowercase_  							=			jsonl_path
					elif issubclass(UpperCAmelCase__						,							UpperCAmelCase__  ):
										lowercase_  							=			[jsonl_path]
					lowercase_  							=			tmp_path / """cache"""
					lowercase_  							=			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
					lowercase_  							=			JsonDatasetReader(UpperCAmelCase__						,							cache_dir=UpperCAmelCase__  ).read()
					_check_json_dataset(UpperCAmelCase__						,							UpperCAmelCase__  )


def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__=("train",)  ):
					assert isinstance(UpperCAmelCase__						,							UpperCAmelCase__  )
					for split in splits:
										lowercase_  							=			dataset_dict[split]
										assert dataset.num_rows == 4
										assert dataset.num_columns == 3
										assert dataset.column_names == ["col_1", "col_2", "col_3"]
										for feature, expected_dtype in expected_features.items():
															assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("""keep_in_memory"""						,							[False, True]  )
def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__  ):
					lowercase_  							=			tmp_path / """cache"""
					lowercase_  							=			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
					with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
										lowercase_  							=			JsonDatasetReader({"""train""": jsonl_path}						,							cache_dir=UpperCAmelCase__						,							keep_in_memory=UpperCAmelCase__  ).read()
					_check_json_datasetdict(UpperCAmelCase__						,							UpperCAmelCase__  )


@pytest.mark.parametrize(
    """features"""						,							[
        None,
        {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""},
        {"""col_1""": """string""", """col_2""": """string""", """col_3""": """string"""},
        {"""col_1""": """int32""", """col_2""": """int32""", """col_3""": """int32"""},
        {"""col_1""": """float32""", """col_2""": """float32""", """col_3""": """float32"""},
    ]						,							)
def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__  ):
					lowercase_  							=			tmp_path / """cache"""
					lowercase_  							=			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
					lowercase_  							=			features.copy() if features else default_expected_features
					lowercase_  							=			(
					    Features({feature: Value(UpperCAmelCase__  ) for feature, dtype in features.items()}  ) if features is not None else None
					)
					lowercase_  							=			JsonDatasetReader({"""train""": jsonl_path}						,							features=UpperCAmelCase__						,							cache_dir=UpperCAmelCase__  ).read()
					_check_json_datasetdict(UpperCAmelCase__						,							UpperCAmelCase__  )


@pytest.mark.parametrize("""split"""						,							[None, NamedSplit("""train"""  ), """train""", """test"""]  )
def 			UpperCAmelCase_				(      UpperCAmelCase__						,							UpperCAmelCase__						,							UpperCAmelCase__  ):
					if split:
										lowercase_  							=			{split: jsonl_path}
					else:
										lowercase_  							=			"""train"""
										lowercase_  							=			{"""train""": jsonl_path, """test""": jsonl_path}
					lowercase_  							=			tmp_path / """cache"""
					lowercase_  							=			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
					lowercase_  							=			JsonDatasetReader(UpperCAmelCase__						,							cache_dir=UpperCAmelCase__  ).read()
					_check_json_datasetdict(UpperCAmelCase__						,							UpperCAmelCase__						,							splits=list(path.keys()  )  )
					assert all(dataset[split].split == split for split in path.keys()  )


def 			UpperCAmelCase_				(      UpperCAmelCase__  ):
					return json.load(UpperCAmelCase__  )


def 			UpperCAmelCase_				(      UpperCAmelCase__  ):
					return [json.loads(UpperCAmelCase__  ) for line in buffer]







class        UpperCamelCase__   :

					@pytest.mark.parametrize("""lines, load_json_function"""							,					[(True, load_json_lines), (False, load_json)]       )
					def   UpperCAmelCase__       (     self       :   List[Any]							,					UpperCamelCase__       :   Optional[int]							,					UpperCamelCase__       :   Optional[int]							,					UpperCamelCase__       :   Optional[Any]       ):






										'''simple docstring'''






										with io.BytesIO() as buffer:
															JsonDatasetWriter(UpperCamelCase__							,					UpperCamelCase__							,					lines=UpperCamelCase__       ).write()
															buffer.seek(0       )
															lowercase_  							=			load_json_function(UpperCamelCase__       )
										assert isinstance(UpperCamelCase__							,					UpperCamelCase__       )
										assert isinstance(exported_content[0]							,					UpperCamelCase__       )
										assert len(UpperCamelCase__       ) == 10

					@pytest.mark.parametrize(
					    """orient, container, keys, len_at"""							,					[
					        ("""records""", list, {"""tokens""", """labels""", """answers""", """id"""}, None),
					        ("""split""", dict, {"""columns""", """data"""}, """data"""),
					        ("""index""", dict, set("""0123456789"""       ), None),
					        ("""columns""", dict, {"""tokens""", """labels""", """answers""", """id"""}, """tokens"""),
					        ("""values""", list, None, None),
					        ("""table""", dict, {"""schema""", """data"""}, """data"""),
					    ]							,					)
					def   UpperCAmelCase__       (     self       :   Any							,					UpperCamelCase__       :   Union[str, Any]							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   Optional[int]							,					UpperCamelCase__       :   int       ):






										'''simple docstring'''






										with io.BytesIO() as buffer:
															JsonDatasetWriter(UpperCamelCase__							,					UpperCamelCase__							,					lines=UpperCamelCase__							,					orient=UpperCamelCase__       ).write()
															buffer.seek(0       )
															lowercase_  							=			load_json(UpperCamelCase__       )
										assert isinstance(UpperCamelCase__							,					UpperCamelCase__       )
										if keys:
															if container is dict:
																				assert exported_content.keys() == keys
															else:
																				assert exported_content[0].keys() == keys
										else:
															assert not hasattr(UpperCamelCase__							,					"""keys"""       ) and not hasattr(exported_content[0]							,					"""keys"""       )
										if len_at:
															assert len(exported_content[len_at]       ) == 10
										else:
															assert len(UpperCamelCase__       ) == 10

					@pytest.mark.parametrize("""lines, load_json_function"""							,					[(True, load_json_lines), (False, load_json)]       )
					def   UpperCAmelCase__       (     self       :   Optional[Any]							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   Dict       ):






										'''simple docstring'''






										with io.BytesIO() as buffer:
															JsonDatasetWriter(UpperCamelCase__							,					UpperCamelCase__							,					lines=UpperCamelCase__							,					num_proc=2       ).write()
															buffer.seek(0       )
															lowercase_  							=			load_json_function(UpperCamelCase__       )
										assert isinstance(UpperCamelCase__							,					UpperCamelCase__       )
										assert isinstance(exported_content[0]							,					UpperCamelCase__       )
										assert len(UpperCamelCase__       ) == 10

					@pytest.mark.parametrize(
					    """orient, container, keys, len_at"""							,					[
					        ("""records""", list, {"""tokens""", """labels""", """answers""", """id"""}, None),
					        ("""split""", dict, {"""columns""", """data"""}, """data"""),
					        ("""index""", dict, set("""0123456789"""       ), None),
					        ("""columns""", dict, {"""tokens""", """labels""", """answers""", """id"""}, """tokens"""),
					        ("""values""", list, None, None),
					        ("""table""", dict, {"""schema""", """data"""}, """data"""),
					    ]							,					)
					def   UpperCAmelCase__       (     self       :   Optional[Any]							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   Tuple							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   Any							,					UpperCamelCase__       :   Optional[int]       ):






										'''simple docstring'''






										with io.BytesIO() as buffer:
															JsonDatasetWriter(UpperCamelCase__							,					UpperCamelCase__							,					lines=UpperCamelCase__							,					orient=UpperCamelCase__							,					num_proc=2       ).write()
															buffer.seek(0       )
															lowercase_  							=			load_json(UpperCamelCase__       )
										assert isinstance(UpperCamelCase__							,					UpperCamelCase__       )
										if keys:
															if container is dict:
																				assert exported_content.keys() == keys
															else:
																				assert exported_content[0].keys() == keys
										else:
															assert not hasattr(UpperCamelCase__							,					"""keys"""       ) and not hasattr(exported_content[0]							,					"""keys"""       )
										if len_at:
															assert len(exported_content[len_at]       ) == 10
										else:
															assert len(UpperCamelCase__       ) == 10

					def   UpperCAmelCase__       (     self       :   Dict							,					UpperCamelCase__       :   Tuple       ):






										'''simple docstring'''






										with pytest.raises(UpperCamelCase__       ):
															with io.BytesIO() as buffer:
																				JsonDatasetWriter(UpperCamelCase__							,					UpperCamelCase__							,					num_proc=0       )







					@pytest.mark.parametrize("""compression, extension"""							,					[("""gzip""", """gz"""), ("""bz2""", """bz2"""), ("""xz""", """xz""")]       )
					def   UpperCAmelCase__       (     self       :   Any							,					UpperCamelCase__       :   Optional[Any]							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   List[Any]							,					UpperCamelCase__       :   Tuple       ):






										'''simple docstring'''






										lowercase_  							=			tmp_path_factory.mktemp("""data"""       ) / F'''test.json.{extension}'''
										lowercase_  							=			str(shared_datadir / F'''test_file.json.{extension}'''       )
										JsonDatasetWriter(UpperCamelCase__							,					UpperCamelCase__							,					compression=UpperCamelCase__       ).write()

										with fsspec.open(UpperCamelCase__							,					"""rb"""							,					compression="""infer"""       ) as f:
															lowercase_  							=			f.read()
										with fsspec.open(UpperCamelCase__							,					"""rb"""							,					compression="""infer"""       ) as f:
															lowercase_  							=			f.read()
										assert exported_content == original_content






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__   (	__magic_name__							,				__magic_name__							,				unittest.TestCase   ):
					__SCREAMING_SNAKE_CASE       :  Optional[int]						=   StableDiffusionDiffEditPipeline
					__SCREAMING_SNAKE_CASE       :  Optional[Any]						=   TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {'height', 'width', 'image'} | {'image_latents'}
					__SCREAMING_SNAKE_CASE       :  List[str]						=   TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {'image'} | {'image_latents'}
					__SCREAMING_SNAKE_CASE       :  int						=   frozenset(
					    []   )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
					__SCREAMING_SNAKE_CASE       :  Any						=   frozenset([]   )

					def   UpperCAmelCase__       (     self       :   Any       ):






										'''simple docstring'''






										torch.manual_seed(0       )
										lowercase_  							=			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=UpperCamelCase__							,					)
										lowercase_  							=			DDIMScheduler(
										    beta_start=0.00_085							,					beta_end=0.012							,					beta_schedule="""scaled_linear"""							,					clip_sample=UpperCamelCase__							,					set_alpha_to_one=UpperCamelCase__							,					)
										lowercase_  							=			DDIMInverseScheduler(
										    beta_start=0.00_085							,					beta_end=0.012							,					beta_schedule="""scaled_linear"""							,					clip_sample=UpperCamelCase__							,					set_alpha_to_zero=UpperCamelCase__							,					)
										torch.manual_seed(0       )
										lowercase_  							=			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       )
										lowercase_  							=			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							,					)
										lowercase_  							=			CLIPTextModel(UpperCamelCase__       )
										lowercase_  							=			CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       )

										lowercase_  							=			{
										    """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       :   int							,					UpperCamelCase__       :   Any							,					UpperCamelCase__       :   Any=0       ):






										'''simple docstring'''






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






										'''simple docstring'''






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






										'''simple docstring'''






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






										'''simple docstring'''






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

										lowercase_  							=			self.get_dummy_components()
										lowercase_  							=			self.pipeline_class(**UpperCamelCase__       )
										pipe.to(UpperCamelCase__       )
										pipe.set_progress_bar_config(disable=UpperCamelCase__       )

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

										lowercase_  							=			self.get_dummy_inputs(UpperCamelCase__       )
										lowercase_  							=			pipe(**UpperCamelCase__       )[0]

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

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

										lowercase_  							=			self.get_dummy_inputs(UpperCamelCase__       )
										lowercase_  							=			pipe_loaded(**UpperCamelCase__       )[0]

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

					def   UpperCAmelCase__       (     self       :   Optional[int]       ):






										'''simple docstring'''






										lowercase_  							=			"""cpu"""

										lowercase_  							=			self.get_dummy_components()
										lowercase_  							=			self.pipeline_class(**UpperCamelCase__       )
										pipe.to(UpperCamelCase__       )
										pipe.set_progress_bar_config(disable=UpperCamelCase__       )

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

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

					def   UpperCAmelCase__       (     self       :   Any       ):






										'''simple docstring'''






										lowercase_  							=			"""cpu"""

										lowercase_  							=			self.get_dummy_components()
										lowercase_  							=			self.pipeline_class(**UpperCamelCase__       )
										pipe.to(UpperCamelCase__       )
										pipe.set_progress_bar_config(disable=UpperCamelCase__       )

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

										self.assertEqual(image.shape							,					(2, 32, 32, 3)       )
										lowercase_  							=			np.array(
										    [0.5_150, 0.5_134, 0.5_043, 0.5_376, 0.4_694, 0.51_050, 0.5_015, 0.4_407, 0.4_799]							,					)
										lowercase_  							=			np.abs(image_slice.flatten() - expected_slice       ).max()
										self.assertLessEqual(UpperCamelCase__							,					1e-3       )

					def   UpperCAmelCase__       (     self       :   Optional[Any]       ):






										'''simple docstring'''






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







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






										'''simple docstring'''






										lowercase_  							=			"""cpu"""

										lowercase_  							=			self.get_dummy_components()

										lowercase_  							=			{"""beta_start""": 0.00_085, """beta_end""": 0.012, """beta_schedule""": """scaled_linear"""}
										lowercase_  							=			DPMSolverMultistepScheduler(**UpperCamelCase__       )
										lowercase_  							=			DPMSolverMultistepInverseScheduler(**UpperCamelCase__       )

										lowercase_  							=			self.pipeline_class(**UpperCamelCase__       )
										pipe.to(UpperCamelCase__       )
										pipe.set_progress_bar_config(disable=UpperCamelCase__       )

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

										self.assertEqual(image.shape							,					(2, 32, 32, 3)       )
										lowercase_  							=			np.array(
										    [0.5_150, 0.5_134, 0.5_043, 0.5_376, 0.4_694, 0.51_050, 0.5_015, 0.4_407, 0.4_799]							,					)
										lowercase_  							=			np.abs(image_slice.flatten() - expected_slice       ).max()
										self.assertLessEqual(UpperCamelCase__							,					1e-3       )







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

					def   UpperCAmelCase__       (     self       :   Tuple       ):






										'''simple docstring'''






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

					@classmethod
					def   UpperCAmelCase__       (     cls       :   Dict       ):






										'''simple docstring'''






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

										lowercase_  							=			raw_image.convert("""RGB"""       ).resize((768, 768)       )

										lowercase_  							=			raw_image

					def   UpperCAmelCase__       (     self       :   Dict       ):






										'''simple docstring'''






										lowercase_  							=			torch.manual_seed(0       )

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

										lowercase_  							=			"""a bowl of fruit"""
										lowercase_  							=			"""a bowl of pears"""

										lowercase_  							=			pipe.generate_mask(
										    image=self.raw_image							,					source_prompt=UpperCamelCase__							,					target_prompt=UpperCamelCase__							,					generator=UpperCamelCase__							,					)

										lowercase_  							=			pipe.invert(
										    prompt=UpperCamelCase__							,					image=self.raw_image							,					inpaint_strength=0.7							,					generator=UpperCamelCase__       ).latents

										lowercase_  							=			pipe(
										    prompt=UpperCamelCase__							,					mask_image=UpperCamelCase__							,					image_latents=UpperCamelCase__							,					generator=UpperCamelCase__							,					negative_prompt=UpperCamelCase__							,					inpaint_strength=0.7							,					output_type="""numpy"""							,					).images[0]

										lowercase_  							=			(
										    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       :   Any       ):






										'''simple docstring'''






										lowercase_  							=			torch.manual_seed(0       )

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

										lowercase_  							=			"""a bowl of fruit"""
										lowercase_  							=			"""a bowl of pears"""

										lowercase_  							=			pipe.generate_mask(
										    image=self.raw_image							,					source_prompt=UpperCamelCase__							,					target_prompt=UpperCamelCase__							,					generator=UpperCamelCase__							,					)

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

										lowercase_  							=			pipe(
										    prompt=UpperCamelCase__							,					mask_image=UpperCamelCase__							,					image_latents=UpperCamelCase__							,					generator=UpperCamelCase__							,					negative_prompt=UpperCamelCase__							,					inpaint_strength=0.7							,					num_inference_steps=25							,					output_type="""numpy"""							,					).images[0]

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


