






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,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


__UpperCamelCase					:		List[Any]	  =						logging.get_logger(__name__)


class   __SCREAMING_SNAKE_CASE(	a_	):
							_UpperCAmelCase     =       ["pixel_values"]
							def __init__(     self:      Optional[Any]       ,	UpperCamelCase:      bool = True       ,	UpperCamelCase:      Optional[Dict[str, int]] = None       ,	UpperCamelCase:      PILImageResampling = PILImageResampling.BILINEAR       ,	UpperCamelCase:      bool = True       ,	UpperCamelCase:      Dict[str, int] = None       ,	UpperCamelCase:      bool = True       ,	UpperCamelCase:      Union[int, float] = 1 / 2_55       ,	UpperCamelCase:      bool = True       ,	UpperCamelCase:      Optional[Union[float, List[float]]] = None       ,	UpperCamelCase:      Optional[Union[float, List[float]]] = None       ,	**UpperCamelCase:      Any       ,	)		->		None:
											super().__init__(**UpperCamelCase						)
											snake_case__									=     size if size is not None else {'shortest_edge': 2_56}
											snake_case__									=     get_size_dict(UpperCamelCase       ,	default_to_square=UpperCamelCase						)
											snake_case__									=     crop_size if crop_size is not None else {'height': 2_24, 'width': 2_24}
											snake_case__									=     get_size_dict(UpperCamelCase						)
											snake_case__									=     do_resize
											snake_case__									=     size
											snake_case__									=     resample
											snake_case__									=     do_center_crop
											snake_case__									=     crop_size
											snake_case__									=     do_rescale
											snake_case__									=     rescale_factor
											snake_case__									=     do_normalize
											snake_case__									=     image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
											snake_case__									=     image_std if image_std is not None else IMAGENET_STANDARD_STD
							def lowerCAmelCase_ (     self:      Optional[Any]       ,	UpperCamelCase:      np.ndarray       ,	UpperCamelCase:      Dict[str, int]       ,	UpperCamelCase:      PILImageResampling = PILImageResampling.BICUBIC       ,	UpperCamelCase:      Optional[Union[str, ChannelDimension]] = None       ,	**UpperCamelCase:      List[Any]       ,	)		->		np.ndarray:
											snake_case__									=     get_size_dict(UpperCamelCase       ,	default_to_square=UpperCamelCase						)
											if "shortest_edge" not in size:
															raise ValueError(F'''The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}'''						)
											snake_case__									=     get_resize_output_image_size(UpperCamelCase       ,	size=size['shortest_edge']       ,	default_to_square=UpperCamelCase						)
											return resize(UpperCamelCase       ,	size=UpperCamelCase       ,	resample=UpperCamelCase       ,	data_format=UpperCamelCase       ,	**UpperCamelCase						)
							def lowerCAmelCase_ (     self:      str       ,	UpperCamelCase:      np.ndarray       ,	UpperCamelCase:      Dict[str, int]       ,	UpperCamelCase:      Optional[Union[str, ChannelDimension]] = None       ,	**UpperCamelCase:      Any       ,	)		->		np.ndarray:
											snake_case__									=     get_size_dict(UpperCamelCase						)
											return center_crop(UpperCamelCase       ,	size=(size['height'], size['width'])       ,	data_format=UpperCamelCase       ,	**UpperCamelCase						)
							def lowerCAmelCase_ (     self:      int       ,	UpperCamelCase:      np.ndarray       ,	UpperCamelCase:      float       ,	UpperCamelCase:      Optional[Union[str, ChannelDimension]] = None       ,	**UpperCamelCase:      Tuple						)		->		np.ndarray:
											return rescale(UpperCamelCase       ,	scale=UpperCamelCase       ,	data_format=UpperCamelCase       ,	**UpperCamelCase						)
							def lowerCAmelCase_ (     self:      int       ,	UpperCamelCase:      np.ndarray       ,	UpperCamelCase:      Union[float, List[float]]       ,	UpperCamelCase:      Union[float, List[float]]       ,	UpperCamelCase:      Optional[Union[str, ChannelDimension]] = None       ,	**UpperCamelCase:      List[Any]       ,	)		->		np.ndarray:
											return normalize(UpperCamelCase       ,	mean=UpperCamelCase       ,	std=UpperCamelCase       ,	data_format=UpperCamelCase       ,	**UpperCamelCase						)


							def lowerCAmelCase_ (     self:      Optional[int]       ,	UpperCamelCase:      ImageInput       ,	UpperCamelCase:      Optional[bool] = None       ,	UpperCamelCase:      Dict[str, int] = None       ,	UpperCamelCase:      PILImageResampling = None       ,	UpperCamelCase:      bool = None       ,	UpperCamelCase:      Dict[str, int] = None       ,	UpperCamelCase:      Optional[bool] = None       ,	UpperCamelCase:      Optional[float] = None       ,	UpperCamelCase:      Optional[bool] = None       ,	UpperCamelCase:      Optional[Union[float, List[float]]] = None       ,	UpperCamelCase:      Optional[Union[float, List[float]]] = None       ,	UpperCamelCase:      Optional[Union[str, TensorType]] = None       ,	UpperCamelCase:      Union[str, ChannelDimension] = ChannelDimension.FIRST       ,	**UpperCamelCase:      Tuple       ,	)		->		Tuple:
											snake_case__									=     do_resize if do_resize is not None else self.do_resize
											snake_case__									=     size if size is not None else self.size
											snake_case__									=     get_size_dict(UpperCamelCase       ,	default_to_square=UpperCamelCase						)
											snake_case__									=     resample if resample is not None else self.resample
											snake_case__									=     do_center_crop if do_center_crop is not None else self.do_center_crop
											snake_case__									=     crop_size if crop_size is not None else self.crop_size
											snake_case__									=     get_size_dict(UpperCamelCase						)
											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_normalize if do_normalize is not None else self.do_normalize
											snake_case__									=     image_mean if image_mean is not None else self.image_mean
											snake_case__									=     image_std if image_std is not None else self.image_std

											snake_case__									=     make_list_of_images(UpperCamelCase						)

											if not valid_images(UpperCamelCase						):
															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.'						)

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

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

											if do_center_crop:
															snake_case__									=     [self.center_crop(image=UpperCamelCase       ,	size=UpperCamelCase						) for image in images]

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

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

											snake_case__									=     [to_channel_dimension_format(UpperCamelCase       ,	UpperCamelCase						) for image in images]

											snake_case__									=     {'pixel_values': images}
											return BatchFeature(data=UpperCamelCase       ,	tensor_type=UpperCamelCase						)












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   __SCREAMING_SNAKE_CASE(	a_    ,  a_    ,  unittest.TestCase	):
							_UpperCAmelCase     =       StableDiffusionSAGPipeline
							_UpperCAmelCase     =       TEXT_TO_IMAGE_PARAMS
							_UpperCAmelCase     =       TEXT_TO_IMAGE_BATCH_PARAMS
							_UpperCAmelCase     =       TEXT_TO_IMAGE_IMAGE_PARAMS
							_UpperCAmelCase     =       TEXT_TO_IMAGE_IMAGE_PARAMS
							_UpperCAmelCase     =       False
							def lowerCAmelCase_ (     self:      int						)		->		Any:
											torch.manual_seed(0						)
											snake_case__									=     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       ,	)
											snake_case__									=     DDIMScheduler(
											    beta_start=0.00_085       ,	beta_end=0.012       ,	beta_schedule='scaled_linear'       ,	clip_sample=UpperCamelCase       ,	set_alpha_to_one=UpperCamelCase       ,	)
											torch.manual_seed(0						)
											snake_case__									=     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						)
											snake_case__									=     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       ,	)
											snake_case__									=     CLIPTextModel(UpperCamelCase						)
											snake_case__									=     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)

											snake_case__									=     {
											    'unet': unet,
											    'scheduler': scheduler,
											    'vae': vae,
											    'text_encoder': text_encoder,
											    'tokenizer': tokenizer,
											    'safety_checker': None,
											    'feature_extractor': None,
											}
											return components
							def lowerCAmelCase_ (     self:      Optional[int]       ,	UpperCamelCase:      Dict       ,	UpperCamelCase:      Optional[int]=0						)		->		Union[str, Any]:
											if str(UpperCamelCase						).startswith('mps'						):
															snake_case__									=     torch.manual_seed(UpperCamelCase						)
											else:
															snake_case__									=     torch.Generator(device=UpperCamelCase						).manual_seed(UpperCamelCase						)
											snake_case__									=     {
											    'prompt': '.',
											    'generator': generator,
											    'num_inference_steps': 2,
											    'guidance_scale': 1.0,
											    'sag_scale': 1.0,
											    'output_type': 'numpy',
											}
											return inputs


							def lowerCAmelCase_ (     self:      Union[str, Any]						)		->		Dict:
											super().test_inference_batch_single_identical(expected_max_diff=3e-3						)


