



'''simple docstring'''



import math
from typing import Dict, Iterable, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    is_torch_available,
    is_torch_tensor,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_torch_available():
		import torch

if is_vision_available():
		import PIL


lowerCamelCase	 =       logging.get_logger(__name__)



def   _A       (   _lowerCAmelCase  ,   _lowerCAmelCase  ,   _lowerCAmelCase  ,   _lowerCAmelCase					):



							"""simple docstring"""



							def constraint_to_multiple_of(_lowerCAmelCase  ,   _lowerCAmelCase  ,   _lowerCAmelCase=0  ,   _lowerCAmelCase=None					):
														__lowercase       	=round(val / multiple					) * multiple

														if max_val is not None and x > max_val:
																					__lowercase       	=math.floor(val / multiple					) * multiple

														if x < min_val:
																					__lowercase       	=math.ceil(val / multiple					) * multiple

														return x

							__lowercase       	=(output_size, output_size) if isinstance(_lowerCAmelCase  ,   _lowerCAmelCase					) else output_size

							__lowercase	,				__lowercase       	=get_image_size(_lowerCAmelCase					)
							__lowercase	,				__lowercase       	=output_size

							# determine new height and width
							__lowercase       	=output_height / input_height
							__lowercase       	=output_width / input_width

							if keep_aspect_ratio:
														# scale as little as possible
														if abs(1 - scale_width					) < abs(1 - scale_height					):
																					# fit width
																					__lowercase       	=scale_width
														else:
																					# fit height
																					__lowercase       	=scale_height

							__lowercase       	=constraint_to_multiple_of(scale_height * input_height  ,   multiple=_lowerCAmelCase					)
							__lowercase       	=constraint_to_multiple_of(scale_width * input_width  ,   multiple=_lowerCAmelCase					)

							return (new_height, new_width)





