'''simple docstring'''




import warnings
from typing import List, Optional, Union

import numpy as np
import PIL
import torch
from PIL import Image

from .configuration_utils import ConfigMixin, register_to_config
from .utils import CONFIG_NAME, PIL_INTERPOLATION, deprecate




class       __UpperCAmelCase							(							lowerCamelCase__	):





       '''simple docstring'''



       __lowercase   :   List[Any]  					=							CONFIG_NAME

       @register_to_config
       def __init__( self       ,	_SCREAMING_SNAKE_CASE = True       ,	_SCREAMING_SNAKE_CASE = 8       ,	_SCREAMING_SNAKE_CASE = "lanczos"       ,	_SCREAMING_SNAKE_CASE = True       ,	_SCREAMING_SNAKE_CASE = False       ,	)						->	Any:
           super().__init__()

       @staticmethod
       def    __A			( _SCREAMING_SNAKE_CASE					)						->	PIL.Image.Image:





           if images.ndim == 3:
               A_      =  images[None, ...]
           A_      =  (images * 255).round().astype('''uint8'''					)
           if images.shape[-1] == 1:
               # special case for grayscale (single channel) images
               A_      =  [Image.fromarray(image.squeeze()       ,	mode='''L'''					) for image in images]
           else:
               A_      =  [Image.fromarray(snake_case__					) for image in images]

           return pil_images

       @staticmethod
       def    __A			( _SCREAMING_SNAKE_CASE					)						->	np.ndarray:





           if not isinstance(snake_case__       ,	snake_case__					):
               A_      =  [images]
           A_      =  [np.array(snake_case__					).astype(np.floataa					) / 255.0 for image in images]
           A_      =  np.stack(snake_case__       ,	axis=0					)

           return images

       @staticmethod
       def    __A			( _SCREAMING_SNAKE_CASE					)						->	torch.FloatTensor:





           if images.ndim == 3:
               A_      =  images[..., None]

           A_      =  torch.from_numpy(images.transpose(0       ,	3       ,	1       ,	2					)					)
           return images

       @staticmethod
       def    __A			( _SCREAMING_SNAKE_CASE					)						->	np.ndarray:
           A_      =  images.cpu().permute(0       ,	2       ,	3       ,	1					).float().numpy()
           return images

       @staticmethod
       def    __A			( _SCREAMING_SNAKE_CASE					)						->	Any:
           return 2.0 * images - 1.0

       @staticmethod
       def    __A			( _SCREAMING_SNAKE_CASE					)						->	Tuple:
           return (images / 2 + 0.5).clamp(0       ,	1					)

       @staticmethod
       def    __A			( _SCREAMING_SNAKE_CASE					)						->	PIL.Image.Image:
           A_      =  image.convert('''RGB'''					)
           return image

       def    __A			( self       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE = None       ,	_SCREAMING_SNAKE_CASE = None       ,	)						->	PIL.Image.Image:





           if height is None:
               A_      =  image.height
           if width is None:
               A_      =  image.width

           A_      =  (
               x - x % self.config.vae_scale_factor for x in (width, height)
           )  # resize to integer multiple of vae_scale_factor
           A_      =  image.resize((width, height)       ,	resample=PIL_INTERPOLATION[self.config.resample]					)
           return image

       def    __A			( self       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE = None       ,	_SCREAMING_SNAKE_CASE = None       ,	)						->	torch.Tensor:
           A_      =  (PIL.Image.Image, np.ndarray, torch.Tensor)
           if isinstance(snake_case__       ,	snake_case__					):
               A_      =  [image]
           elif not (isinstance(snake_case__       ,	snake_case__					) and all(isinstance(snake_case__       ,	snake_case__					) for i in image					)):
               raise ValueError(
                   F'''Input is in incorrect format: {[type(snake_case__					) for i in image]}. Currently, we only support {', '.join(snake_case__					)}'''					)

           if isinstance(image[0]       ,	PIL.Image.Image					):
               if self.config.do_convert_rgb:
                   A_      =  [self.convert_to_rgb(snake_case__					) for i in image]
               if self.config.do_resize:
                   A_      =  [self.resize(snake_case__       ,	snake_case__       ,	snake_case__					) for i in image]
               A_      =  self.pil_to_numpy(snake_case__					)  # to np
               A_      =  self.numpy_to_pt(snake_case__					)  # to pt

           elif isinstance(image[0]       ,	np.ndarray					):
               A_      =  np.concatenate(snake_case__       ,	axis=0					) if image[0].ndim == 4 else np.stack(snake_case__       ,	axis=0					)
               A_      =  self.numpy_to_pt(snake_case__					)
               A_      =  image.shape
               if self.config.do_resize and (
                   height % self.config.vae_scale_factor != 0 or width % self.config.vae_scale_factor != 0
               ):
                   raise ValueError(
                       F'''Currently we only support resizing for PIL image - please resize your numpy array to be divisible by {self.config.vae_scale_factor}'''
                       F'''currently the sizes are {height} and {width}. You can also pass a PIL image instead to use resize option in VAEImageProcessor'''					)

           elif isinstance(image[0]       ,	torch.Tensor					):
               A_      =  torch.cat(snake_case__       ,	axis=0					) if image[0].ndim == 4 else torch.stack(snake_case__       ,	axis=0					)
               A_      =  image.shape

               # don't need any preprocess if the image is latents
               if channel == 4:
                   return image

               if self.config.do_resize and (
                   height % self.config.vae_scale_factor != 0 or width % self.config.vae_scale_factor != 0
               ):
                   raise ValueError(
                       F'''Currently we only support resizing for PIL image - please resize your pytorch tensor to be divisible by {self.config.vae_scale_factor}'''
                       F'''currently the sizes are {height} and {width}. You can also pass a PIL image instead to use resize option in VAEImageProcessor'''					)

        # expected range [0,1], normalize to [-1,1]
           A_      =  self.config.do_normalize
           if image.min() < 0:
               warnings.warn(
                   '''Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] '''
                   F'''when passing as pytorch tensor or numpy Array. You passed `image` with value range [{image.min()},{image.max()}]'''       ,	snake_case__       ,	)
               A_      =  False

           if do_normalize:
               A_      =  self.normalize(snake_case__					)

           return image

       def    __A			( self       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE = "pil"       ,	_SCREAMING_SNAKE_CASE = None       ,	)						->	str:





           if not isinstance(snake_case__       ,	torch.Tensor					):
               raise ValueError(
                   F'''Input for postprocessing is in incorrect format: {type(snake_case__					)}. We only support pytorch tensor'''					)
           if output_type not in ["latent", "pt", "np", "pil"]:
               A_      =  (
                   F'''the output_type {output_type} is outdated and has been set to `np`. Please make sure to set it to one of these instead: '''
                   '''`pil`, `np`, `pt`, `latent`'''
               )
               deprecate('''Unsupported output_type'''       ,	'''1.0.0'''       ,	snake_case__       ,	standard_warn=snake_case__					)
               A_      =  '''np'''

           if output_type == "latent":
               return image

           if do_denormalize is None:
               A_      =  [self.config.do_normalize] * image.shape[0]

           A_      =  torch.stack(
               [self.denormalize(image[i]					) if do_denormalize[i] else image[i] for i in range(image.shape[0]					)]					)

           if output_type == "pt":
               return image

           A_      =  self.pt_to_numpy(snake_case__					)

           if output_type == "np":
               return image

           if output_type == "pil":
               return self.numpy_to_pil(snake_case__					)




