import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






@skip_mps
class 				_SCREAMING_SNAKE_CASE( __lowerCamelCase							,   __lowerCamelCase							,   __lowerCamelCase							,   unittest.TestCase						):
   UpperCamelCase_					:    Union[str, Any]			      = StableDiffusionModelEditingPipeline
   UpperCamelCase_					:    Union[str, Any]			      = TEXT_TO_IMAGE_PARAMS
   UpperCamelCase_					:    Optional[Any]			      = TEXT_TO_IMAGE_BATCH_PARAMS
   UpperCamelCase_					:    Optional[Any]			      = TEXT_TO_IMAGE_IMAGE_PARAMS
   UpperCamelCase_					:    Optional[int]			      = TEXT_TO_IMAGE_IMAGE_PARAMS
   def _lowercase(     self     )						->Union[str, Any]:
        torch.manual_seed(0     )
        a	:     List[str]						      =    UNetaDConditionModel(
            block_out_channels=(32, 64)				,  layers_per_block=2				,  sample_size=32				,  in_channels=4				,  out_channels=4				,  down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')				,  up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')				,  cross_attention_dim=32				,  )
        a	:     Tuple						      =    DDIMScheduler()
        torch.manual_seed(0     )
        a	:     Optional[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				,  )
        torch.manual_seed(0     )
        a	:     List[str]						      =    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=10_00				,  )
        a	:     Tuple						      =    CLIPTextModel(_lowercase     )
        a	:     Union[str, Any]						      =    CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )

        a	:     Optional[int]						      =    {
            '''unet''': unet,
            '''scheduler''': scheduler,
            '''vae''': vae,
            '''text_encoder''': text_encoder,
            '''tokenizer''': tokenizer,
            '''safety_checker''': None,
            '''feature_extractor''': None,
        }
        return components
   def _lowercase(     self				,  lowerCAmelCase_				,  lowerCAmelCase_=0     )						->List[Any]:
        a	:     str						      =    torch.manual_seed(_lowercase     )
        a	:     int						      =    {
            '''prompt''': '''A field of roses''',
            '''generator''': generator,
            # Setting height and width to None to prevent OOMs on CPU.
            '''height''': None,
            '''width''': None,
            '''num_inference_steps''': 2,
            '''guidance_scale''': 6.0,
            '''output_type''': '''numpy''',
        }
        return inputs
   def _lowercase(     self     )						->Union[str, Any]:
        a	:     Any						      =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        a	:     List[str]						      =    self.get_dummy_components()
        a	:     Tuple						      =    StableDiffusionModelEditingPipeline(**_lowercase     )
        a	:     Tuple						      =    sd_pipe.to(_lowercase     )
        sd_pipe.set_progress_bar_config(disable=_lowercase     )

        a	:     str						      =    self.get_dummy_inputs(_lowercase     )
        a	:     List[Any]						      =    sd_pipe(**_lowercase     ).images
        a	:     Any						      =    image[0, -3:, -3:, -1]
        assert image.shape == (1, 64, 64, 3)

        a	:     Any						      =    np.array([0.47_55, 0.51_32, 0.49_76, 0.39_04, 0.35_54, 0.47_65, 0.51_39, 0.51_58, 0.48_89]     )

        assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1e-2
   def _lowercase(     self     )						->List[str]:
        a	:     List[Any]						      =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        a	:     Dict						      =    self.get_dummy_components()
        a	:     int						      =    StableDiffusionModelEditingPipeline(**_lowercase     )
        a	:     Union[str, Any]						      =    sd_pipe.to(_lowercase     )
        sd_pipe.set_progress_bar_config(disable=_lowercase     )

        a	:     List[Any]						      =    self.get_dummy_inputs(_lowercase     )
        a	:     List[Any]						      =    '''french fries'''
        a	:     Tuple						      =    sd_pipe(**_lowercase				,  negative_prompt=_lowercase     )
        a	:     List[str]						      =    output.images
        a	:     Union[str, Any]						      =    image[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)

        a	:     Any						      =    np.array([0.49_92, 0.51_01, 0.50_04, 0.39_49, 0.36_04, 0.47_35, 0.52_16, 0.52_04, 0.49_13]     )

        assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1e-2
   def _lowercase(     self     )						->Tuple:
        a	:     int						      =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        a	:     Tuple						      =    self.get_dummy_components()
        a	:     Optional[int]						      =    EulerAncestralDiscreteScheduler(
            beta_start=0.0_00_85				,  beta_end=0.0_12				,  beta_schedule='''scaled_linear'''     )
        a	:     Any						      =    StableDiffusionModelEditingPipeline(**_lowercase     )
        a	:     Any						      =    sd_pipe.to(_lowercase     )
        sd_pipe.set_progress_bar_config(disable=_lowercase     )

        a	:     str						      =    self.get_dummy_inputs(_lowercase     )
        a	:     Dict						      =    sd_pipe(**_lowercase     ).images
        a	:     str						      =    image[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)

        a	:     Dict						      =    np.array([0.47_47, 0.53_72, 0.47_79, 0.49_82, 0.55_43, 0.48_16, 0.52_38, 0.49_04, 0.50_27]     )

        assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1e-2
   def _lowercase(     self     )						->Optional[int]:
        a	:     Any						      =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        a	:     str						      =    self.get_dummy_components()
        a	:     List[Any]						      =    PNDMScheduler()
        a	:     str						      =    StableDiffusionModelEditingPipeline(**_lowercase     )
        a	:     str						      =    sd_pipe.to(_lowercase     )
        sd_pipe.set_progress_bar_config(disable=_lowercase     )

        a	:     Any						      =    self.get_dummy_inputs(_lowercase     )
        # the pipeline does not expect pndm so test if it raises error.
        with self.assertRaises(_lowercase     ):
             a	:     int						      =    sd_pipe(**_lowercase     ).images
   def _lowercase(     self     )						->Optional[Any]:
        super().test_inference_batch_single_identical(expected_max_diff=5e-3     )


   def _lowercase(     self     )						->Optional[int]:
        super().test_attention_slicing_forward_pass(expected_max_diff=5e-3     )






