

import numpy as np
import torch
from torch.utils.data import Dataset, IterableDataset

from ..utils.generic import ModelOutput



class   __lowercase    (			UpperCAmelCase__     ):




			"""simple docstring"""







			def __init__(			self  :					Union[str, Any]       ,						lowerCAmelCase__  :					List[Any]       ,						lowerCAmelCase__  :					List[str]       ,						lowerCAmelCase__  :					int):
									SCREAMING_SNAKE_CASE_:			Any  		=	dataset
									SCREAMING_SNAKE_CASE_:			List[str]  		=	process
									SCREAMING_SNAKE_CASE_:			int  		=	params



			def __len__(			self  :					List[str]):
									return len(self.dataset)




			def __getitem__(			self  :					Tuple       ,						lowerCAmelCase__  :					str):
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	self.dataset[i]
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	self.process(_SCREAMING_SNAKE_CASE       ,						**self.params)
									return processed







class   __lowercase    (			UpperCAmelCase__     ):




			"""simple docstring"""







			def __init__(			self  :					List[Any]       ,						lowerCAmelCase__  :					str       ,						lowerCAmelCase__  :					int       ,						lowerCAmelCase__  :					str       ,						lowerCAmelCase__  :					Optional[int]=None):
									SCREAMING_SNAKE_CASE_:			Any  		=	loader
									SCREAMING_SNAKE_CASE_:			Tuple  		=	infer
									SCREAMING_SNAKE_CASE_:			List[Any]  		=	params
									if loader_batch_size == 1:
															# Let's spare some time by deactivating altogether
															SCREAMING_SNAKE_CASE_:			str  		=	None
									SCREAMING_SNAKE_CASE_:			List[str]  		=	loader_batch_size

									# Internal bookkeeping
									SCREAMING_SNAKE_CASE_:			Dict  		=	None
									SCREAMING_SNAKE_CASE_:			Any  		=	None



			def __len__(			self  :					List[Any]):
									return len(self.loader)



			def __iter__(			self  :					Any):
									SCREAMING_SNAKE_CASE_:			Tuple  		=	iter(self.loader)
									return self



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[Any]):




									if isinstance(self._loader_batch_data       ,						torch.Tensor):
															# Batch data is simple tensor, just fetch the slice
															SCREAMING_SNAKE_CASE_:			Any  		=	self._loader_batch_data[self._loader_batch_index]
									else:
															# Batch data is assumed to be BaseModelOutput (or dict)
															SCREAMING_SNAKE_CASE_:			List[str]  		=	{}
															for k, element in self._loader_batch_data.items():
																					if isinstance(_SCREAMING_SNAKE_CASE       ,						_SCREAMING_SNAKE_CASE):
																											# Convert ModelOutput to tuple first
																											SCREAMING_SNAKE_CASE_:			List[str]  		=	element.to_tuple()
																											if isinstance(element[0]       ,						torch.Tensor):
																																	SCREAMING_SNAKE_CASE_:			Dict  		=	tuple(el[self._loader_batch_index].unsqueeze(0) for el in element)
																											elif isinstance(element[0]       ,						np.ndarray):
																																	SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	tuple(np.expand_dims(el[self._loader_batch_index]       ,						0) for el in element)
																											continue
																					if k in {"hidden_states", "past_key_values", "attentions"} and isinstance(_SCREAMING_SNAKE_CASE       ,						_SCREAMING_SNAKE_CASE):
																											# Those are stored as lists of tensors so need specific unbatching.
																											if isinstance(element[0]       ,						torch.Tensor):
																																	SCREAMING_SNAKE_CASE_:			Dict  		=	tuple(el[self._loader_batch_index].unsqueeze(0) for el in element)
																											elif isinstance(element[0]       ,						np.ndarray):
																																	SCREAMING_SNAKE_CASE_:			List[str]  		=	tuple(np.expand_dims(el[self._loader_batch_index]       ,						0) for el in element)
																											continue
																					if element is None:
																											# This can happen for optional data that get passed around
																											SCREAMING_SNAKE_CASE_:			List[Any]  		=	None
																					elif isinstance(element[self._loader_batch_index]       ,						torch.Tensor):
																											# Take correct batch data, but make it looked like batch_size=1
																											# For compatibility with other methods within transformers

																											SCREAMING_SNAKE_CASE_:			Tuple  		=	element[self._loader_batch_index].unsqueeze(0)
																					elif isinstance(element[self._loader_batch_index]       ,						np.ndarray):
																											# Take correct batch data, but make it looked like batch_size=1
																											# For compatibility with other methods within transformers
																											SCREAMING_SNAKE_CASE_:			Any  		=	np.expand_dims(element[self._loader_batch_index]       ,						0)
																					else:
																											# This is typically a list, so no need to `unsqueeze`.
																											SCREAMING_SNAKE_CASE_:			Any  		=	element[self._loader_batch_index]
																											     # Recreate the element by reusing the original class to make it look
																											     # batch_size=1
															SCREAMING_SNAKE_CASE_:			List[str]  		=	self._loader_batch_data.__class__(_SCREAMING_SNAKE_CASE)
									self._loader_batch_index += 1
									return result




			def 	_SCREAMING_SNAKE_CASE      (			self  :					Dict):




									if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
															# We are currently unrolling a batch so we just need to return
															# the current item within a batch
															return self.loader_batch_item()

									# We're out of items within a batch
									SCREAMING_SNAKE_CASE_:			Dict  		=	next(self.iterator)
									SCREAMING_SNAKE_CASE_:			List[str]  		=	self.infer(_SCREAMING_SNAKE_CASE       ,						**self.params)
									# We now have a batch of "inferred things".
									if self.loader_batch_size is not None:
															# Try to infer the size of the batch
															if isinstance(_SCREAMING_SNAKE_CASE       ,						torch.Tensor):
																					SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	processed
															else:
																					SCREAMING_SNAKE_CASE_:			Dict  		=	list(processed.keys())[0]
																					SCREAMING_SNAKE_CASE_:			int  		=	processed[key]
															if isinstance(_SCREAMING_SNAKE_CASE       ,						_SCREAMING_SNAKE_CASE):
																					SCREAMING_SNAKE_CASE_:			Tuple  		=	len(_SCREAMING_SNAKE_CASE)
															else:
																					SCREAMING_SNAKE_CASE_:			int  		=	first_tensor.shape[0]
															if 0 < observed_batch_size < self.loader_batch_size:
																					# could be last batch so we can't unroll as many
																					# elements.
																					SCREAMING_SNAKE_CASE_:			Any  		=	observed_batch_size
															# Setting internal index to unwrap the batch
															SCREAMING_SNAKE_CASE_:			Tuple  		=	processed
															SCREAMING_SNAKE_CASE_:			List[Any]  		=	0
															return self.loader_batch_item()
									else:
															# We're not unrolling batches
															return processed