class _UpperCamelCase		(     A		):





				'''simple docstring'''





				lowerCAmelCase__					    =    ["""pixel_values"""]

				def __init__(							self  :  Any  ,       _lowerCAmelCase  :  bool = True  ,       _lowerCAmelCase  :  Dict[str, int] = None  ,       _lowerCAmelCase  :  PILImageResampling = PILImageResampling.BILINEAR  ,       _lowerCAmelCase  :  bool = False  ,       _lowerCAmelCase  :  int = 1  ,       _lowerCAmelCase  :  bool = True  ,       _lowerCAmelCase  :  Union[int, float] = 1 / 2_5_5  ,       _lowerCAmelCase  :  bool = True  ,       _lowerCAmelCase  :  Optional[Union[float, List[float]]] = None  ,       _lowerCAmelCase  :  Optional[Union[float, List[float]]] = None  ,       **_lowerCAmelCase  :  List[Any]  ,       ):







											'''simple docstring'''

											super().__init__(**_lowerCAmelCase)
											__lowercase       	=size if size is not None else {'height': 3_8_4, 'width': 3_8_4}
											__lowercase       	=get_size_dict(_lowerCAmelCase)
											__lowercase       	=do_resize
											__lowercase       	=size
											__lowercase       	=keep_aspect_ratio
											__lowercase       	=ensure_multiple_of
											__lowercase       	=resample
											__lowercase       	=do_rescale
											__lowercase       	=rescale_factor
											__lowercase       	=do_normalize
											__lowercase       	=image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
											__lowercase       	=image_std if image_std is not None else IMAGENET_STANDARD_STD

				def 				__lowerCamelCase     (							self  :  Any  ,       _lowerCAmelCase  :  np.ndarray  ,       _lowerCAmelCase  :  Dict[str, int]  ,       _lowerCAmelCase  :  bool = False  ,       _lowerCAmelCase  :  int = 1  ,       _lowerCAmelCase  :  PILImageResampling = PILImageResampling.BICUBIC  ,       _lowerCAmelCase  :  Optional[Union[str, ChannelDimension]] = None  ,       **_lowerCAmelCase  :  Tuple  ,       ):







											'''simple docstring'''

											__lowercase       	=get_size_dict(_lowerCAmelCase)
											if "height" not in size or "width" not in size:
																		raise ValueError(f"""The size dictionary must contain the keys 'height' and 'width'. Got {size.keys()}""")
											__lowercase       	=get_resize_output_image_size(
											    _lowerCAmelCase  ,       output_size=(size['height'], size['width'])  ,       keep_aspect_ratio=_lowerCAmelCase  ,       multiple=_lowerCAmelCase  ,       )
											return resize(_lowerCAmelCase  ,       size=_lowerCAmelCase  ,       resample=_lowerCAmelCase  ,       data_format=_lowerCAmelCase  ,       **_lowerCAmelCase)

				def 				__lowerCamelCase     (							self  :  Union[str, Any]  ,       _lowerCAmelCase  :  np.ndarray  ,       _lowerCAmelCase  :  Union[int, float]  ,       _lowerCAmelCase  :  Optional[Union[str, ChannelDimension]] = None  ,       **_lowerCAmelCase  :  Union[str, Any]  ,       ):







											'''simple docstring'''

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

				def 				__lowerCamelCase     (							self  :  List[Any]  ,       _lowerCAmelCase  :  np.ndarray  ,       _lowerCAmelCase  :  Union[float, List[float]]  ,       _lowerCAmelCase  :  Union[float, List[float]]  ,       _lowerCAmelCase  :  Optional[Union[str, ChannelDimension]] = None  ,       **_lowerCAmelCase  :  Union[str, Any]  ,       ):







											'''simple docstring'''

											return normalize(_lowerCAmelCase  ,       mean=_lowerCAmelCase  ,       std=_lowerCAmelCase  ,       data_format=_lowerCAmelCase  ,       **_lowerCAmelCase)

				def 				__lowerCamelCase     (							self  :  List[Any]  ,       _lowerCAmelCase  :  ImageInput  ,       _lowerCAmelCase  :  bool = None  ,       _lowerCAmelCase  :  int = None  ,       _lowerCAmelCase  :  bool = None  ,       _lowerCAmelCase  :  int = None  ,       _lowerCAmelCase  :  PILImageResampling = None  ,       _lowerCAmelCase  :  bool = None  ,       _lowerCAmelCase  :  float = None  ,       _lowerCAmelCase  :  bool = None  ,       _lowerCAmelCase  :  Optional[Union[float, List[float]]] = None  ,       _lowerCAmelCase  :  Optional[Union[float, List[float]]] = None  ,       _lowerCAmelCase  :  Optional[Union[str, TensorType]] = None  ,       _lowerCAmelCase  :  ChannelDimension = ChannelDimension.FIRST  ,       **_lowerCAmelCase  :  Tuple  ,       ):







											'''simple docstring'''

											__lowercase       	=do_resize if do_resize is not None else self.do_resize
											__lowercase       	=size if size is not None else self.size
											__lowercase       	=get_size_dict(_lowerCAmelCase)
											__lowercase       	=keep_aspect_ratio if keep_aspect_ratio is not None else self.keep_aspect_ratio
											__lowercase       	=ensure_multiple_of if ensure_multiple_of is not None else self.ensure_multiple_of
											__lowercase       	=resample if resample is not None else self.resample
											__lowercase       	=do_rescale if do_rescale is not None else self.do_rescale
											__lowercase       	=rescale_factor if rescale_factor is not None else self.rescale_factor
											__lowercase       	=do_normalize if do_normalize is not None else self.do_normalize
											__lowercase       	=image_mean if image_mean is not None else self.image_mean
											__lowercase       	=image_std if image_std is not None else self.image_std

											__lowercase       	=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 or resample is None:
																		raise ValueError('Size and resample must be specified if do_resize 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.
											__lowercase       	=[to_numpy_array(_lowerCAmelCase) for image in images]

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

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

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

											__lowercase       	=[to_channel_dimension_format(_lowerCAmelCase  ,       _lowerCAmelCase) for image in images]

											__lowercase       	={'pixel_values': images}
											return BatchFeature(data=_lowerCAmelCase  ,       tensor_type=_lowerCAmelCase)

				def 				__lowerCamelCase     (							self  :  int  ,       _lowerCAmelCase  :  Tuple  ,       _lowerCAmelCase  :  List[Tuple] = None):







											'''simple docstring'''

											__lowercase       	=outputs.logits

											# Resize logits and compute semantic segmentation maps
											if target_sizes is not None:
																		if len(_lowerCAmelCase) != len(_lowerCAmelCase):
																									raise ValueError(
																									    'Make sure that you pass in as many target sizes as the batch dimension of the logits')

																		if is_torch_tensor(_lowerCAmelCase):
																									__lowercase       	=target_sizes.numpy()

																		__lowercase       	=[]

																		for idx in range(len(_lowerCAmelCase)):
																									__lowercase       	=torch.nn.functional.interpolate(
																									    logits[idx].unsqueeze(dim=0)  ,       size=target_sizes[idx]  ,       mode='bilinear'  ,       align_corners=_lowerCAmelCase)
																									__lowercase       	=resized_logits[0].argmax(dim=0)
																									semantic_segmentation.append(_lowerCAmelCase)
											else:
																		__lowercase       	=logits.argmax(dim=1)
																		__lowercase       	=[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]

											return semantic_segmentation








'''simple docstring'''



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 _UpperCamelCase		(     A      ,       A      ,       A      ,       unittest.TestCase		):





				'''simple docstring'''





				lowerCAmelCase__					    =    AltDiffusionPipeline
				lowerCAmelCase__					    =    TEXT_TO_IMAGE_PARAMS
				lowerCAmelCase__					    =    TEXT_TO_IMAGE_BATCH_PARAMS
				lowerCAmelCase__					    =    TEXT_TO_IMAGE_IMAGE_PARAMS
				lowerCAmelCase__					    =    TEXT_TO_IMAGE_IMAGE_PARAMS

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







											'''simple docstring'''

											torch.manual_seed(0)
											__lowercase       	=UNetaDConditionModel(
											    block_out_channels=(3_2, 6_4)  ,       layers_per_block=2  ,       sample_size=3_2  ,       in_channels=4  ,       out_channels=4  ,       down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')  ,       up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')  ,       cross_attention_dim=3_2  ,       )
											__lowercase       	=DDIMScheduler(
											    beta_start=0.0_0085  ,       beta_end=0.012  ,       beta_schedule='scaled_linear'  ,       clip_sample=_lowerCAmelCase  ,       set_alpha_to_one=_lowerCAmelCase  ,       )
											torch.manual_seed(0)
											__lowercase       	=AutoencoderKL(
											    block_out_channels=[3_2, 6_4]  ,       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)
											__lowercase       	=CLIPTextConfig(
											    bos_token_id=0  ,       eos_token_id=2  ,       hidden_size=3_2  ,       projection_dim=3_2  ,       intermediate_size=3_7  ,       layer_norm_eps=1e-05  ,       num_attention_heads=4  ,       num_hidden_layers=5  ,       pad_token_id=1  ,       vocab_size=5_0_0_2  ,       )
											__lowercase       	=CLIPTextModel(_lowerCAmelCase)

											__lowercase       	=XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta')
											__lowercase       	=7_7

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

				def 				__lowerCamelCase     (							self  :  List[Any]  ,       _lowerCAmelCase  :  Union[str, Any]  ,       _lowerCAmelCase  :  Dict=0):







											'''simple docstring'''

											if str(_lowerCAmelCase).startswith('mps'):
																		__lowercase       	=torch.manual_seed(_lowerCAmelCase)
											else:
																		__lowercase       	=torch.Generator(device=_lowerCAmelCase).manual_seed(_lowerCAmelCase)
											__lowercase       	={
											    '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 				__lowerCamelCase     (							self  :  List[str]):







											'''simple docstring'''

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

				def 				__lowerCamelCase     (							self  :  Optional[int]):







											'''simple docstring'''

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

				def 				__lowerCamelCase     (							self  :  Optional[Any]):







											'''simple docstring'''

											__lowercase       	='cpu'  # ensure determinism for the device-dependent torch.Generator

											__lowercase       	=self.get_dummy_components()
											torch.manual_seed(0)
											__lowercase       	=RobertaSeriesConfig(
											    hidden_size=3_2  ,       project_dim=3_2  ,       intermediate_size=3_7  ,       layer_norm_eps=1e-05  ,       num_attention_heads=4  ,       num_hidden_layers=5  ,       vocab_size=5_0_0_2  ,       )
											# TODO: remove after fixing the non-deterministic text encoder
											__lowercase       	=RobertaSeriesModelWithTransformation(_lowerCAmelCase)
											__lowercase       	=text_encoder

											__lowercase       	=AltDiffusionPipeline(**_lowerCAmelCase)
											__lowercase       	=alt_pipe.to(_lowerCAmelCase)
											alt_pipe.set_progress_bar_config(disable=_lowerCAmelCase)

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

											assert image.shape == (1, 6_4, 6_4, 3)
											__lowercase       	=np.array(
											    [0.574_8162, 0.6044_7145, 0.4882_1217, 0.5010_0636, 0.543_1185, 0.4576_3683, 0.4965_7696, 0.4813_2733, 0.4757_3093])

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

				def 				__lowerCamelCase     (							self  :  str):







											'''simple docstring'''

											__lowercase       	='cpu'  # ensure determinism for the device-dependent torch.Generator

											__lowercase       	=self.get_dummy_components()
											__lowercase       	=PNDMScheduler(skip_prk_steps=_lowerCAmelCase)
											torch.manual_seed(0)
											__lowercase       	=RobertaSeriesConfig(
											    hidden_size=3_2  ,       project_dim=3_2  ,       intermediate_size=3_7  ,       layer_norm_eps=1e-05  ,       num_attention_heads=4  ,       num_hidden_layers=5  ,       vocab_size=5_0_0_2  ,       )
											# TODO: remove after fixing the non-deterministic text encoder
											__lowercase       	=RobertaSeriesModelWithTransformation(_lowerCAmelCase)
											__lowercase       	=text_encoder
											__lowercase       	=AltDiffusionPipeline(**_lowerCAmelCase)
											__lowercase       	=alt_pipe.to(_lowerCAmelCase)
											alt_pipe.set_progress_bar_config(disable=_lowerCAmelCase)

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

											assert image.shape == (1, 6_4, 6_4, 3)
											__lowercase       	=np.array(
											    [0.5160_5093, 0.570_7241, 0.4736_5507, 0.5057_8886, 0.563_3877, 0.464_2503, 0.518_2081, 0.4876_3484, 0.4908_4237])

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





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





				'''simple docstring'''





				def 				__lowerCamelCase     (							self  :  str):







											'''simple docstring'''

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

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







											'''simple docstring'''

											__lowercase       	=AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'  ,       safety_checker=_lowerCAmelCase)
											__lowercase       	=alt_pipe.to(_lowerCAmelCase)
											alt_pipe.set_progress_bar_config(disable=_lowerCAmelCase)

											__lowercase       	='A painting of a squirrel eating a burger'
											__lowercase       	=torch.manual_seed(0)
											__lowercase       	=alt_pipe([prompt]  ,       generator=_lowerCAmelCase  ,       guidance_scale=6.0  ,       num_inference_steps=2_0  ,       output_type='np')

											__lowercase       	=output.images

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

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											__lowercase       	=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 				__lowerCamelCase     (							self  :  Dict):







											'''simple docstring'''

											__lowercase       	=DDIMScheduler.from_pretrained('BAAI/AltDiffusion'  ,       subfolder='scheduler')

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

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

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

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

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											__lowercase       	=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