@slow
@require_torch_gpu
class 				_SCREAMING_SNAKE_CASE( unittest.TestCase						):
   def _lowercase(     self     )						->List[str]:
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()
   def _lowercase(     self				,  lowerCAmelCase_=0     )						->int:
        a	:     List[Any]						      =    torch.manual_seed(_lowercase     )
        a	:     Union[str, Any]						      =    {
            '''prompt''': '''A field of roses''',
            '''generator''': generator,
            '''num_inference_steps''': 3,
            '''guidance_scale''': 7.5,
            '''output_type''': '''numpy''',
        }
        return inputs
   def _lowercase(     self     )						->List[Any]:
        a	:     Any						      =    '''CompVis/stable-diffusion-v1-4'''
        a	:     str						      =    StableDiffusionModelEditingPipeline.from_pretrained(_lowercase				,  safety_checker=_lowercase     )
        pipe.to(_lowercase     )
        pipe.set_progress_bar_config(disable=_lowercase     )
        pipe.enable_attention_slicing()

        a	:     int						      =    self.get_inputs()
        a	:     int						      =    pipe(**_lowercase     ).images
        a	:     Any						      =    image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 5_12, 5_12, 3)

        a	:     int						      =    np.array(
            [0.6_74_94_96, 0.6_38_64_53, 0.51_44_32_67, 0.66_09_49_05, 0.61_92_12_15, 0.5_49_13_32, 0.5_74_44_17, 0.58_07_51_06, 0.5_17_46_58]     )

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

        # make sure image changes after editing
        pipe.edit_model('''A pack of roses'''				,  '''A pack of blue roses'''     )

        a	:     Optional[Any]						      =    pipe(**_lowercase     ).images
        a	:     Optional[int]						      =    image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 5_12, 5_12, 3)

        assert np.abs(expected_slice - image_slice     ).max() > 1e-1


   def _lowercase(     self     )						->Dict:
        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        a	:     str						      =    '''CompVis/stable-diffusion-v1-4'''
        a	:     Any						      =    DDIMScheduler.from_pretrained(_lowercase				,  subfolder='''scheduler'''     )
        a	:     int						      =    StableDiffusionModelEditingPipeline.from_pretrained(
            _lowercase				,  scheduler=_lowercase				,  safety_checker=_lowercase     )
        a	:     int						      =    pipe.to(_lowercase     )
        pipe.set_progress_bar_config(disable=_lowercase     )
        pipe.enable_attention_slicing(1     )
        pipe.enable_sequential_cpu_offload()

        a	:     int						      =    self.get_inputs()
        a	:     List[str]						      =    pipe(**_lowercase     )

        a	:     List[str]						      =    torch.cuda.max_memory_allocated()
        # make sure that less than 4.4 GB is allocated
        assert mem_bytes < 4.4 * 10**9
