





'''simple docstring'''







import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import (
    DiffusionPipeline,
    UnCLIPImageVariationPipeline,
    UnCLIPScheduler,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, load_image, require_torch_gpu, skip_mps

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()
class 				__SCREAMING_SNAKE_CASE  (lowerCamelCase_       ,							unittest.TestCase      ):






				"""simple docstring"""







				__a            =UnCLIPImageVariationPipeline
				__a            =IMAGE_VARIATION_PARAMS - {'height', 'width', 'guidance_scale'}
				__a            =IMAGE_VARIATION_BATCH_PARAMS

				__a            =[
				    'generator',
				    'return_dict',
				    'decoder_num_inference_steps',
				    'super_res_num_inference_steps',
				]
				__a            =False




				@property
				def   UpperCamelCase__   (  self     :					Union[str, Any]					):
								return 32




				@property
				def   UpperCamelCase__   (  self     :					Dict					):
								return 32




				@property
				def   UpperCamelCase__   (  self     :					Union[str, Any]					):
								return self.time_input_dim




				@property
				def   UpperCamelCase__   (  self     :					Any					):
								return self.time_input_dim * 4




				@property
				def   UpperCamelCase__   (  self     :					Any					):
								return 1_00




				@property
				def   UpperCamelCase__   (  self     :					int					):
								_a								= CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"					)
								return tokenizer




				@property
				def   UpperCamelCase__   (  self     :					Dict					):
								torch.manual_seed(0					)
								_a								= CLIPTextConfig(
								    bos_token_id=0					,				eos_token_id=2					,				hidden_size=self.text_embedder_hidden_size					,				projection_dim=self.text_embedder_hidden_size					,				intermediate_size=37					,				layer_norm_eps=1e-0_5					,				num_attention_heads=4					,				num_hidden_layers=5					,				pad_token_id=1					,				vocab_size=10_00					,				)
								return CLIPTextModelWithProjection(__a					)




				@property
				def   UpperCamelCase__   (  self     :					List[str]					):
								torch.manual_seed(0					)
								_a								= CLIPVisionConfig(
								    hidden_size=self.text_embedder_hidden_size					,				projection_dim=self.text_embedder_hidden_size					,				num_hidden_layers=5					,				num_attention_heads=4					,				image_size=32					,				intermediate_size=37					,				patch_size=1					,				)
								return CLIPVisionModelWithProjection(__a					)




				@property
				def   UpperCamelCase__   (  self     :					List[str]					):
								torch.manual_seed(0					)

								_a								= {
								    "clip_embeddings_dim": self.text_embedder_hidden_size,
								    "time_embed_dim": self.time_embed_dim,
								    "cross_attention_dim": self.cross_attention_dim,
								}

								_a								= UnCLIPTextProjModel(**__a					)
								return model




				@property
				def   UpperCamelCase__   (  self     :					Optional[int]					):
								torch.manual_seed(0					)

								_a								= {
								    "sample_size": 32,
								    # RGB in channels
								    "in_channels": 3,
								    # Out channels is double in channels because predicts mean and variance
								    "out_channels": 6,
								    "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
								    "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
								    "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
								    "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
								    "layers_per_block": 1,
								    "cross_attention_dim": self.cross_attention_dim,
								    "attention_head_dim": 4,
								    "resnet_time_scale_shift": "scale_shift",
								    "class_embed_type": "identity",
								}

								_a								= UNetaDConditionModel(**__a					)
								return model




				@property
				def   UpperCamelCase__   (  self     :					int					):
								return {
								    "sample_size": 64,
								    "layers_per_block": 1,
								    "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
								    "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
								    "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
								    "in_channels": 6,
								    "out_channels": 3,
								}




				@property
				def   UpperCamelCase__   (  self     :					Optional[Any]					):
								torch.manual_seed(0					)

								_a								= UNetaDModel(**self.dummy_super_res_kwargs					)
								return model




				@property
				def   UpperCamelCase__   (  self     :					Union[str, Any]					):
								# seeded differently to get different unet than `self.dummy_super_res_first`
								torch.manual_seed(1					)

								_a								= UNetaDModel(**self.dummy_super_res_kwargs					)
								return model




				def   UpperCamelCase__   (  self     :					Any					):
								_a								= self.dummy_decoder
								_a								= self.dummy_text_proj
								_a								= self.dummy_text_encoder
								_a								= self.dummy_tokenizer
								_a								= self.dummy_super_res_first
								_a								= self.dummy_super_res_last

								_a								= UnCLIPScheduler(
								    variance_type="learned_range"					,				prediction_type="epsilon"					,				num_train_timesteps=10_00					,				)

								_a								= UnCLIPScheduler(
								    variance_type="fixed_small_log"					,				prediction_type="epsilon"					,				num_train_timesteps=10_00					,				)

								_a								= CLIPImageProcessor(crop_size=32					,				size=32					)

								_a								= self.dummy_image_encoder

								return {
								    "decoder": decoder,
								    "text_encoder": text_encoder,
								    "tokenizer": tokenizer,
								    "text_proj": text_proj,
								    "feature_extractor": feature_extractor,
								    "image_encoder": image_encoder,
								    "super_res_first": super_res_first,
								    "super_res_last": super_res_last,
								    "decoder_scheduler": decoder_scheduler,
								    "super_res_scheduler": super_res_scheduler,
								}




				def   UpperCamelCase__   (  self     :					Dict					,				__a     :					Any					,				__a     :					Optional[int]=0					,				__a     :					List[str]=True					):
								_a								= floats_tensor((1, 3, 32, 32)					,				rng=random.Random(__a					)					).to(__a					)
								if str(__a					).startswith("mps"					):
												_a								= torch.manual_seed(__a					)
								else:
												_a								= torch.Generator(device=__a					).manual_seed(__a					)

								if pil_image:
												_a								= input_image * 0.5 + 0.5
												_a								= input_image.clamp(0					,				1					)
												_a								= input_image.cpu().permute(0					,				2					,				3					,				1					).float().numpy()
												_a								= DiffusionPipeline.numpy_to_pil(__a					)[0]

								return {
								    "image": input_image,
								    "generator": generator,
								    "decoder_num_inference_steps": 2,
								    "super_res_num_inference_steps": 2,
								    "output_type": "np",
								}




				def   UpperCamelCase__   (  self     :					int					):
								_a								= "cpu"

								_a								= self.get_dummy_components()

								_a								= self.pipeline_class(**__a					)
								_a								= pipe.to(__a					)

								pipe.set_progress_bar_config(disable=__a					)

								_a								= self.get_dummy_inputs(__a					,				pil_image=__a					)

								_a								= pipe(**__a					)
								_a								= output.images

								_a								= self.get_dummy_inputs(__a					,				pil_image=__a					)

								_a								= pipe(
								    **__a					,				return_dict=__a					,				)[0]

								_a								= image[0, -3:, -3:, -1]
								_a								= image_from_tuple[0, -3:, -3:, -1]

								assert image.shape == (1, 64, 64, 3)

								_a								= np.array(
								    [
								        0.9997,
								        0.0002,
								        0.9997,
								        0.9997,
								        0.9969,
								        0.0023,
								        0.9997,
								        0.9969,
								        0.9970,
								    ]					)

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




				def   UpperCamelCase__   (  self     :					Dict					):
								_a								= "cpu"

								_a								= self.get_dummy_components()

								_a								= self.pipeline_class(**__a					)
								_a								= pipe.to(__a					)

								pipe.set_progress_bar_config(disable=__a					)

								_a								= self.get_dummy_inputs(__a					,				pil_image=__a					)

								_a								= pipe(**__a					)
								_a								= output.images

								_a								= self.get_dummy_inputs(__a					,				pil_image=__a					)

								_a								= pipe(
								    **__a					,				return_dict=__a					,				)[0]

								_a								= image[0, -3:, -3:, -1]
								_a								= image_from_tuple[0, -3:, -3:, -1]

								assert image.shape == (1, 64, 64, 3)

								_a								= np.array([0.9997, 0.0003, 0.9997, 0.9997, 0.9970, 0.0024, 0.9997, 0.9971, 0.9971]					)

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




				def   UpperCamelCase__   (  self     :					Optional[int]					):
								_a								= "cpu"

								_a								= self.get_dummy_components()

								_a								= self.pipeline_class(**__a					)
								_a								= pipe.to(__a					)

								pipe.set_progress_bar_config(disable=__a					)

								_a								= self.get_dummy_inputs(__a					,				pil_image=__a					)
								_a								= [
								    pipeline_inputs["image"],
								    pipeline_inputs["image"],
								]

								_a								= pipe(**__a					)
								_a								= output.images

								_a								= self.get_dummy_inputs(__a					,				pil_image=__a					)
								_a								= [
								    tuple_pipeline_inputs["image"],
								    tuple_pipeline_inputs["image"],
								]

								_a								= pipe(
								    **__a					,				return_dict=__a					,				)[0]

								_a								= image[0, -3:, -3:, -1]
								_a								= image_from_tuple[0, -3:, -3:, -1]

								assert image.shape == (2, 64, 64, 3)

								_a								= np.array(
								    [
								        0.9997,
								        0.9989,
								        0.0008,
								        0.0021,
								        0.9960,
								        0.0018,
								        0.0014,
								        0.0002,
								        0.9933,
								    ]					)

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




				def   UpperCamelCase__   (  self     :					Union[str, Any]					):
								_a								= torch.device("cpu"					)
								class 				__SCREAMING_SNAKE_CASE  :






												"""simple docstring"""







												__a            =1


								_a								= self.get_dummy_components()

								_a								= self.pipeline_class(**__a					)
								_a								= pipe.to(__a					)

								pipe.set_progress_bar_config(disable=__a					)

								_a								= torch.Generator(device=__a					).manual_seed(0					)
								_a								= pipe.decoder.dtype
								_a								= 1

								_a								= (
								    batch_size,
								    pipe.decoder.config.in_channels,
								    pipe.decoder.config.sample_size,
								    pipe.decoder.config.sample_size,
								)
								_a								= pipe.prepare_latents(
								    __a					,				dtype=__a					,				device=__a					,				generator=__a					,				latents=__a					,				scheduler=DummyScheduler()					)

								_a								= (
								    batch_size,
								    pipe.super_res_first.config.in_channels // 2,
								    pipe.super_res_first.config.sample_size,
								    pipe.super_res_first.config.sample_size,
								)
								_a								= pipe.prepare_latents(
								    __a					,				dtype=__a					,				device=__a					,				generator=__a					,				latents=__a					,				scheduler=DummyScheduler()					)

								_a								= self.get_dummy_inputs(__a					,				pil_image=__a					)

								_a								= pipe(
								    **__a					,				decoder_latents=__a					,				super_res_latents=__a					).images

								_a								= self.get_dummy_inputs(__a					,				pil_image=__a					)
								# Don't pass image, instead pass embedding
								_a								= pipeline_inputs.pop("image"					)
								_a								= pipe.image_encoder(__a					).image_embeds

								_a								= pipe(
								    **__a					,				decoder_latents=__a					,				super_res_latents=__a					,				image_embeddings=__a					,				).images

								# make sure passing text embeddings manually is identical
								assert np.abs(img_out_a - img_out_a					).max() < 1e-4




				@skip_mps
				def   UpperCamelCase__   (  self     :					List[str]					):
								_a								= torch_device == "cpu"

								# Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
								_a								= 1e-2

								self._test_attention_slicing_forward_pass(
								    test_max_difference=__a					,				expected_max_diff=__a					)




				@skip_mps
				def   UpperCamelCase__   (  self     :					Any					):
								_a								= torch_device == "cpu"
								_a								= True
								_a								= [
								    "decoder_num_inference_steps",
								    "super_res_num_inference_steps",
								]

								self._test_inference_batch_single_identical(
								    test_max_difference=__a					,				relax_max_difference=__a					,				additional_params_copy_to_batched_inputs=__a					,				)




				def   UpperCamelCase__   (  self     :					str					):
								_a								= [
								    "decoder_num_inference_steps",
								    "super_res_num_inference_steps",
								]

								if torch_device == "mps":
												# TODO: MPS errors with larger batch sizes
												_a								= [2, 3]
												self._test_inference_batch_consistent(
												    batch_sizes=__a					,				additional_params_copy_to_batched_inputs=__a					,				)
								else:
												self._test_inference_batch_consistent(
												    additional_params_copy_to_batched_inputs=__a					)




				@skip_mps
				def   UpperCamelCase__   (  self     :					Union[str, Any]					):
								return super().test_dict_tuple_outputs_equivalent()




				@skip_mps
				def   UpperCamelCase__   (  self     :					Tuple					):
								return super().test_save_load_local()




				@skip_mps
				def   UpperCamelCase__   (  self     :					Optional[int]					):
								return super().test_save_load_optional_components()


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






				"""simple docstring"""







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




				def   UpperCamelCase__   (  self     :					int					):
								_a								= load_image(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png"					)
								_a								= load_numpy(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
								    "/unclip/karlo_v1_alpha_cat_variation_fp16.npy"					)

								_a								= UnCLIPImageVariationPipeline.from_pretrained(
								    "kakaobrain/karlo-v1-alpha-image-variations"					,				torch_dtype=torch.floataa					)
								_a								= pipeline.to(__a					)
								pipeline.set_progress_bar_config(disable=__a					)

								_a								= torch.Generator(device="cpu"					).manual_seed(0					)
								_a								= pipeline(
								    __a					,				generator=__a					,				output_type="np"					,				)

								_a								= output.images[0]

								assert image.shape == (2_56, 2_56, 3)

								assert_mean_pixel_difference(__a					,				__a					,				15					)







'''simple docstring'''







from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


lowerCAmelCase_					:    int      	=			logging.get_logger(__name__)

lowerCAmelCase_					:    Tuple      	=			{
    'google/bigbird-roberta-base': 'https://huggingface.co/google/bigbird-roberta-base/resolve/main/config.json',
    'google/bigbird-roberta-large': 'https://huggingface.co/google/bigbird-roberta-large/resolve/main/config.json',
    'google/bigbird-base-trivia-itc': 'https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/config.json',
    # See all BigBird models at https://huggingface.co/models?filter=big_bird
}
class 				__SCREAMING_SNAKE_CASE  (lowerCamelCase_      ):






				"""simple docstring"""







				__a            ='big_bird'




				def __init__(  self     :					Optional[int]					,				__a     :					Dict=5_03_58					,				__a     :					str=7_68					,				__a     :					List[Any]=12					,				__a     :					List[str]=12					,				__a     :					Union[str, Any]=30_72					,				__a     :					str="gelu_new"					,				__a     :					Dict=0.1					,				__a     :					Union[str, Any]=0.1					,				__a     :					Any=40_96					,				__a     :					int=2					,				__a     :					Tuple=0.02					,				__a     :					List[Any]=1e-1_2					,				__a     :					int=True					,				__a     :					List[str]=0					,				__a     :					Tuple=1					,				__a     :					Optional[Any]=2					,				__a     :					Tuple=66					,				__a     :					str="block_sparse"					,				__a     :					Tuple=True					,				__a     :					Optional[int]=False					,				__a     :					str=64					,				__a     :					Tuple=3					,				__a     :					Any=None					,				**__a     :					Dict					,				):
								super().__init__(
								    pad_token_id=__a					,				bos_token_id=__a					,				eos_token_id=__a					,				sep_token_id=__a					,				**__a					,				)

								_a								= vocab_size
								_a								= max_position_embeddings
								_a								= hidden_size
								_a								= num_hidden_layers
								_a								= num_attention_heads
								_a								= intermediate_size
								_a								= hidden_act
								_a								= hidden_dropout_prob
								_a								= attention_probs_dropout_prob
								_a								= initializer_range
								_a								= type_vocab_size
								_a								= layer_norm_eps
								_a								= use_cache

								_a								= rescale_embeddings
								_a								= attention_type
								_a								= use_bias
								_a								= block_size
								_a								= num_random_blocks
								_a								= classifier_dropout


class 				__SCREAMING_SNAKE_CASE  (lowerCamelCase_      ):






				"""simple docstring"""







				@property
				def   UpperCamelCase__   (  self     :					Optional[int]					):
								if self.task == "multiple-choice":
												_a								= {0: "batch", 1: "choice", 2: "sequence"}
								else:
												_a								= {0: "batch", 1: "sequence"}
								return OrderedDict(
								    [
								        ("input_ids", dynamic_axis),
								        ("attention_mask", dynamic_axis),
								    ]					)
