


from typing import Any, Dict, List, Union

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging, requires_backends
from .base import PIPELINE_INIT_ARGS, ChunkPipeline


if is_vision_available():
							from PIL import Image

							from ..image_utils import load_image

if is_torch_available():
							import torch

							from transformers.modeling_outputs import BaseModelOutput

							from ..models.auto.modeling_auto import MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING

a					  =     logging.get_logger(__name__)







@add_end_docstrings(__magic_name__   )
class        UpperCamelCase__   (	__magic_name__   ):

					def __init__(     self       :   Union[str, Any]							,					**UpperCamelCase__       :   Union[str, Any]       ):






										'''simple docstring'''






										super().__init__(**UpperCamelCase__       )

										if self.framework == "tf":
															raise ValueError(F'''The {self.__class__} is only available in PyTorch.'''       )

										requires_backends(self							,					"""vision"""       )
										self.check_model_type(UpperCamelCase__       )

					def __call__(     self       :   Dict							,					UpperCamelCase__       :   Union[str, "Image.Image", List[Dict[str, Any]]]							,					UpperCamelCase__       :   Union[str, List[str]] = None							,					**UpperCamelCase__       :   Any							,					):






										'''simple docstring'''






										if "text_queries" in kwargs:
															lowercase_  							=			kwargs.pop("""text_queries"""       )

										if isinstance(UpperCamelCase__							,					(str, Image.Image)       ):
															lowercase_  							=			{"""image""": image, """candidate_labels""": candidate_labels}
										else:
															lowercase_  							=			image
										lowercase_  							=			super().__call__(UpperCamelCase__							,					**UpperCamelCase__       )
										return results

					def   UpperCAmelCase__       (     self       :   int							,					**UpperCamelCase__       :   Tuple       ):






										'''simple docstring'''






										lowercase_  							=			{}
										if "threshold" in kwargs:
															lowercase_  							=			kwargs["""threshold"""]
										if "top_k" in kwargs:
															lowercase_  							=			kwargs["""top_k"""]
										return {}, {}, postprocess_params

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






										'''simple docstring'''






										lowercase_  							=			load_image(inputs["""image"""]       )
										lowercase_  							=			inputs["""candidate_labels"""]
										if isinstance(UpperCamelCase__							,					UpperCamelCase__       ):
															lowercase_  							=			candidate_labels.split(""","""       )

										lowercase_  							=			torch.tensor([[image.height, image.width]]							,					dtype=torch.intaa       )
										for i, candidate_label in enumerate(UpperCamelCase__       ):
															lowercase_  							=			self.tokenizer(UpperCamelCase__							,					return_tensors=self.framework       )
															lowercase_  							=			self.image_processor(UpperCamelCase__							,					return_tensors=self.framework       )
															yield {
															    "is_last": i == len(UpperCamelCase__       ) - 1,
															    "target_size": target_size,
															    "candidate_label": candidate_label,
															    **text_inputs,
															    **image_features,
															}

					def   UpperCAmelCase__       (     self       :   Optional[Any]							,					UpperCamelCase__       :   Any       ):






										'''simple docstring'''






										lowercase_  							=			model_inputs.pop("""target_size"""       )
										lowercase_  							=			model_inputs.pop("""candidate_label"""       )
										lowercase_  							=			model_inputs.pop("""is_last"""       )

										lowercase_  							=			self.model(**UpperCamelCase__       )

										lowercase_  							=			{"""target_size""": target_size, """candidate_label""": candidate_label, """is_last""": is_last, **outputs}
										return model_outputs

					def   UpperCAmelCase__       (     self       :   Union[str, Any]							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   Dict=0.1							,					UpperCamelCase__       :   List[str]=None       ):






										'''simple docstring'''






										lowercase_  							=			[]
										for model_output in model_outputs:
															lowercase_  							=			model_output["""candidate_label"""]
															lowercase_  							=			BaseModelOutput(UpperCamelCase__       )
															lowercase_  							=			self.image_processor.post_process_object_detection(
															    outputs=UpperCamelCase__							,					threshold=UpperCamelCase__							,					target_sizes=model_output["""target_size"""]       )[0]

															for index in outputs["scores"].nonzero():
																				lowercase_  							=			outputs["""scores"""][index].item()
																				lowercase_  							=			self._get_bounding_box(outputs["""boxes"""][index][0]       )

																				lowercase_  							=			{"""score""": score, """label""": label, """box""": box}
																				results.append(UpperCamelCase__       )

										lowercase_  							=			sorted(UpperCamelCase__							,					key=lambda UpperCamelCase__       : x["score"]							,					reverse=UpperCamelCase__       )
										if top_k:
															lowercase_  							=			results[:top_k]

										return results







					def   UpperCAmelCase__       (     self       :   List[Any]							,					UpperCamelCase__       :   "torch.Tensor"       ):






										'''simple docstring'''






										if self.framework != "pt":
															raise ValueError("""The ZeroShotObjectDetectionPipeline is only available in PyTorch."""       )
										lowercase_							,					lowercase_							,					lowercase_							,					lowercase_  							=			box.int().tolist()
										lowercase_  							=			{
										    """xmin""": xmin,
										    """ymin""": ymin,
										    """xmax""": xmax,
										    """ymax""": ymax,
										}
										return bbox






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