class       __UpperCAmelCase							(							lowerCamelCase__	):





       '''simple docstring'''



       __lowercase   :   int  					=							CONFIG_NAME

       @register_to_config
       def __init__( self       ,	_SCREAMING_SNAKE_CASE = True       ,	_SCREAMING_SNAKE_CASE = 8       ,	_SCREAMING_SNAKE_CASE = "lanczos"       ,	_SCREAMING_SNAKE_CASE = True       ,	)						->	Dict:
           super().__init__()

       @staticmethod
       def    __A			( _SCREAMING_SNAKE_CASE					)						->	List[Any]:





           if images.ndim == 3:
               A_      =  images[None, ...]
           A_      =  (images * 255).round().astype('''uint8'''					)
           if images.shape[-1] == 1:
               # special case for grayscale (single channel) images
               A_      =  [Image.fromarray(image.squeeze()       ,	mode='''L'''					) for image in images]
           else:
               A_      =  [Image.fromarray(image[:, :, :3]					) for image in images]

           return pil_images

       @staticmethod
       def    __A			( _SCREAMING_SNAKE_CASE					)						->	Optional[Any]:
           return image[:, :, 1] * 2**8 + image[:, :, 2]

       def    __A			( self       ,	_SCREAMING_SNAKE_CASE					)						->	Tuple:





           if images.ndim == 3:
               A_      =  images[None, ...]
           A_      =  images[:, :, :, 3:]
           if images.shape[-1] == 6:
               A_      =  (images_depth * 255).round().astype('''uint8'''					)
               A_      =  [
                   Image.fromarray(self.rgblike_to_depthmap(snake_case__					)       ,	mode='''I;16'''					) for image_depth in images_depth
               ]
           elif images.shape[-1] == 4:
               A_      =  (images_depth * 6_5535.0).astype(np.uintaa					)
               A_      =  [Image.fromarray(snake_case__       ,	mode='''I;16'''					) for image_depth in images_depth]
           else:
               raise Exception('''Not supported'''					)

           return pil_images

       def    __A			( self       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE = "pil"       ,	_SCREAMING_SNAKE_CASE = None       ,	)						->	Dict:





           if not isinstance(snake_case__       ,	torch.Tensor					):
               raise ValueError(
                   F'''Input for postprocessing is in incorrect format: {type(snake_case__					)}. We only support pytorch tensor'''					)
           if output_type not in ["latent", "pt", "np", "pil"]:
               A_      =  (
                   F'''the output_type {output_type} is outdated and has been set to `np`. Please make sure to set it to one of these instead: '''
                   '''`pil`, `np`, `pt`, `latent`'''
               )
               deprecate('''Unsupported output_type'''       ,	'''1.0.0'''       ,	snake_case__       ,	standard_warn=snake_case__					)
               A_      =  '''np'''

           if do_denormalize is None:
               A_      =  [self.config.do_normalize] * image.shape[0]

           A_      =  torch.stack(
               [self.denormalize(image[i]					) if do_denormalize[i] else image[i] for i in range(image.shape[0]					)]					)

           A_      =  self.pt_to_numpy(snake_case__					)

           if output_type == "np":
               if image.shape[-1] == 6:
                   A_      =  np.stack([self.rgblike_to_depthmap(im[:, :, 3:]					) for im in image]       ,	axis=0					)
               else:
                   A_      =  image[:, :, :, 3:]
               return image[:, :, :, :3], image_depth

           if output_type == "pil":
               return self.numpy_to_pil(snake_case__					), self.numpy_to_depth(snake_case__					)
           else:
               raise Exception(F'''This type {output_type} is not supported'''					)







