

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,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


__lowercase     : Any       =						logging.get_logger(__name__)


if is_vision_available():
     import PIL



class 						_A     (				snake_case				):





  '''simple docstring'''



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





  def __init__(  self   ,SCREAMING_SNAKE_CASE_ = True   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = PILImageResampling.BICUBIC   ,SCREAMING_SNAKE_CASE_ = True   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = True   ,SCREAMING_SNAKE_CASE_ = 1 / 255   ,SCREAMING_SNAKE_CASE_ = True   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = True   ,**SCREAMING_SNAKE_CASE_   ,):

      '''simple docstring'''
      super().__init__(**SCREAMING_SNAKE_CASE_    )
      snake_case							:		Optional[int]        =		size if size is not None else {"""shortest_edge""": 224}
      snake_case							:		List[Any]        =		get_size_dict(SCREAMING_SNAKE_CASE_   ,default_to_square=SCREAMING_SNAKE_CASE_    )
      snake_case							:		List[Any]        =		crop_size if crop_size is not None else {"""height""": 224, """width""": 224}
      snake_case							:		str        =		get_size_dict(SCREAMING_SNAKE_CASE_   ,default_to_square=SCREAMING_SNAKE_CASE_   ,param_name="""crop_size"""    )

      snake_case							:		Any        =		do_resize
      snake_case							:		int        =		size
      snake_case							:		List[Any]        =		resample
      snake_case							:		int        =		do_center_crop
      snake_case							:		str        =		crop_size
      snake_case							:		List[Any]        =		do_rescale
      snake_case							:		Union[str, Any]        =		rescale_factor
      snake_case							:		Any        =		do_normalize
      snake_case							:		Tuple        =		image_mean if image_mean is not None else OPENAI_CLIP_MEAN
      snake_case							:		Dict        =		image_std if image_std is not None else OPENAI_CLIP_STD
      snake_case							:		Any        =		do_convert_rgb





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_ = PILImageResampling.BICUBIC   ,SCREAMING_SNAKE_CASE_ = None   ,**SCREAMING_SNAKE_CASE_   ,):

      '''simple docstring'''
      snake_case							:		List[Any]        =		get_size_dict(SCREAMING_SNAKE_CASE_   ,default_to_square=SCREAMING_SNAKE_CASE_    )
      if "shortest_edge" not in size:
          raise ValueError(F"""The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}"""    )
      snake_case							:		Any        =		get_resize_output_image_size(SCREAMING_SNAKE_CASE_   ,size=size["""shortest_edge"""]   ,default_to_square=SCREAMING_SNAKE_CASE_    )
      return resize(SCREAMING_SNAKE_CASE_   ,size=SCREAMING_SNAKE_CASE_   ,resample=SCREAMING_SNAKE_CASE_   ,data_format=SCREAMING_SNAKE_CASE_   ,**SCREAMING_SNAKE_CASE_    )





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_ = None   ,**SCREAMING_SNAKE_CASE_   ,):

      '''simple docstring'''
      snake_case							:		List[str]        =		get_size_dict(SCREAMING_SNAKE_CASE_    )
      if "height" not in size or "width" not in size:
          raise ValueError(F"""The `size` parameter must contain the keys (height, width). Got {size.keys()}"""    )
      return center_crop(SCREAMING_SNAKE_CASE_   ,size=(size["""height"""], size["""width"""])   ,data_format=SCREAMING_SNAKE_CASE_   ,**SCREAMING_SNAKE_CASE_    )





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_ = None   ,**SCREAMING_SNAKE_CASE_   ,):

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





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_ = None   ,**SCREAMING_SNAKE_CASE_   ,):

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





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = None   ,SCREAMING_SNAKE_CASE_ = ChannelDimension.FIRST   ,**SCREAMING_SNAKE_CASE_   ,):

      '''simple docstring'''
      snake_case							:		Optional[int]        =		do_resize if do_resize is not None else self.do_resize
      snake_case							:		Union[str, Any]        =		size if size is not None else self.size
      snake_case							:		List[Any]        =		get_size_dict(SCREAMING_SNAKE_CASE_   ,param_name="""size"""   ,default_to_square=SCREAMING_SNAKE_CASE_    )
      snake_case							:		Any        =		resample if resample is not None else self.resample
      snake_case							:		Union[str, Any]        =		do_center_crop if do_center_crop is not None else self.do_center_crop
      snake_case							:		List[Any]        =		crop_size if crop_size is not None else self.crop_size
      snake_case							:		List[str]        =		get_size_dict(SCREAMING_SNAKE_CASE_   ,param_name="""crop_size"""   ,default_to_square=SCREAMING_SNAKE_CASE_    )
      snake_case							:		str        =		do_rescale if do_rescale is not None else self.do_rescale
      snake_case							:		Tuple        =		rescale_factor if rescale_factor is not None else self.rescale_factor
      snake_case							:		Optional[Any]        =		do_normalize if do_normalize is not None else self.do_normalize
      snake_case							:		Tuple        =		image_mean if image_mean is not None else self.image_mean
      snake_case							:		Tuple        =		image_std if image_std is not None else self.image_std
      snake_case							:		Dict        =		do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

      snake_case							:		List[Any]        =		make_list_of_images(SCREAMING_SNAKE_CASE_    )

      if not valid_images(SCREAMING_SNAKE_CASE_    ):
          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_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."""    )

      # PIL RGBA images are converted to RGB
      if do_convert_rgb:
          snake_case							:		List[str]        =		[convert_to_rgb(SCREAMING_SNAKE_CASE_    ) for image in images]

      # All transformations expect numpy arrays.
      snake_case							:		Optional[Any]        =		[to_numpy_array(SCREAMING_SNAKE_CASE_    ) for image in images]

      if do_resize:
          snake_case							:		Union[str, Any]        =		[self.resize(image=SCREAMING_SNAKE_CASE_   ,size=SCREAMING_SNAKE_CASE_   ,resample=SCREAMING_SNAKE_CASE_    ) for image in images]

      if do_center_crop:
          snake_case							:		List[Any]        =		[self.center_crop(image=SCREAMING_SNAKE_CASE_   ,size=SCREAMING_SNAKE_CASE_    ) for image in images]

      if do_rescale:
          snake_case							:		List[Any]        =		[self.rescale(image=SCREAMING_SNAKE_CASE_   ,scale=SCREAMING_SNAKE_CASE_    ) for image in images]

      if do_normalize:
          snake_case							:		List[Any]        =		[self.normalize(image=SCREAMING_SNAKE_CASE_   ,mean=SCREAMING_SNAKE_CASE_   ,std=SCREAMING_SNAKE_CASE_    ) for image in images]

      snake_case							:		Optional[Any]        =		[to_channel_dimension_format(SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_    ) for image in images]

      snake_case							:		Optional[int]        =		{"""pixel_values""": images}
      return BatchFeature(data=SCREAMING_SNAKE_CASE_   ,tensor_type=SCREAMING_SNAKE_CASE_    )




import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

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


enable_full_determinism()




class 			__magic_name__	(		__a   ,						__a   ,						unittest.TestCase     ):





   """simple docstring"""




   lowerCAmelCase   :     List[str]	      = StableDiffusionSAGPipeline
   lowerCAmelCase   :     List[Any]	      = TEXT_TO_IMAGE_PARAMS
   lowerCAmelCase   :     List[str]	      = TEXT_TO_IMAGE_BATCH_PARAMS
   lowerCAmelCase   :     str	      = TEXT_TO_IMAGE_IMAGE_PARAMS
   lowerCAmelCase   :     int	      = TEXT_TO_IMAGE_IMAGE_PARAMS
   lowerCAmelCase   :     int	      = False



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    torch.manual_seed(0				)
    _UpperCamelCase:  Dict    						=			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		,	)
    _UpperCamelCase:  List[Any]    						=			DDIMScheduler(
        beta_start=0.00085		,	beta_end=0.012		,	beta_schedule='''scaled_linear'''		,	clip_sample=_lowercase		,	set_alpha_to_one=_lowercase		,	)
    torch.manual_seed(0				)
    _UpperCamelCase:  Optional[int]    						=			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				)
    _UpperCamelCase:  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=1_000		,	)
    _UpperCamelCase:  Any    						=			CLIPTextModel(_lowercase				)
    _UpperCamelCase:  List[str]    						=			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				)

    _UpperCamelCase:  int    						=			{
        '''unet''': unet,
        '''scheduler''': scheduler,
        '''vae''': vae,
        '''text_encoder''': text_encoder,
        '''tokenizer''': tokenizer,
        '''safety_checker''': None,
        '''feature_extractor''': None,
    }
    return components



   def 	lowerCAmelCase    (						self       :	Dict		,	_lowercase       :	List[Any]		,	_lowercase       :	int=0				):



    """simple docstring"""



    if str(_lowercase				).startswith('''mps'''				):
     _UpperCamelCase:  Tuple    						=			torch.manual_seed(_lowercase				)
    else:
     _UpperCamelCase:  List[str]    						=			torch.Generator(device=_lowercase				).manual_seed(_lowercase				)
    _UpperCamelCase:  Dict    						=			{
        '''prompt''': '''.''',
        '''generator''': generator,
        '''num_inference_steps''': 2,
        '''guidance_scale''': 1.0,
        '''sag_scale''': 1.0,
        '''output_type''': '''numpy''',
    }
    return inputs



   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    super().test_inference_batch_single_identical(expected_max_diff=3E-3				)




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





   """simple docstring"""




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



    """simple docstring"""



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



   def 	lowerCAmelCase    (						self       :	int				):



    """simple docstring"""



    _UpperCamelCase:  Any    						=			StableDiffusionSAGPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''				)
    _UpperCamelCase:  Any    						=			sag_pipe.to(_lowercase				)
    sag_pipe.set_progress_bar_config(disable=_lowercase				)

    _UpperCamelCase:  Tuple    						=			'''.'''
    _UpperCamelCase:  Tuple    						=			torch.manual_seed(0				)
    _UpperCamelCase:  Union[str, Any]    						=			sag_pipe(
        [prompt]		,	generator=_lowercase		,	guidance_scale=7.5		,	sag_scale=1.0		,	num_inference_steps=20		,	output_type='''np'''				)

    _UpperCamelCase:  Optional[Any]    						=			output.images

    _UpperCamelCase:  Tuple    						=			image[0, -3:, -3:, -1]

    assert image.shape == (1, 512, 512, 3)
    _UpperCamelCase:  Any    						=			np.array([0.1568, 0.1738, 0.1695, 0.1693, 0.1507, 0.1705, 0.1547, 0.1751, 0.1949]				)

    assert np.abs(image_slice.flatten() - expected_slice				).max() < 5E-2



   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    _UpperCamelCase:  Optional[Any]    						=			StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''				)
    _UpperCamelCase:  int    						=			sag_pipe.to(_lowercase				)
    sag_pipe.set_progress_bar_config(disable=_lowercase				)

    _UpperCamelCase:  List[Any]    						=			'''.'''
    _UpperCamelCase:  str    						=			torch.manual_seed(0				)
    _UpperCamelCase:  Dict    						=			sag_pipe(
        [prompt]		,	generator=_lowercase		,	guidance_scale=7.5		,	sag_scale=1.0		,	num_inference_steps=20		,	output_type='''np'''				)

    _UpperCamelCase:  Union[str, Any]    						=			output.images

    _UpperCamelCase:  List[str]    						=			image[0, -3:, -3:, -1]

    assert image.shape == (1, 512, 512, 3)
    _UpperCamelCase:  int    						=			np.array([0.3459, 0.2876, 0.2537, 0.3002, 0.2671, 0.2160, 0.3026, 0.2262, 0.2371]				)

    assert np.abs(image_slice.flatten() - expected_slice				).max() < 5E-2



   def 	lowerCAmelCase    (						self       :	Any				):



    """simple docstring"""



    _UpperCamelCase:  Tuple    						=			StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''				)
    _UpperCamelCase:  int    						=			sag_pipe.to(_lowercase				)
    sag_pipe.set_progress_bar_config(disable=_lowercase				)

    _UpperCamelCase:  Any    						=			'''.'''
    _UpperCamelCase:  List[Any]    						=			torch.manual_seed(0				)
    _UpperCamelCase:  str    						=			sag_pipe(
        [prompt]		,	width=768		,	height=512		,	generator=_lowercase		,	guidance_scale=7.5		,	sag_scale=1.0		,	num_inference_steps=20		,	output_type='''np'''		,	)

    _UpperCamelCase:  str    						=			output.images

    assert image.shape == (1, 512, 768, 3)