class   __lowercase    (			UpperCAmelCase__     ):




			"""simple docstring"""







			def __init__(			self  :					Union[str, Any]       ,						lowerCAmelCase__  :					int       ,						lowerCAmelCase__  :					List[str]       ,						lowerCAmelCase__  :					str       ,						lowerCAmelCase__  :					Tuple=None):
									super().__init__(_SCREAMING_SNAKE_CASE       ,						_SCREAMING_SNAKE_CASE       ,						_SCREAMING_SNAKE_CASE)



			def __iter__(			self  :					Optional[int]):
									SCREAMING_SNAKE_CASE_:			Any  		=	iter(self.loader)
									SCREAMING_SNAKE_CASE_:			str  		=	None
									return self




			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[int]):




									if self.subiterator is None:
															SCREAMING_SNAKE_CASE_:			List[Any]  		=	self.infer(next(self.iterator)       ,						**self.params)
									try:
															# Try to return next item
															SCREAMING_SNAKE_CASE_:			Any  		=	next(self.subiterator)
									except StopIteration:
															# When a preprocess iterator ends, we can start lookig at the next item
															# ChunkIterator will keep feeding until ALL elements of iterator
															# all have created their subiterator and have been iterating against.
															#
															# Another way to look at it, is we're basically flattening lists of lists
															# into a single list, but with generators
															SCREAMING_SNAKE_CASE_:			int  		=	self.infer(next(self.iterator)       ,						**self.params)
															SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	next(self.subiterator)
									return processed