import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLPipeline,
    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       __snake_case (  lowerCamelCase__       , lowerCamelCase__       , unittest.TestCase ):
 __lowerCamelCase     : int 	=   StableDiffusionXLPipeline
 __lowerCamelCase     : Any 	=   TEXT_TO_IMAGE_PARAMS
 __lowerCamelCase     : List[str] 	=   TEXT_TO_IMAGE_BATCH_PARAMS
 __lowerCamelCase     : Optional[int] 	=   TEXT_TO_IMAGE_IMAGE_PARAMS
 __lowerCamelCase     : Union[str, Any] 	=   TEXT_TO_IMAGE_IMAGE_PARAMS
 def 						UpperCAmelCase__							(	self    )						->				Optional[Any]:
       '''simple docstring'''





       torch.manual_seed(0    )
       UpperCAmelCase     :			Union[str, Any]     	=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''')		, attention_head_dim=(2, 4)		, use_linear_projection=snake_case__		, addition_embed_type='''text_time'''		, addition_time_embed_dim=8		, transformer_layers_per_block=(1, 2)		, projection_class_embeddings_input_dim=80		, cross_attention_dim=64		, )
       UpperCAmelCase     :			Optional[int]     	=EulerDiscreteScheduler(
           beta_start=0.0_0085		, beta_end=0.012		, steps_offset=1		, beta_schedule='''scaled_linear'''		, timestep_spacing='''leading'''		, )
       torch.manual_seed(0    )
       UpperCAmelCase     :			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		, sample_size=128		, )
       torch.manual_seed(0    )
       UpperCAmelCase     :			Optional[Any]     	=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=32		, )
       UpperCAmelCase     :			List[str]     	=CLIPTextModel(snake_case__    )
       UpperCAmelCase     :			int     	=CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		, local_files_only=snake_case__    )

       UpperCAmelCase     :			Tuple     	=CLIPTextModelWithProjection(snake_case__    )
       UpperCAmelCase     :			int     	=CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		, local_files_only=snake_case__    )

       UpperCAmelCase     :			Optional[int]     	={
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''vae''': vae,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''text_encoder_2''': text_encoder_a,
           '''tokenizer_2''': tokenizer_a,
           # "safety_checker": None,
           # "feature_extractor": None,
       }
       return components
 def 						UpperCAmelCase__							(	self		, snake_case__		, snake_case__=0    )						->				Optional[Any]:
       '''simple docstring'''





       if str(snake_case__    ).startswith('''mps'''    ):
             UpperCAmelCase     :			Union[str, Any]     	=torch.manual_seed(snake_case__    )
       else:
             UpperCAmelCase     :			Optional[int]     	=torch.Generator(device=snake_case__    ).manual_seed(snake_case__    )
       UpperCAmelCase     :			Any     	={
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''guidance_scale''': 5.0,
           '''output_type''': '''numpy''',
       }
       return inputs
 def 						UpperCAmelCase__							(	self    )						->				List[Any]:
       '''simple docstring'''





       UpperCAmelCase     :			Optional[Any]     	='''cpu'''  # ensure determinism for the device-dependent torch.Generator
       UpperCAmelCase     :			List[Any]     	=self.get_dummy_components()
       UpperCAmelCase     :			Tuple     	=StableDiffusionXLPipeline(**snake_case__    )
       UpperCAmelCase     :			Optional[Any]     	=sd_pipe.to(snake_case__    )
       sd_pipe.set_progress_bar_config(disable=snake_case__    )

       UpperCAmelCase     :			Optional[Any]     	=self.get_dummy_inputs(snake_case__    )
       UpperCAmelCase     :			Optional[int]     	=sd_pipe(**snake_case__    ).images
       UpperCAmelCase     :			str     	=image[0, -3:, -3:, -1]

       assert image.shape == (1, 64, 64, 3)
       UpperCAmelCase     :			Dict     	=np.array([0.5873, 0.6128, 0.4797, 0.5122, 0.5674, 0.4639, 0.5227, 0.5149, 0.4747]    )

       assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1e-2
 def 						UpperCAmelCase__							(	self    )						->				List[str]:
       '''simple docstring'''





       UpperCAmelCase     :			Optional[Any]     	=self.get_dummy_components()
       UpperCAmelCase     :			List[Any]     	=StableDiffusionXLPipeline(**snake_case__    )
       UpperCAmelCase     :			List[str]     	=sd_pipe.to(snake_case__    )
       UpperCAmelCase     :			Optional[int]     	=sd_pipe.to(snake_case__    )
       sd_pipe.set_progress_bar_config(disable=snake_case__    )

       # forward without prompt embeds
       UpperCAmelCase     :			int     	=self.get_dummy_inputs(snake_case__    )
       UpperCAmelCase     :			int     	=3 * ['''this is a negative prompt''']
       UpperCAmelCase     :			Union[str, Any]     	=negative_prompt
       UpperCAmelCase     :			Union[str, Any]     	=3 * [inputs['''prompt''']]

       UpperCAmelCase     :			int     	=sd_pipe(**snake_case__    )
       UpperCAmelCase     :			str     	=output.images[0, -3:, -3:, -1]

       # forward with prompt embeds
       UpperCAmelCase     :			Any     	=self.get_dummy_inputs(snake_case__    )
       UpperCAmelCase     :			List[str]     	=3 * ['''this is a negative prompt''']
       UpperCAmelCase     :			int     	=3 * [inputs.pop('''prompt'''    )]

       (
           (
           UpperCAmelCase
       )	,    (
           UpperCAmelCase
       )	,    (
           UpperCAmelCase
       )	,    (
           UpperCAmelCase
       )	,    
       )     :			Union[str, Any]     	=sd_pipe.encode_prompt(snake_case__		, negative_prompt=snake_case__    )

       UpperCAmelCase     :			Union[str, Any]     	=sd_pipe(
           **snake_case__		, prompt_embeds=snake_case__		, negative_prompt_embeds=snake_case__		, pooled_prompt_embeds=snake_case__		, negative_pooled_prompt_embeds=snake_case__		, )
       UpperCAmelCase     :			Union[str, Any]     	=output.images[0, -3:, -3:, -1]

       # make sure that it's equal
       assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()    ).max() < 1e-4
 def 						UpperCAmelCase__							(	self    )						->				List[Any]:
       '''simple docstring'''





       super().test_attention_slicing_forward_pass(expected_max_diff=3e-3    )






 def 						UpperCAmelCase__							(	self    )						->				Any:
       '''simple docstring'''





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