"""simple docstring"""







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 (
    get_resize_output_image_size,
    normalize,
    pad,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging
from ...utils.import_utils import is_vision_available


UpperCamelCase_							:							Optional[int]				 =		logging.get_logger(__name__)


if is_vision_available():
						import PIL


class       _lowercase		( __lowerCamelCase	):
				lowerCamelCase__							=							['pixel_values']

				def __init__(  self							:     str							,     _lowercase							:     bool = True							,     _lowercase							:     Dict[str, int] = None							,     _lowercase							:     PILImageResampling = PILImageResampling.BILINEAR							,     _lowercase							:     bool = True							,     _lowercase							:     bool = False							,     _lowercase							:     bool = True							,     _lowercase							:     bool = True							,     _lowercase							:     Union[int, float] = 1 / 2_5_5							,     _lowercase							:     bool = True							,     _lowercase							:     Optional[Union[float, List[float]]] = None							,     _lowercase							:     Optional[Union[float, List[float]]] = None							,     **_lowercase							:     List[str]							,     ):


										'''simple docstring'''
										super().__init__(**_lowercase  )

										UpperCAmelCase__       						=		size if size is not None else {"""height""": 2_5_6_0, """width""": 1_9_2_0}
										if isinstance(_lowercase							,     (tuple, list)  ):
																# The previous feature extractor size parameter was in (width, height) format
																UpperCAmelCase__       						=		size[::-1]
										UpperCAmelCase__       						=		get_size_dict(_lowercase  )

										UpperCAmelCase__       						=		do_resize
										UpperCAmelCase__       						=		size
										UpperCAmelCase__       						=		resample
										UpperCAmelCase__       						=		do_thumbnail
										UpperCAmelCase__       						=		do_align_long_axis
										UpperCAmelCase__       						=		do_pad
										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							:     str							,     _lowercase							:     np.ndarray							,     _lowercase							:     Dict[str, int]							,     _lowercase							:     Optional[Union[str, ChannelDimension]] = None  ):


										'''simple docstring'''
										UpperCAmelCase__						,   UpperCAmelCase__       						=		get_image_size(_lowercase  )
										UpperCAmelCase__						,   UpperCAmelCase__       						=		size["""height"""], size["""width"""]

										if (output_width < output_height and input_width > input_height) or (
										    output_width > output_height and input_width < input_height
										):
																UpperCAmelCase__       						=		np.rotaa(_lowercase							,     3  )

										if data_format is not None:
																UpperCAmelCase__       						=		to_channel_dimension_format(_lowercase							,     _lowercase  )

										return image

				def 							A			(  self							:     Tuple							,     _lowercase							:     np.ndarray							,     _lowercase							:     Dict[str, int]							,     _lowercase							:     bool = False							,     _lowercase							:     Optional[Union[str, ChannelDimension]] = None							,     ):


										'''simple docstring'''
										UpperCAmelCase__						,   UpperCAmelCase__       						=		size["""height"""], size["""width"""]
										UpperCAmelCase__						,   UpperCAmelCase__       						=		get_image_size(_lowercase  )

										UpperCAmelCase__       						=		output_width - input_width
										UpperCAmelCase__       						=		output_height - input_height

										if random_padding:
																UpperCAmelCase__       						=		np.random.randint(low=0							,     high=delta_height + 1  )
																UpperCAmelCase__       						=		np.random.randint(low=0							,     high=delta_width + 1  )
										else:
																UpperCAmelCase__       						=		delta_height // 2
																UpperCAmelCase__       						=		delta_width // 2

										UpperCAmelCase__       						=		delta_height - pad_top
										UpperCAmelCase__       						=		delta_width - pad_left

										UpperCAmelCase__       						=		((pad_top, pad_bottom), (pad_left, pad_right))
										return pad(_lowercase							,     _lowercase							,     data_format=_lowercase  )

				def 							A			(  self							:     Union[str, Any]							,     *_lowercase							:     List[str]							,     **_lowercase							:     int  ):


										'''simple docstring'''
										logger.info("""pad is deprecated and will be removed in version 4.27. Please use pad_image instead."""  )
										return self.pad_image(*_lowercase							,     **_lowercase  )

				def 							A			(  self							:     Optional[int]							,     _lowercase							:     np.ndarray							,     _lowercase							:     Dict[str, int]							,     _lowercase							:     PILImageResampling = PILImageResampling.BICUBIC							,     _lowercase							:     Optional[Union[str, ChannelDimension]] = None							,     **_lowercase							:     Dict							,     ):


										'''simple docstring'''
										UpperCAmelCase__						,   UpperCAmelCase__       						=		get_image_size(_lowercase  )
										UpperCAmelCase__						,   UpperCAmelCase__       						=		size["""height"""], size["""width"""]

										# We always resize to the smallest of either the input or output size.
										UpperCAmelCase__       						=		min(_lowercase							,     _lowercase  )
										UpperCAmelCase__       						=		min(_lowercase							,     _lowercase  )

										if height == input_height and width == input_width:
																return image

										if input_height > input_width:
																UpperCAmelCase__       						=		int(input_width * height / input_height  )
										elif input_width > input_height:
																UpperCAmelCase__       						=		int(input_height * width / input_width  )

										return resize(
										    _lowercase							,     size=(height, width)							,     resample=_lowercase							,     reducing_gap=2.0							,     data_format=_lowercase							,     **_lowercase  )

				def 							A			(  self							:     List[str]							,     _lowercase							:     np.ndarray							,     _lowercase							:     Dict[str, int]							,     _lowercase							:     PILImageResampling = PILImageResampling.BICUBIC							,     _lowercase							:     Optional[Union[str, ChannelDimension]] = None							,     **_lowercase							:     Optional[Any]							,     ):


										'''simple docstring'''
										UpperCAmelCase__       						=		get_size_dict(_lowercase  )
										UpperCAmelCase__       						=		min(size["""height"""]							,     size["""width"""]  )
										UpperCAmelCase__       						=		get_resize_output_image_size(_lowercase							,     size=_lowercase							,     default_to_square=_lowercase  )
										UpperCAmelCase__       						=		resize(_lowercase							,     size=_lowercase							,     resample=_lowercase							,     data_format=_lowercase							,     **_lowercase  )
										return resized_image

				def 							A			(  self							:     int							,     _lowercase							:     np.ndarray							,     _lowercase							:     Union[int, float]							,     _lowercase							:     Optional[Union[str, ChannelDimension]] = None							,     **_lowercase							:     Union[str, Any]							,     ):


										'''simple docstring'''
										return rescale(_lowercase							,     scale=_lowercase							,     data_format=_lowercase							,     **_lowercase  )

				def 							A			(  self							:     Optional[int]							,     _lowercase							:     np.ndarray							,     _lowercase							:     Union[float, List[float]]							,     _lowercase							:     Union[float, List[float]]							,     _lowercase							:     Optional[Union[str, ChannelDimension]] = None							,     **_lowercase							:     Tuple							,     ):


										'''simple docstring'''
										return normalize(_lowercase							,     mean=_lowercase							,     std=_lowercase							,     data_format=_lowercase							,     **_lowercase  )

				def 							A			(  self							:     List[str]							,     _lowercase							:     ImageInput							,     _lowercase							:     bool = None							,     _lowercase							:     Dict[str, int] = None							,     _lowercase							:     PILImageResampling = None							,     _lowercase							:     bool = None							,     _lowercase							:     bool = None							,     _lowercase							:     bool = None							,     _lowercase							:     bool = False							,     _lowercase							:     bool = None							,     _lowercase							:     float = None							,     _lowercase							:     bool = None							,     _lowercase							:     Optional[Union[float, List[float]]] = None							,     _lowercase							:     Optional[Union[float, List[float]]] = None							,     _lowercase							:     Optional[Union[str, TensorType]] = None							,     _lowercase							:     Optional[ChannelDimension] = ChannelDimension.FIRST							,     **_lowercase							:     Dict							,     ):


										'''simple docstring'''
										UpperCAmelCase__       						=		do_resize if do_resize is not None else self.do_resize
										UpperCAmelCase__       						=		size if size is not None else self.size
										if isinstance(_lowercase							,     (tuple, list)  ):
																# Previous feature extractor had size in (width, height) format
																UpperCAmelCase__       						=		size[::-1]
										UpperCAmelCase__       						=		get_size_dict(_lowercase  )
										UpperCAmelCase__       						=		resample if resample is not None else self.resample
										UpperCAmelCase__       						=		do_thumbnail if do_thumbnail is not None else self.do_thumbnail
										UpperCAmelCase__       						=		do_align_long_axis if do_align_long_axis is not None else self.do_align_long_axis
										UpperCAmelCase__       						=		do_pad if do_pad is not None else self.do_pad
										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__       						=		make_list_of_images(_lowercase  )

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

										if do_resize and size is None:
																raise ValueError("""Size must be specified if do_resize is True."""  )

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

										if do_pad and size is None:
																raise ValueError("""Size must be specified if do_pad 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(_lowercase  ) for image in images]

										if do_align_long_axis:
																UpperCAmelCase__       						=		[self.align_long_axis(_lowercase							,     size=_lowercase  ) for image in images]

										if do_resize:
																UpperCAmelCase__       						=		[self.resize(image=_lowercase							,     size=_lowercase							,     resample=_lowercase  ) for image in images]

										if do_thumbnail:
																UpperCAmelCase__       						=		[self.thumbnail(image=_lowercase							,     size=_lowercase  ) for image in images]

										if do_pad:
																UpperCAmelCase__       						=		[self.pad_image(image=_lowercase							,     size=_lowercase							,     random_padding=_lowercase  ) for image in images]

										if do_rescale:
																UpperCAmelCase__       						=		[self.rescale(image=_lowercase							,     scale=_lowercase  ) for image in images]

										if do_normalize:
																UpperCAmelCase__       						=		[self.normalize(image=_lowercase							,     mean=_lowercase							,     std=_lowercase  ) for image in images]

										UpperCAmelCase__       						=		[to_channel_dimension_format(_lowercase							,     _lowercase  ) for image in images]

										UpperCAmelCase__       						=		{"""pixel_values""": images}
										return BatchFeature(data=_lowercase							,     tensor_type=_lowercase  )

