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


UpperCAmelCase_								=							logging.get_logger(__name__)


if is_vision_available():
						import PIL





class      lowerCAmelCase						(					_a			):
					_SCREAMING_SNAKE_CASE     :     str		=["""pixel_values"""]



					def __init__(					self				,      lowerCAmelCase__ = True				,      lowerCAmelCase__ = None				,      lowerCAmelCase__ = PILImageResampling.BICUBIC				,      lowerCAmelCase__ = True				,      lowerCAmelCase__ = None				,      lowerCAmelCase__ = True				,      lowerCAmelCase__ = 1 / 255				,      lowerCAmelCase__ = True				,      lowerCAmelCase__ = None				,      lowerCAmelCase__ = None				,      lowerCAmelCase__ = True				,      **lowerCAmelCase__				,      ):
												super().__init__(**lowerCAmelCase__       )
												_A=	size if size is not None else {'shortest_edge': 224}
												_A=	get_size_dict(lowerCAmelCase__				,      default_to_square=lowerCAmelCase__       )
												_A=	crop_size if crop_size is not None else {'height': 224, 'width': 224}
												_A=	get_size_dict(lowerCAmelCase__				,      default_to_square=lowerCAmelCase__				,      param_name='crop_size'       )

												_A=	do_resize
												_A=	size
												_A=	resample
												_A=	do_center_crop
												_A=	crop_size
												_A=	do_rescale
												_A=	rescale_factor
												_A=	do_normalize
												_A=	image_mean if image_mean is not None else OPENAI_CLIP_MEAN
												_A=	image_std if image_std is not None else OPENAI_CLIP_STD
												_A=	do_convert_rgb



					def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__ = PILImageResampling.BICUBIC				,      lowerCAmelCase__ = None				,      **lowerCAmelCase__				,      ):
												_A=	get_size_dict(lowerCAmelCase__				,      default_to_square=lowerCAmelCase__       )
												if "shortest_edge" not in size:
																			raise ValueError(f"The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}"       )
												_A=	get_resize_output_image_size(lowerCAmelCase__				,      size=size['shortest_edge']				,      default_to_square=lowerCAmelCase__       )
												return resize(lowerCAmelCase__				,      size=lowerCAmelCase__				,      resample=lowerCAmelCase__				,      data_format=lowerCAmelCase__				,      **lowerCAmelCase__       )



					def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__ = None				,      **lowerCAmelCase__				,      ):
												_A=	get_size_dict(lowerCAmelCase__       )
												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(lowerCAmelCase__				,      size=(size['height'], size['width'])				,      data_format=lowerCAmelCase__				,      **lowerCAmelCase__       )



					def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__ = None				,      **lowerCAmelCase__				,      ):
												return rescale(lowerCAmelCase__				,      scale=lowerCAmelCase__				,      data_format=lowerCAmelCase__				,      **lowerCAmelCase__       )



					def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__				,      lowerCAmelCase__ = None				,      **lowerCAmelCase__				,      ):
												return normalize(lowerCAmelCase__				,      mean=lowerCAmelCase__				,      std=lowerCAmelCase__				,      data_format=lowerCAmelCase__				,      **lowerCAmelCase__       )





					def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__ = None				,      lowerCAmelCase__ = None				,      lowerCAmelCase__ = None				,      lowerCAmelCase__ = None				,      lowerCAmelCase__ = None				,      lowerCAmelCase__ = None				,      lowerCAmelCase__ = None				,      lowerCAmelCase__ = None				,      lowerCAmelCase__ = None				,      lowerCAmelCase__ = None				,      lowerCAmelCase__ = None				,      lowerCAmelCase__ = None				,      lowerCAmelCase__ = ChannelDimension.FIRST				,      **lowerCAmelCase__				,      ):
												_A=	do_resize if do_resize is not None else self.do_resize
												_A=	size if size is not None else self.size
												_A=	get_size_dict(lowerCAmelCase__				,      param_name='size'				,      default_to_square=lowerCAmelCase__       )
												_A=	resample if resample is not None else self.resample
												_A=	do_center_crop if do_center_crop is not None else self.do_center_crop
												_A=	crop_size if crop_size is not None else self.crop_size
												_A=	get_size_dict(lowerCAmelCase__				,      param_name='crop_size'				,      default_to_square=lowerCAmelCase__       )
												_A=	do_rescale if do_rescale is not None else self.do_rescale
												_A=	rescale_factor if rescale_factor is not None else self.rescale_factor
												_A=	do_normalize if do_normalize is not None else self.do_normalize
												_A=	image_mean if image_mean is not None else self.image_mean
												_A=	image_std if image_std is not None else self.image_std
												_A=	do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

												_A=	make_list_of_images(lowerCAmelCase__       )

												if not valid_images(lowerCAmelCase__       ):
																			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:
																			_A=	[convert_to_rgb(lowerCAmelCase__       ) for image in images]

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

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

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

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

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

												_A=	[to_channel_dimension_format(lowerCAmelCase__				,      lowerCAmelCase__       ) for image in images]

												_A=	{'pixel_values': images}
												return BatchFeature(data=lowerCAmelCase__				,      tensor_type=lowerCAmelCase__       )