@slow
@require_torch_gpu
class       __snake_case (  unittest.TestCase ):
 def 						UpperCAmelCase__							(	self    )						->				List[str]:
       '''simple docstring'''





       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
 def 						UpperCAmelCase__							(	self		, snake_case__		, snake_case__="cpu"		, snake_case__=torch.floataa		, snake_case__=0    )						->				Any:
       '''simple docstring'''





       UpperCAmelCase     :			str     	=torch.Generator(device=snake_case__    ).manual_seed(snake_case__    )
       UpperCAmelCase     :			Optional[Any]     	=np.random.RandomState(snake_case__    ).standard_normal((1, 4, 64, 64)    )
       UpperCAmelCase     :			List[Any]     	=torch.from_numpy(snake_case__    ).to(device=snake_case__		, dtype=snake_case__    )
       UpperCAmelCase     :			Optional[Any]     	={
           '''prompt''': '''a photograph of an astronaut riding a horse''',
           '''latents''': latents,
           '''generator''': generator,
           '''num_inference_steps''': 3,
           '''guidance_scale''': 7.5,
           '''output_type''': '''numpy''',
       }
       return inputs






 def 						UpperCAmelCase__							(	self    )						->				str:
       '''simple docstring'''





       UpperCAmelCase     :			List[Any]     	=DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''    )
       pipe.to(snake_case__    )
       pipe.set_progress_bar_config(disable=snake_case__    )

       UpperCAmelCase     :			List[Any]     	=self.get_inputs(snake_case__    )
       UpperCAmelCase     :			str     	=pipe(**snake_case__    ).images
       UpperCAmelCase     :			int     	=image[0, -3:, -3:, -1].flatten()

       assert image.shape == (1, 512, 512, 3)
       UpperCAmelCase     :			Any     	=np.array([0.4_9493, 0.4_7896, 0.4_0798, 0.5_4214, 0.5_3212, 0.4_8202, 0.4_7656, 0.4_6329, 0.4_8506]    )
       assert np.abs(image_slice - expected_slice    ).max() < 7e-3
