



'''simple docstring'''





import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, PNDMScheduler, StableDiffusionInpaintPipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class 	lowerCamelCase      (						_UpperCAmelCase,     _UpperCAmelCase,     _UpperCAmelCase,     unittest.TestCase  ):







					'''simple docstring'''







					lowerCAmelCase_				:    Dict									=  StableDiffusionInpaintPipeline
					lowerCAmelCase_				:    Any									=  TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
					lowerCAmelCase_				:    List[Any]									=  TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
					lowerCAmelCase_				:    Optional[Any]									=  frozenset(
					    []  )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
					lowerCAmelCase_				:    List[str]									=  frozenset([]  )



					def    A__		(			self     ):
							torch.manual_seed(0     )
							UpperCAmelCase_						=				UNetaDConditionModel(
							    block_out_channels=(32, 64)    ,	layers_per_block=2    ,	sample_size=32    ,	in_channels=9    ,	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    ,	)
							UpperCAmelCase_						=				PNDMScheduler(skip_prk_steps=__UpperCamelCase     )
							torch.manual_seed(0     )
							UpperCAmelCase_						=				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     )
							UpperCAmelCase_						=				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    ,	)
							UpperCAmelCase_						=				CLIPTextModel(__UpperCamelCase     )
							UpperCAmelCase_						=				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )

							UpperCAmelCase_						=				{
							    "unet": unet,
							    "scheduler": scheduler,
							    "vae": vae,
							    "text_encoder": text_encoder,
							    "tokenizer": tokenizer,
							    "safety_checker": None,
							    "feature_extractor": None,
							}
							return components



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase=0     ):
							UpperCAmelCase_						=				floats_tensor((1, 3, 32, 32)    ,	rng=random.Random(__UpperCamelCase     )     ).to(__UpperCamelCase     )
							UpperCAmelCase_						=				image.cpu().permute(0    ,	2    ,	3    ,	1     )[0]
							UpperCAmelCase_						=				Image.fromarray(np.uinta(__UpperCamelCase     )     ).convert("RGB"     ).resize((64, 64)     )
							UpperCAmelCase_						=				Image.fromarray(np.uinta(image + 4     )     ).convert("RGB"     ).resize((64, 64)     )
							if str(__UpperCamelCase     ).startswith("mps"     ):
									UpperCAmelCase_						=				torch.manual_seed(__UpperCamelCase     )
							else:
									UpperCAmelCase_						=				torch.Generator(device=__UpperCamelCase     ).manual_seed(__UpperCamelCase     )
							UpperCAmelCase_						=				{
							    "prompt": "A painting of a squirrel eating a burger",
							    "image": init_image,
							    "mask_image": mask_image,
							    "generator": generator,
							    "num_inference_steps": 2,
							    "guidance_scale": 6.0,
							    "output_type": "numpy",
							}
							return inputs



					def    A__		(			self     ):
							UpperCAmelCase_						=				"cpu"  # ensure determinism for the device-dependent torch.Generator
							UpperCAmelCase_						=				self.get_dummy_components()
							UpperCAmelCase_						=				StableDiffusionInpaintPipeline(**__UpperCamelCase     )
							UpperCAmelCase_						=				sd_pipe.to(__UpperCamelCase     )
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase     )

							UpperCAmelCase_						=				self.get_dummy_inputs(__UpperCamelCase     )
							UpperCAmelCase_						=				sd_pipe(**__UpperCamelCase     ).images
							UpperCAmelCase_						=				image[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)
							UpperCAmelCase_						=				np.array([0.4727, 0.5735, 0.3941, 0.5446, 0.5926, 0.4394, 0.5062, 0.4654, 0.4476]     )

							assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1e-2



					def    A__		(			self     ):
							super().test_inference_batch_single_identical(expected_max_diff=3e-3     )





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







					'''simple docstring'''







					def    A__		(			self     ):
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()



					def    A__		(			self     ):
							UpperCAmelCase_						=				load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
							    "/sd2-inpaint/init_image.png"     )
							UpperCAmelCase_						=				load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"     )
							UpperCAmelCase_						=				load_numpy(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint"
							    "/yellow_cat_sitting_on_a_park_bench.npy"     )

							UpperCAmelCase_						=				"stabilityai/stable-diffusion-2-inpainting"
							UpperCAmelCase_						=				StableDiffusionInpaintPipeline.from_pretrained(__UpperCamelCase    ,	safety_checker=__UpperCamelCase     )
							pipe.to(__UpperCamelCase     )
							pipe.set_progress_bar_config(disable=__UpperCamelCase     )
							pipe.enable_attention_slicing()

							UpperCAmelCase_						=				"Face of a yellow cat, high resolution, sitting on a park bench"

							UpperCAmelCase_						=				torch.manual_seed(0     )
							UpperCAmelCase_						=				pipe(
							    prompt=__UpperCamelCase    ,	image=__UpperCamelCase    ,	mask_image=__UpperCamelCase    ,	generator=__UpperCamelCase    ,	output_type="np"    ,	)
							UpperCAmelCase_						=				output.images[0]

							assert image.shape == (512, 512, 3)
							assert np.abs(expected_image - image     ).max() < 9e-3



					def    A__		(			self     ):
							UpperCAmelCase_						=				load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
							    "/sd2-inpaint/init_image.png"     )
							UpperCAmelCase_						=				load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"     )
							UpperCAmelCase_						=				load_numpy(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint"
							    "/yellow_cat_sitting_on_a_park_bench_fp16.npy"     )

							UpperCAmelCase_						=				"stabilityai/stable-diffusion-2-inpainting"
							UpperCAmelCase_						=				StableDiffusionInpaintPipeline.from_pretrained(
							    __UpperCamelCase    ,	torch_dtype=torch.floataa    ,	safety_checker=__UpperCamelCase    ,	)
							pipe.to(__UpperCamelCase     )
							pipe.set_progress_bar_config(disable=__UpperCamelCase     )
							pipe.enable_attention_slicing()

							UpperCAmelCase_						=				"Face of a yellow cat, high resolution, sitting on a park bench"

							UpperCAmelCase_						=				torch.manual_seed(0     )
							UpperCAmelCase_						=				pipe(
							    prompt=__UpperCamelCase    ,	image=__UpperCamelCase    ,	mask_image=__UpperCamelCase    ,	generator=__UpperCamelCase    ,	output_type="np"    ,	)
							UpperCAmelCase_						=				output.images[0]

							assert image.shape == (512, 512, 3)
							assert np.abs(expected_image - image     ).max() < 5e-1



					def    A__		(			self     ):
							torch.cuda.empty_cache()
							torch.cuda.reset_max_memory_allocated()
							torch.cuda.reset_peak_memory_stats()

							UpperCAmelCase_						=				load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
							    "/sd2-inpaint/init_image.png"     )
							UpperCAmelCase_						=				load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"     )

							UpperCAmelCase_						=				"stabilityai/stable-diffusion-2-inpainting"
							UpperCAmelCase_						=				PNDMScheduler.from_pretrained(__UpperCamelCase    ,	subfolder="scheduler"     )
							UpperCAmelCase_						=				StableDiffusionInpaintPipeline.from_pretrained(
							    __UpperCamelCase    ,	safety_checker=__UpperCamelCase    ,	scheduler=__UpperCamelCase    ,	torch_dtype=torch.floataa    ,	)
							pipe.to(__UpperCamelCase     )
							pipe.set_progress_bar_config(disable=__UpperCamelCase     )
							pipe.enable_attention_slicing(1     )
							pipe.enable_sequential_cpu_offload()

							UpperCAmelCase_						=				"Face of a yellow cat, high resolution, sitting on a park bench"

							UpperCAmelCase_						=				torch.manual_seed(0     )
							UpperCAmelCase_						=				pipe(
							    prompt=__UpperCamelCase    ,	image=__UpperCamelCase    ,	mask_image=__UpperCamelCase    ,	generator=__UpperCamelCase    ,	num_inference_steps=2    ,	output_type="np"    ,	)

							UpperCAmelCase_						=				torch.cuda.max_memory_allocated()
							# make sure that less than 2.65 GB is allocated
							assert mem_bytes < 2.65 * 10**9