@slow
@require_torch_gpu
class   __SCREAMING_SNAKE_CASE(	unittest.TestCase	):
							def lowerCAmelCase_ (     self:      Optional[int]						)		->		Optional[int]:
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()
							def lowerCAmelCase_ (     self:      Dict						)		->		Tuple:
											snake_case__									=     StableDiffusionSAGPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'						)
											snake_case__									=     sag_pipe.to(UpperCamelCase						)
											sag_pipe.set_progress_bar_config(disable=UpperCamelCase						)

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

											snake_case__									=     output.images

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

											assert image.shape == (1, 5_12, 5_12, 3)
											snake_case__									=     np.array([0.1_568, 0.1_738, 0.1_695, 0.1_693, 0.1_507, 0.1_705, 0.1_547, 0.1_751, 0.1_949]						)

											assert np.abs(image_slice.flatten() - expected_slice						).max() < 5e-2
							def lowerCAmelCase_ (     self:      Union[str, Any]						)		->		List[Any]:
											snake_case__									=     StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'						)
											snake_case__									=     sag_pipe.to(UpperCamelCase						)
											sag_pipe.set_progress_bar_config(disable=UpperCamelCase						)

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

											snake_case__									=     output.images

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

											assert image.shape == (1, 5_12, 5_12, 3)
											snake_case__									=     np.array([0.3_459, 0.2_876, 0.2_537, 0.3_002, 0.2_671, 0.2_160, 0.3_026, 0.2_262, 0.2_371]						)

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


							def lowerCAmelCase_ (     self:      Union[str, Any]						)		->		Optional[Any]:
											snake_case__									=     StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'						)
											snake_case__									=     sag_pipe.to(UpperCamelCase						)
											sag_pipe.set_progress_bar_config(disable=UpperCamelCase						)

											snake_case__									=     '.'
											snake_case__									=     torch.manual_seed(0						)
											snake_case__									=     sag_pipe(
											    [prompt]       ,	width=7_68       ,	height=5_12       ,	generator=UpperCamelCase       ,	guidance_scale=7.5       ,	sag_scale=1.0       ,	num_inference_steps=20       ,	output_type='np'       ,	)

											snake_case__									=     output.images

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




