




'''simple docstring'''

from typing import Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature
from ...image_transforms import get_image_size, pad, rescale, to_channel_dimension_format
from ...image_utils import ChannelDimension, ImageInput, make_list_of_images, to_numpy_array, valid_images
from ...utils import TensorType, logging


a       :						Any								=		logging.get_logger(__name__)



class 		SCREAMING_SNAKE_CASE__  (			_UpperCamelCase	):
			__SCREAMING_SNAKE_CASE    =			["""pixel_values"""]





			def __init__(      self						:    str      ,		a_						:    bool = True      ,		a_						:    Union[int, float] = 1 / 255      ,		a_						:    bool = True      ,		a_						:    int = 8      ,		**a_						:    Any      ,		):



						"""simple docstring"""




						super().__init__(**a_							)

						__snake_case			=		do_rescale
						__snake_case			=		rescale_factor
						__snake_case			=		do_pad
						__snake_case			=		pad_size





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



						"""simple docstring"""




						return rescale(a_      ,		scale=a_      ,		data_format=a_      ,		**a_							)





			def 							A (      self						:    Optional[int]      ,		a_						:    np.ndarray      ,		a_						:    int      ,		a_						:    Optional[Union[str, ChannelDimension]] = None							):



						"""simple docstring"""




						__snake_case		,     __snake_case			=		get_image_size(a_							)
						__snake_case			=		(old_height // size + 1) * size - old_height
						__snake_case			=		(old_width // size + 1) * size - old_width

						return pad(a_      ,		((0, pad_height), (0, pad_width))      ,		mode="symmetric"      ,		data_format=a_							)






			def 							A (      self						:    int      ,		a_						:    ImageInput      ,		a_						:    Optional[bool] = None      ,		a_						:    Optional[float] = None      ,		a_						:    Optional[bool] = None      ,		a_						:    Optional[int] = None      ,		a_						:    Optional[Union[str, TensorType]] = None      ,		a_						:    Union[str, ChannelDimension] = ChannelDimension.FIRST      ,		**a_						:    str      ,		):



						"""simple docstring"""




						__snake_case			=		do_rescale if do_rescale is not None else self.do_rescale
						__snake_case			=		rescale_factor if rescale_factor is not None else self.rescale_factor
						__snake_case			=		do_pad if do_pad is not None else self.do_pad
						__snake_case			=		pad_size if pad_size is not None else self.pad_size

						__snake_case			=		make_list_of_images(a_							)

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

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

						# All transformations expect numpy arrays.
						__snake_case			=		[to_numpy_array(a_							) for image in images]

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

						if do_pad:
									__snake_case			=		[self.pad(a_      ,		size=a_							) for image in images]

						__snake_case			=		[to_channel_dimension_format(a_      ,		a_							) for image in images]

						__snake_case			=		{"pixel_values": images}
						return BatchFeature(data=a_      ,		tensor_type=a_							)













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     _a							(	UpperCAmelCase__    ,							UpperCAmelCase__    ,							unittest.TestCase     ):






  """simple docstring"""





  A_						       = StableDiffusionSAGPipeline
  A_						       = TEXT_TO_IMAGE_PARAMS
  A_						       = TEXT_TO_IMAGE_BATCH_PARAMS
  A_						       = TEXT_TO_IMAGE_IMAGE_PARAMS
  A_						       = TEXT_TO_IMAGE_IMAGE_PARAMS
  A_						       = False




  def  _UpperCAmelCase  (  self			)      ->			Optional[Any]:
     torch.manual_seed(0			)
     UpperCamelCase_						=  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_						=  DDIMScheduler(
         beta_start=0.0_0_0_8_5		,		beta_end=0.0_1_2		,		beta_schedule='scaled_linear'		,		clip_sample=_UpperCAmelCase		,		set_alpha_to_one=_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		,		)
     torch.manual_seed(0			)
     UpperCamelCase_						=  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		,		)
     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  _UpperCAmelCase  (  self		,		_UpperCAmelCase		,		_UpperCAmelCase=0			)      ->			List[Any]:
     if str(_UpperCAmelCase			).startswith('mps'			):
        UpperCamelCase_						=  torch.manual_seed(_UpperCAmelCase			)
     else:
        UpperCamelCase_						=  torch.Generator(device=_UpperCAmelCase			).manual_seed(_UpperCAmelCase			)
     UpperCamelCase_						=  {
         'prompt': '.',
         'generator': generator,
         'num_inference_steps': 2,
         'guidance_scale': 1.0,
         'sag_scale': 1.0,
         'output_type': 'numpy',
     }
     return inputs




  def  _UpperCAmelCase  (  self			)      ->			Tuple:
     super().test_inference_batch_single_identical(expected_max_diff=3e-3			)




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






  """simple docstring"""





  def  _UpperCAmelCase  (  self			)      ->			Tuple:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()




  def  _UpperCAmelCase  (  self			)      ->			str:
     UpperCamelCase_						=  StableDiffusionSAGPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'			)
     UpperCamelCase_						=  sag_pipe.to(_UpperCAmelCase			)
     sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase			)

     UpperCamelCase_						=  '.'
     UpperCamelCase_						=  torch.manual_seed(0			)
     UpperCamelCase_						=  sag_pipe(
         [prompt]		,		generator=_UpperCAmelCase		,		guidance_scale=7.5		,		sag_scale=1.0		,		num_inference_steps=20		,		output_type='np'			)

     UpperCamelCase_						=  output.images

     UpperCamelCase_						=  image[0, -3:, -3:, -1]

     assert image.shape == (1, 512, 512, 3)
     UpperCamelCase_						=  np.array([0.1_5_6_8, 0.1_7_3_8, 0.1_6_9_5, 0.1_6_9_3, 0.1_5_0_7, 0.1_7_0_5, 0.1_5_4_7, 0.1_7_5_1, 0.1_9_4_9]			)

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




  def  _UpperCAmelCase  (  self			)      ->			List[Any]:
     UpperCamelCase_						=  StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'			)
     UpperCamelCase_						=  sag_pipe.to(_UpperCAmelCase			)
     sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase			)

     UpperCamelCase_						=  '.'
     UpperCamelCase_						=  torch.manual_seed(0			)
     UpperCamelCase_						=  sag_pipe(
         [prompt]		,		generator=_UpperCAmelCase		,		guidance_scale=7.5		,		sag_scale=1.0		,		num_inference_steps=20		,		output_type='np'			)

     UpperCamelCase_						=  output.images

     UpperCamelCase_						=  image[0, -3:, -3:, -1]

     assert image.shape == (1, 512, 512, 3)
     UpperCamelCase_						=  np.array([0.3_4_5_9, 0.2_8_7_6, 0.2_5_3_7, 0.3_0_0_2, 0.2_6_7_1, 0.2_1_6_0, 0.3_0_2_6, 0.2_2_6_2, 0.2_3_7_1]			)

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




  def  _UpperCAmelCase  (  self			)      ->			Dict:
     UpperCamelCase_						=  StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'			)
     UpperCamelCase_						=  sag_pipe.to(_UpperCAmelCase			)
     sag_pipe.set_progress_bar_config(disable=_UpperCAmelCase			)

     UpperCamelCase_						=  '.'
     UpperCamelCase_						=  torch.manual_seed(0			)
     UpperCamelCase_						=  sag_pipe(
         [prompt]		,		width=768		,		height=512		,		generator=_UpperCAmelCase		,		guidance_scale=7.5		,		sag_scale=1.0		,		num_inference_steps=20		,		output_type='np'		,		)

     UpperCamelCase_						=  output.images

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