import gc
import unittest

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

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
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 PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





class      lowerCAmelCase						(					_a			,		_a			,		_a			,		unittest.TestCase			):
     _SCREAMING_SNAKE_CASE     :     Optional[int]		=AltDiffusionPipeline
     _SCREAMING_SNAKE_CASE     :     int		=TEXT_TO_IMAGE_PARAMS
     _SCREAMING_SNAKE_CASE     :     List[Any]		=TEXT_TO_IMAGE_BATCH_PARAMS
     _SCREAMING_SNAKE_CASE     :     Dict		=TEXT_TO_IMAGE_IMAGE_PARAMS
     _SCREAMING_SNAKE_CASE     :     Any		=TEXT_TO_IMAGE_IMAGE_PARAMS



     def 			a__  (					self       ):
            torch.manual_seed(0       )
            _A=	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=	DDIMScheduler(
                beta_start=0.00085				,      beta_end=0.012				,      beta_schedule='scaled_linear'				,      clip_sample=lowerCAmelCase__				,      set_alpha_to_one=lowerCAmelCase__				,      )
            torch.manual_seed(0       )
            _A=	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				,      )

            # TODO: address the non-deterministic text encoder (fails for save-load tests)
            # torch.manual_seed(0)
            # text_encoder_config = RobertaSeriesConfig(
            #     hidden_size=32,
            #     project_dim=32,
            #     intermediate_size=37,
            #     layer_norm_eps=1e-05,
            #     num_attention_heads=4,
            #     num_hidden_layers=5,
            #     vocab_size=5002,
            # )
            # text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

            torch.manual_seed(0       )
            _A=	CLIPTextConfig(
                bos_token_id=0				,      eos_token_id=2				,      hidden_size=32				,      projection_dim=32				,      intermediate_size=37				,      layer_norm_eps=1E-05				,      num_attention_heads=4				,      num_hidden_layers=5				,      pad_token_id=1				,      vocab_size=5002				,      )
            _A=	CLIPTextModel(lowerCAmelCase__       )

            _A=	XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'       )
            _A=	77

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



     def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__=0       ):
            if str(lowerCAmelCase__       ).startswith('mps'       ):
                   _A=	torch.manual_seed(lowerCAmelCase__       )
            else:
                   _A=	torch.Generator(device=lowerCAmelCase__       ).manual_seed(lowerCAmelCase__       )
            _A=	{
                'prompt': 'A painting of a squirrel eating a burger',
                'generator': generator,
                'num_inference_steps': 2,
                'guidance_scale': 6.0,
                'output_type': 'numpy',
            }
            return inputs



     def 			a__  (					self       ):
            super().test_attention_slicing_forward_pass(expected_max_diff=3E-3       )



     def 			a__  (					self       ):
            super().test_inference_batch_single_identical(expected_max_diff=3E-3       )



     def 			a__  (					self       ):
            _A=	'cpu'  # ensure determinism for the device-dependent torch.Generator

            _A=	self.get_dummy_components()
            torch.manual_seed(0       )
            _A=	RobertaSeriesConfig(
                hidden_size=32				,      project_dim=32				,      intermediate_size=37				,      layer_norm_eps=1E-05				,      num_attention_heads=4				,      num_hidden_layers=5				,      vocab_size=5002				,      )
            # TODO: remove after fixing the non-deterministic text encoder
            _A=	RobertaSeriesModelWithTransformation(lowerCAmelCase__       )
            _A=	text_encoder

            _A=	AltDiffusionPipeline(**lowerCAmelCase__       )
            _A=	alt_pipe.to(lowerCAmelCase__       )
            alt_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	self.get_dummy_inputs(lowerCAmelCase__       )
            _A=	'A photo of an astronaut'
            _A=	alt_pipe(**lowerCAmelCase__       )
            _A=	output.images
            _A=	image[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)
            _A=	np.array(
                [0.5748162, 0.60447145, 0.48821217, 0.50100636, 0.5431185, 0.45763683, 0.49657696, 0.48132733, 0.47573093]       )

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





     def 			a__  (					self       ):
            _A=	'cpu'  # ensure determinism for the device-dependent torch.Generator

            _A=	self.get_dummy_components()
            _A=	PNDMScheduler(skip_prk_steps=lowerCAmelCase__       )
            torch.manual_seed(0       )
            _A=	RobertaSeriesConfig(
                hidden_size=32				,      project_dim=32				,      intermediate_size=37				,      layer_norm_eps=1E-05				,      num_attention_heads=4				,      num_hidden_layers=5				,      vocab_size=5002				,      )
            # TODO: remove after fixing the non-deterministic text encoder
            _A=	RobertaSeriesModelWithTransformation(lowerCAmelCase__       )
            _A=	text_encoder
            _A=	AltDiffusionPipeline(**lowerCAmelCase__       )
            _A=	alt_pipe.to(lowerCAmelCase__       )
            alt_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	self.get_dummy_inputs(lowerCAmelCase__       )
            _A=	alt_pipe(**lowerCAmelCase__       )
            _A=	output.images
            _A=	image[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)
            _A=	np.array(
                [0.51605093, 0.5707241, 0.47365507, 0.50578886, 0.5633877, 0.4642503, 0.5182081, 0.48763484, 0.49084237]       )

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





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



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



     def 			a__  (					self       ):
            # make sure here that pndm scheduler skips prk
            _A=	AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'				,      safety_checker=lowerCAmelCase__       )
            _A=	alt_pipe.to(lowerCAmelCase__       )
            alt_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	'A painting of a squirrel eating a burger'
            _A=	torch.manual_seed(0       )
            _A=	alt_pipe([prompt]				,      generator=lowerCAmelCase__				,      guidance_scale=6.0				,      num_inference_steps=20				,      output_type='np'       )

            _A=	output.images

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

            assert image.shape == (1, 512, 512, 3)
            _A=	np.array([0.1010, 0.0800, 0.0794, 0.0885, 0.0843, 0.0762, 0.0769, 0.0729, 0.0586]       )

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





     def 			a__  (					self       ):
            _A=	DDIMScheduler.from_pretrained('BAAI/AltDiffusion'				,      subfolder='scheduler'       )

            _A=	AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'				,      scheduler=lowerCAmelCase__				,      safety_checker=lowerCAmelCase__       )
            _A=	alt_pipe.to(lowerCAmelCase__       )
            alt_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	'A painting of a squirrel eating a burger'
            _A=	torch.manual_seed(0       )

            _A=	alt_pipe([prompt]				,      generator=lowerCAmelCase__				,      num_inference_steps=2				,      output_type='numpy'       )
            _A=	output.images

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

            assert image.shape == (1, 512, 512, 3)
            _A=	np.array([0.4019, 0.4052, 0.3810, 0.4119, 0.3916, 0.3982, 0.4651, 0.4195, 0.5323]       )

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