class   __lowercase    (			UpperCAmelCase__     ):




			"""simple docstring"""







			def __iter__(			self  :					List[str]):
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	iter(self.loader)
									return self




			def 	_SCREAMING_SNAKE_CASE      (			self  :					int):
									SCREAMING_SNAKE_CASE_:			int  		=	False
									SCREAMING_SNAKE_CASE_:			List[str]  		=	[]
									if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
															while self._loader_batch_index < self.loader_batch_size:
																					SCREAMING_SNAKE_CASE_:			Tuple  		=	self.loader_batch_item()
																					SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	item.pop("is_last")
																					accumulator.append(_SCREAMING_SNAKE_CASE)
																					if is_last:
																											return accumulator

									while not is_last:
															SCREAMING_SNAKE_CASE_:			Optional[int]  		=	self.infer(next(self.iterator)       ,						**self.params)
															if self.loader_batch_size is not None:
																					if isinstance(_SCREAMING_SNAKE_CASE       ,						torch.Tensor):
																											SCREAMING_SNAKE_CASE_:			Any  		=	processed
																					else:
																											SCREAMING_SNAKE_CASE_:			List[Any]  		=	list(processed.keys())[0]
																											SCREAMING_SNAKE_CASE_:			int  		=	processed[key]
																					if isinstance(_SCREAMING_SNAKE_CASE       ,						_SCREAMING_SNAKE_CASE):
																											SCREAMING_SNAKE_CASE_:			Dict  		=	len(_SCREAMING_SNAKE_CASE)
																					else:
																											SCREAMING_SNAKE_CASE_:			List[str]  		=	first_tensor.shape[0]
																					if 0 < observed_batch_size < self.loader_batch_size:
																											# could be last batch so we can't unroll as many
																											# elements.
																											SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	observed_batch_size
																					SCREAMING_SNAKE_CASE_:			Any  		=	processed
																					SCREAMING_SNAKE_CASE_:			int  		=	0
																					while self._loader_batch_index < self.loader_batch_size:
																											SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	self.loader_batch_item()
																											SCREAMING_SNAKE_CASE_:			Any  		=	item.pop("is_last")
																											accumulator.append(_SCREAMING_SNAKE_CASE)
																											if is_last:
																																	return accumulator
															else:
																					SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	processed
																					SCREAMING_SNAKE_CASE_:			Any  		=	item.pop("is_last")
																					accumulator.append(_SCREAMING_SNAKE_CASE)
									return accumulator







class   __lowercase    (			UpperCAmelCase__     ):




			"""simple docstring"""







			def __init__(			self  :					Any       ,						lowerCAmelCase__  :					List[str]       ,						lowerCAmelCase__  :					Optional[Any]):
									SCREAMING_SNAKE_CASE_:			Tuple  		=	dataset
									SCREAMING_SNAKE_CASE_:			Dict  		=	key



			def __len__(			self  :					Optional[Any]):
									return len(self.dataset)




			def __getitem__(			self  :					Any       ,						lowerCAmelCase__  :					Any):
									return self.dataset[i][self.key]