from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    is_valid_image,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
			import PIL


SCREAMING_SNAKE_CASE     				=						logging.get_logger(__name__)







def snake_case__  (			__SCREAMING_SNAKE_CASE )		->							List[List[ImageInput]]:
		if isinstance(__SCREAMING_SNAKE_CASE			,     (list, tuple) ) and isinstance(videos[0]			,     (list, tuple) ) and is_valid_image(videos[0][0] ):
				return videos

		elif isinstance(__SCREAMING_SNAKE_CASE			,     (list, tuple) ) and is_valid_image(videos[0] ):
				return [videos]

		elif is_valid_image(__SCREAMING_SNAKE_CASE ):
				return [[videos]]

		raise ValueError(f'''Could not make batched video from {videos}''' )



class 	lowerCamelCase      (						lowercase__  ):







					'''simple docstring'''







					lowerCAmelCase_				:    List[Any]									=  ['pixel_values']



					def __init__(			self    ,	lowerCAmelCase = True    ,	lowerCAmelCase = None    ,	lowerCAmelCase = PILImageResampling.BILINEAR    ,	lowerCAmelCase = True    ,	lowerCAmelCase = None    ,	lowerCAmelCase = True    ,	lowerCAmelCase = 1 / 255    ,	lowerCAmelCase = True    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	**lowerCAmelCase    ,	):
							super().__init__(**lowerCAmelCase     )
							UpperCAmelCase_						=				size if size is not None else {"shortest_edge": 224}
							UpperCAmelCase_						=				get_size_dict(lowerCAmelCase    ,	default_to_square=lowerCAmelCase     )
							UpperCAmelCase_						=				crop_size if crop_size is not None else {"height": 224, "width": 224}
							UpperCAmelCase_						=				get_size_dict(lowerCAmelCase    ,	param_name="crop_size"     )

							UpperCAmelCase_						=				do_resize
							UpperCAmelCase_						=				size
							UpperCAmelCase_						=				do_center_crop
							UpperCAmelCase_						=				crop_size
							UpperCAmelCase_						=				resample
							UpperCAmelCase_						=				do_rescale
							UpperCAmelCase_						=				rescale_factor
							UpperCAmelCase_						=				do_normalize
							UpperCAmelCase_						=				image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
							UpperCAmelCase_						=				image_std if image_std is not None else IMAGENET_STANDARD_STD



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase = PILImageResampling.BILINEAR    ,	lowerCAmelCase = None    ,	**lowerCAmelCase    ,	):
							UpperCAmelCase_						=				get_size_dict(lowerCAmelCase    ,	default_to_square=lowerCAmelCase     )
							if "shortest_edge" in size:
									UpperCAmelCase_						=				get_resize_output_image_size(lowerCAmelCase    ,	size["shortest_edge"]    ,	default_to_square=lowerCAmelCase     )
							elif "height" in size and "width" in size:
									UpperCAmelCase_						=				(size["height"], size["width"])
							else:
									raise ValueError(f'''Size must have \'height\' and \'width\' or \'shortest_edge\' as keys. Got {size.keys()}'''     )
							return resize(lowerCAmelCase    ,	size=lowerCAmelCase    ,	resample=lowerCAmelCase    ,	data_format=lowerCAmelCase    ,	**lowerCAmelCase     )



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase = None    ,	**lowerCAmelCase    ,	):
							UpperCAmelCase_						=				get_size_dict(lowerCAmelCase     )
							if "height" not in size or "width" not in size:
									raise ValueError(f'''Size must have \'height\' and \'width\' as keys. Got {size.keys()}'''     )
							return center_crop(lowerCAmelCase    ,	size=(size["height"], size["width"])    ,	data_format=lowerCAmelCase    ,	**lowerCAmelCase     )



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase = None    ,	**lowerCAmelCase    ,	):
							return rescale(lowerCAmelCase    ,	scale=lowerCAmelCase    ,	data_format=lowerCAmelCase    ,	**lowerCAmelCase     )



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase = None    ,	**lowerCAmelCase    ,	):
							return normalize(lowerCAmelCase    ,	mean=lowerCAmelCase    ,	std=lowerCAmelCase    ,	data_format=lowerCAmelCase    ,	**lowerCAmelCase     )



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = ChannelDimension.FIRST    ,	):
							if do_resize and size is None or resample is None:
									raise ValueError("Size and resample must be specified if do_resize is True."     )

							if do_center_crop and crop_size is None:
									raise ValueError("Crop size must be specified if do_center_crop is True."     )

							if do_rescale and rescale_factor is None:
									raise ValueError("Rescale factor must be specified if do_rescale is True."     )

							if do_normalize and (image_mean is None or image_std is None):
									raise ValueError("Image mean and std must be specified if do_normalize is True."     )

							# All transformations expect numpy arrays.
							UpperCAmelCase_						=				to_numpy_array(lowerCAmelCase     )

							if do_resize:
									UpperCAmelCase_						=				self.resize(image=lowerCAmelCase    ,	size=lowerCAmelCase    ,	resample=lowerCAmelCase     )

							if do_center_crop:
									UpperCAmelCase_						=				self.center_crop(lowerCAmelCase    ,	size=lowerCAmelCase     )

							if do_rescale:
									UpperCAmelCase_						=				self.rescale(image=lowerCAmelCase    ,	scale=lowerCAmelCase     )

							if do_normalize:
									UpperCAmelCase_						=				self.normalize(image=lowerCAmelCase    ,	mean=lowerCAmelCase    ,	std=lowerCAmelCase     )

							UpperCAmelCase_						=				to_channel_dimension_format(lowerCAmelCase    ,	lowerCAmelCase     )
							return image



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = None    ,	lowerCAmelCase = ChannelDimension.FIRST    ,	**lowerCAmelCase    ,	):
							UpperCAmelCase_						=				do_resize if do_resize is not None else self.do_resize
							UpperCAmelCase_						=				resample if resample is not None else self.resample
							UpperCAmelCase_						=				do_center_crop if do_center_crop is not None else self.do_center_crop
							UpperCAmelCase_						=				do_rescale if do_rescale is not None else self.do_rescale
							UpperCAmelCase_						=				rescale_factor if rescale_factor is not None else self.rescale_factor
							UpperCAmelCase_						=				do_normalize if do_normalize is not None else self.do_normalize
							UpperCAmelCase_						=				image_mean if image_mean is not None else self.image_mean
							UpperCAmelCase_						=				image_std if image_std is not None else self.image_std

							UpperCAmelCase_						=				size if size is not None else self.size
							UpperCAmelCase_						=				get_size_dict(lowerCAmelCase    ,	default_to_square=lowerCAmelCase     )
							UpperCAmelCase_						=				crop_size if crop_size is not None else self.crop_size
							UpperCAmelCase_						=				get_size_dict(lowerCAmelCase    ,	param_name="crop_size"     )

							if not valid_images(lowerCAmelCase     ):
									raise ValueError(
									    "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
									    "torch.Tensor, tf.Tensor or jax.ndarray."     )

							UpperCAmelCase_						=				make_batched(lowerCAmelCase     )

							UpperCAmelCase_						=				[
							    [
							        self._preprocess_image(
							            image=lowerCAmelCase    ,	do_resize=lowerCAmelCase    ,	size=lowerCAmelCase    ,	resample=lowerCAmelCase    ,	do_center_crop=lowerCAmelCase    ,	crop_size=lowerCAmelCase    ,	do_rescale=lowerCAmelCase    ,	rescale_factor=lowerCAmelCase    ,	do_normalize=lowerCAmelCase    ,	image_mean=lowerCAmelCase    ,	image_std=lowerCAmelCase    ,	data_format=lowerCAmelCase    ,	)
							        for img in video
							    ]
							    for video in videos
							]

							UpperCAmelCase_						=				{"pixel_values": videos}
							return BatchFeature(data=lowerCAmelCase    ,	tensor_type=lowerCAmelCase     )