class   __lowercase    (			UpperCAmelCase__     ):




			"""simple docstring"""







			def __init__(			self  :					List[Any]       ,						lowerCAmelCase__  :					Union[str, Any]       ,						lowerCAmelCase__  :					Union[str, Any]       ,						lowerCAmelCase__  :					Optional[Any]):
									SCREAMING_SNAKE_CASE_:			Dict  		=	dataset
									SCREAMING_SNAKE_CASE_:			Dict  		=	keya
									SCREAMING_SNAKE_CASE_:			List[Any]  		=	keya



			def __len__(			self  :					Optional[int]):
									return len(self.dataset)




			def __getitem__(			self  :					Tuple       ,						lowerCAmelCase__  :					Optional[int]):
									return {"text": self.dataset[i][self.keya], "text_pair": self.dataset[i][self.keya]}









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   __lowercase    (			UpperCAmelCase_		,	UpperCAmelCase_		,	UpperCAmelCase_		,	unittest.TestCase     ):




			"""simple docstring"""







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



			def 	_SCREAMING_SNAKE_CASE      (			self  :					int):
									torch.manual_seed(0)
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	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=lowerCAmelCase__       ,						)
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	PNDMScheduler(skip_prk_steps=lowerCAmelCase__)
									torch.manual_seed(0)
									SCREAMING_SNAKE_CASE_:			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)
									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	CLIPTextConfig(
									    bos_token_id=0       ,						eos_token_id=2       ,						hidden_size=32       ,						intermediate_size=37       ,						layer_norm_eps=1E-05       ,						num_attention_heads=4       ,						num_hidden_layers=5       ,						pad_token_id=1       ,						vocab_size=1000       ,						hidden_act="gelu"       ,						projection_dim=512       ,						)
									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	CLIPTextModel(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

									SCREAMING_SNAKE_CASE_:			List[str]  		=	{
									    "unet": unet,
									    "scheduler": scheduler,
									    "vae": vae,
									    "text_encoder": text_encoder,
									    "tokenizer": tokenizer,
									    "safety_checker": None,
									    "feature_extractor": None,
									}
									return components



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Union[str, Any]       ,						lowerCAmelCase__  :					Dict       ,						lowerCAmelCase__  :					Optional[int]=0):
									# TODO: use tensor inputs instead of PIL, this is here just to leave the old expected_slices untouched
									SCREAMING_SNAKE_CASE_:			Dict  		=	floats_tensor((1, 3, 32, 32)       ,						rng=random.Random(lowerCAmelCase__)).to(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			str  		=	image.cpu().permute(0       ,						2       ,						3       ,						1)[0]
									SCREAMING_SNAKE_CASE_:			Tuple  		=	Image.fromarray(np.uinta(lowerCAmelCase__)).convert("RGB").resize((64, 64))
									SCREAMING_SNAKE_CASE_:			List[str]  		=	Image.fromarray(np.uinta(image + 4)).convert("RGB").resize((64, 64))
									if str(lowerCAmelCase__).startswith("mps"):
															SCREAMING_SNAKE_CASE_:			Tuple  		=	torch.manual_seed(lowerCAmelCase__)
									else:
															SCREAMING_SNAKE_CASE_:			Any  		=	torch.Generator(device=lowerCAmelCase__).manual_seed(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			List[Any]  		=	{
									    "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 	_SCREAMING_SNAKE_CASE      (			self  :					Tuple):
									SCREAMING_SNAKE_CASE_:			Tuple  		=	"cpu"  # ensure determinism for the device-dependent torch.Generator
									SCREAMING_SNAKE_CASE_:			int  		=	self.get_dummy_components()
									SCREAMING_SNAKE_CASE_:			int  		=	StableDiffusionInpaintPipeline(**lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			int  		=	sd_pipe.to(lowerCAmelCase__)
									sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__)

									SCREAMING_SNAKE_CASE_:			Any  		=	self.get_dummy_inputs(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			List[Any]  		=	sd_pipe(**lowerCAmelCase__).images
									SCREAMING_SNAKE_CASE_:			Dict  		=	image[0, -3:, -3:, -1]

									assert image.shape == (1, 64, 64, 3)
									SCREAMING_SNAKE_CASE_:			Tuple  		=	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 	_SCREAMING_SNAKE_CASE      (			self  :					List[str]):
									super().test_inference_batch_single_identical(expected_max_diff=3E-3)







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




			"""simple docstring"""







			def 	_SCREAMING_SNAKE_CASE      (			self  :					str):
									# clean up the VRAM after each test
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()



			def 	_SCREAMING_SNAKE_CASE      (			self  :					List[Any]):
									SCREAMING_SNAKE_CASE_:			Any  		=	load_image(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
									    "/sd2-inpaint/init_image.png")
									SCREAMING_SNAKE_CASE_:			int  		=	load_image(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png")
									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	load_numpy(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint"
									    "/yellow_cat_sitting_on_a_park_bench.npy")

									SCREAMING_SNAKE_CASE_:			List[str]  		=	"stabilityai/stable-diffusion-2-inpainting"
									SCREAMING_SNAKE_CASE_:			Any  		=	StableDiffusionInpaintPipeline.from_pretrained(lowerCAmelCase__       ,						safety_checker=lowerCAmelCase__)
									pipe.to(lowerCAmelCase__)
									pipe.set_progress_bar_config(disable=lowerCAmelCase__)
									pipe.enable_attention_slicing()

									SCREAMING_SNAKE_CASE_:			str  		=	"Face of a yellow cat, high resolution, sitting on a park bench"

									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	torch.manual_seed(0)
									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	pipe(
									    prompt=lowerCAmelCase__       ,						image=lowerCAmelCase__       ,						mask_image=lowerCAmelCase__       ,						generator=lowerCAmelCase__       ,						output_type="np"       ,						)
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	output.images[0]

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



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Tuple):
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	load_image(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
									    "/sd2-inpaint/init_image.png")
									SCREAMING_SNAKE_CASE_:			List[Any]  		=	load_image(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png")
									SCREAMING_SNAKE_CASE_:			Any  		=	load_numpy(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint"
									    "/yellow_cat_sitting_on_a_park_bench_fp16.npy")

									SCREAMING_SNAKE_CASE_:			str  		=	"stabilityai/stable-diffusion-2-inpainting"
									SCREAMING_SNAKE_CASE_:			Dict  		=	StableDiffusionInpaintPipeline.from_pretrained(
									    lowerCAmelCase__       ,						torch_dtype=torch.floataa       ,						safety_checker=lowerCAmelCase__       ,						)
									pipe.to(lowerCAmelCase__)
									pipe.set_progress_bar_config(disable=lowerCAmelCase__)
									pipe.enable_attention_slicing()

									SCREAMING_SNAKE_CASE_:			List[str]  		=	"Face of a yellow cat, high resolution, sitting on a park bench"

									SCREAMING_SNAKE_CASE_:			Tuple  		=	torch.manual_seed(0)
									SCREAMING_SNAKE_CASE_:			Dict  		=	pipe(
									    prompt=lowerCAmelCase__       ,						image=lowerCAmelCase__       ,						mask_image=lowerCAmelCase__       ,						generator=lowerCAmelCase__       ,						output_type="np"       ,						)
									SCREAMING_SNAKE_CASE_:			Any  		=	output.images[0]

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




			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[Any]):
									torch.cuda.empty_cache()
									torch.cuda.reset_max_memory_allocated()
									torch.cuda.reset_peak_memory_stats()

									SCREAMING_SNAKE_CASE_:			Tuple  		=	load_image(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
									    "/sd2-inpaint/init_image.png")
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	load_image(
									    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png")

									SCREAMING_SNAKE_CASE_:			List[str]  		=	"stabilityai/stable-diffusion-2-inpainting"
									SCREAMING_SNAKE_CASE_:			Tuple  		=	PNDMScheduler.from_pretrained(lowerCAmelCase__       ,						subfolder="scheduler")
									SCREAMING_SNAKE_CASE_:			Any  		=	StableDiffusionInpaintPipeline.from_pretrained(
									    lowerCAmelCase__       ,						safety_checker=lowerCAmelCase__       ,						scheduler=lowerCAmelCase__       ,						torch_dtype=torch.floataa       ,						)
									pipe.to(lowerCAmelCase__)
									pipe.set_progress_bar_config(disable=lowerCAmelCase__)
									pipe.enable_attention_slicing(1)
									pipe.enable_sequential_cpu_offload()

									SCREAMING_SNAKE_CASE_:			Any  		=	"Face of a yellow cat, high resolution, sitting on a park bench"

									SCREAMING_SNAKE_CASE_:			Any  		=	torch.manual_seed(0)
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	pipe(
									    prompt=lowerCAmelCase__       ,						image=lowerCAmelCase__       ,						mask_image=lowerCAmelCase__       ,						generator=lowerCAmelCase__       ,						num_inference_steps=2       ,						output_type="np"       ,						)

									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	torch.cuda.max_memory_allocated()
									# make sure that less than 2.65 GB is allocated
									assert mem_bytes < 2.65 * 10**9






