
import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()




class 				_lowerCamelCase (      unittest.TestCase ):





     """simple docstring"""







     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Dict:




           """simple docstring"""







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






     @property
     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Optional[Any]:




           """simple docstring"""







           UpperCamelCase__    :   Optional[Any]					=				1
           UpperCamelCase__    :   Union[str, Any]					=				3
           UpperCamelCase__    :   List[Any]					=				(3_2, 3_2)

           UpperCamelCase__    :   Dict					=				floats_tensor((batch_size, num_channels) + sizes						,						rng=random.Random(0  )  ).to(lowercase__  )
           return image






     @property
     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Tuple:




           """simple docstring"""







           torch.manual_seed(0  )
           UpperCamelCase__    :   List[str]					=				UNetaDConditionModel(
               block_out_channels=(3_2, 3_2, 6_4)						,						layers_per_block=2						,						sample_size=3_2						,						in_channels=7						,						out_channels=4						,						down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''', '''CrossAttnDownBlock2D''')						,						up_block_types=('''CrossAttnUpBlock2D''', '''CrossAttnUpBlock2D''', '''UpBlock2D''')						,						cross_attention_dim=3_2						,						attention_head_dim=8						,						use_linear_projection=lowercase__						,						only_cross_attention=(True, True, False)						,						num_class_embeds=1_0_0						,						)
           return model






     @property
     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   List[str]:




           """simple docstring"""







           torch.manual_seed(0  )
           UpperCamelCase__    :   Any					=				AutoencoderKL(
               block_out_channels=[3_2, 3_2, 6_4]						,						in_channels=3						,						out_channels=3						,						down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''', '''DownEncoderBlock2D''']						,						up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''', '''UpDecoderBlock2D''']						,						latent_channels=4						,						)
           return model






     @property
     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   List[str]:




           """simple docstring"""







           torch.manual_seed(0  )
           UpperCamelCase__    :   Union[str, Any]					=				CLIPTextConfig(
               bos_token_id=0						,						eos_token_id=2						,						hidden_size=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=1_0_0_0						,						hidden_act='''gelu'''						,						projection_dim=5_1_2						,						)
           return CLIPTextModel(lowercase__  )






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Any:




           """simple docstring"""







           UpperCamelCase__    :   List[str]					=				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
           UpperCamelCase__    :   Union[str, Any]					=				self.dummy_cond_unet_upscale
           UpperCamelCase__    :   List[Any]					=				DDPMScheduler()
           UpperCamelCase__    :   Dict					=				DDIMScheduler(prediction_type='''v_prediction'''  )
           UpperCamelCase__    :   Union[str, Any]					=				self.dummy_vae
           UpperCamelCase__    :   Optional[Any]					=				self.dummy_text_encoder
           UpperCamelCase__    :   List[Any]					=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )

           UpperCamelCase__    :   int					=				self.dummy_image.cpu().permute(0						,						2						,						3						,						1  )[0]
           UpperCamelCase__    :   List[str]					=				Image.fromarray(np.uinta(lowercase__  )  ).convert('''RGB'''  ).resize((6_4, 6_4)  )

           # make sure here that pndm scheduler skips prk
           UpperCamelCase__    :   int					=				StableDiffusionUpscalePipeline(
               unet=lowercase__						,						low_res_scheduler=lowercase__						,						scheduler=lowercase__						,						vae=lowercase__						,						text_encoder=lowercase__						,						tokenizer=lowercase__						,						max_noise_level=3_5_0						,						)
           UpperCamelCase__    :   Tuple					=				sd_pipe.to(lowercase__  )
           sd_pipe.set_progress_bar_config(disable=lowercase__  )

           UpperCamelCase__    :   Union[str, Any]					=				'''A painting of a squirrel eating a burger'''
           UpperCamelCase__    :   Any					=				torch.Generator(device=lowercase__  ).manual_seed(0  )
           UpperCamelCase__    :   int					=				sd_pipe(
               [prompt]						,						image=lowercase__						,						generator=lowercase__						,						guidance_scale=6.0						,						noise_level=2_0						,						num_inference_steps=2						,						output_type='''np'''						,						)

           UpperCamelCase__    :   Dict					=				output.images

           UpperCamelCase__    :   Optional[int]					=				torch.Generator(device=lowercase__  ).manual_seed(0  )
           UpperCamelCase__    :   List[str]					=				sd_pipe(
               [prompt]						,						image=lowercase__						,						generator=lowercase__						,						guidance_scale=6.0						,						noise_level=2_0						,						num_inference_steps=2						,						output_type='''np'''						,						return_dict=lowercase__						,						)[0]

           UpperCamelCase__    :   Optional[Any]					=				image[0, -3:, -3:, -1]
           UpperCamelCase__    :   Union[str, Any]					=				image_from_tuple[0, -3:, -3:, -1]

           UpperCamelCase__    :   Dict					=				low_res_image.size[0] * 4
           assert image.shape == (1, expected_height_width, expected_height_width, 3)
           UpperCamelCase__    :   int					=				np.array([0.3113, 0.3910, 0.4272, 0.4859, 0.5061, 0.4652, 0.5362, 0.5715, 0.5661]  )

           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    __SCREAMING_SNAKE_CASE			(							self  )  ->   str:




           """simple docstring"""







           UpperCamelCase__    :   Any					=				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
           UpperCamelCase__    :   Dict					=				self.dummy_cond_unet_upscale
           UpperCamelCase__    :   Any					=				DDPMScheduler()
           UpperCamelCase__    :   List[str]					=				DDIMScheduler(prediction_type='''v_prediction'''  )
           UpperCamelCase__    :   Optional[Any]					=				self.dummy_vae
           UpperCamelCase__    :   List[str]					=				self.dummy_text_encoder
           UpperCamelCase__    :   Optional[Any]					=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )

           UpperCamelCase__    :   str					=				self.dummy_image.cpu().permute(0						,						2						,						3						,						1  )[0]
           UpperCamelCase__    :   Optional[Any]					=				Image.fromarray(np.uinta(lowercase__  )  ).convert('''RGB'''  ).resize((6_4, 6_4)  )

           # make sure here that pndm scheduler skips prk
           UpperCamelCase__    :   str					=				StableDiffusionUpscalePipeline(
               unet=lowercase__						,						low_res_scheduler=lowercase__						,						scheduler=lowercase__						,						vae=lowercase__						,						text_encoder=lowercase__						,						tokenizer=lowercase__						,						max_noise_level=3_5_0						,						)
           UpperCamelCase__    :   List[Any]					=				sd_pipe.to(lowercase__  )
           sd_pipe.set_progress_bar_config(disable=lowercase__  )

           UpperCamelCase__    :   str					=				'''A painting of a squirrel eating a burger'''
           UpperCamelCase__    :   Optional[Any]					=				sd_pipe(
               2 * [prompt]						,						image=2 * [low_res_image]						,						guidance_scale=6.0						,						noise_level=2_0						,						num_inference_steps=2						,						output_type='''np'''						,						)
           UpperCamelCase__    :   Any					=				output.images
           assert image.shape[0] == 2

           UpperCamelCase__    :   str					=				torch.Generator(device=lowercase__  ).manual_seed(0  )
           UpperCamelCase__    :   Optional[Any]					=				sd_pipe(
               [prompt]						,						image=lowercase__						,						generator=lowercase__						,						num_images_per_prompt=2						,						guidance_scale=6.0						,						noise_level=2_0						,						num_inference_steps=2						,						output_type='''np'''						,						)
           UpperCamelCase__    :   str					=				output.images
           assert image.shape[0] == 2






     @unittest.skipIf(torch_device != '''cuda'''						,						'''This test requires a GPU'''  )
     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   List[str]:




           """simple docstring"""







           UpperCamelCase__    :   str					=				self.dummy_cond_unet_upscale
           UpperCamelCase__    :   Dict					=				DDPMScheduler()
           UpperCamelCase__    :   int					=				DDIMScheduler(prediction_type='''v_prediction'''  )
           UpperCamelCase__    :   Dict					=				self.dummy_vae
           UpperCamelCase__    :   Optional[int]					=				self.dummy_text_encoder
           UpperCamelCase__    :   List[Any]					=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )

           UpperCamelCase__    :   Dict					=				self.dummy_image.cpu().permute(0						,						2						,						3						,						1  )[0]
           UpperCamelCase__    :   Dict					=				Image.fromarray(np.uinta(lowercase__  )  ).convert('''RGB'''  ).resize((6_4, 6_4)  )

           # put models in fp16, except vae as it overflows in fp16
           UpperCamelCase__    :   Tuple					=				unet.half()
           UpperCamelCase__    :   int					=				text_encoder.half()

           # make sure here that pndm scheduler skips prk
           UpperCamelCase__    :   Optional[Any]					=				StableDiffusionUpscalePipeline(
               unet=lowercase__						,						low_res_scheduler=lowercase__						,						scheduler=lowercase__						,						vae=lowercase__						,						text_encoder=lowercase__						,						tokenizer=lowercase__						,						max_noise_level=3_5_0						,						)
           UpperCamelCase__    :   List[Any]					=				sd_pipe.to(lowercase__  )
           sd_pipe.set_progress_bar_config(disable=lowercase__  )

           UpperCamelCase__    :   List[Any]					=				'''A painting of a squirrel eating a burger'''
           UpperCamelCase__    :   int					=				torch.manual_seed(0  )
           UpperCamelCase__    :   Tuple					=				sd_pipe(
               [prompt]						,						image=lowercase__						,						generator=lowercase__						,						num_inference_steps=2						,						output_type='''np'''						,						).images

           UpperCamelCase__    :   List[Any]					=				low_res_image.size[0] * 4
           assert image.shape == (1, expected_height_width, expected_height_width, 3)




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





     """simple docstring"""







     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Any:




           """simple docstring"""







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






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   List[str]:




           """simple docstring"""







           UpperCamelCase__    :   Optional[Any]					=				load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
               '''/sd2-upscale/low_res_cat.png'''  )
           UpperCamelCase__    :   Optional[Any]					=				load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'''
               '''/upsampled_cat.npy'''  )

           UpperCamelCase__    :   Union[str, Any]					=				'''stabilityai/stable-diffusion-x4-upscaler'''
           UpperCamelCase__    :   Tuple					=				StableDiffusionUpscalePipeline.from_pretrained(lowercase__  )
           pipe.to(lowercase__  )
           pipe.set_progress_bar_config(disable=lowercase__  )
           pipe.enable_attention_slicing()

           UpperCamelCase__    :   Union[str, Any]					=				'''a cat sitting on a park bench'''

           UpperCamelCase__    :   Any					=				torch.manual_seed(0  )
           UpperCamelCase__    :   Tuple					=				pipe(
               prompt=lowercase__						,						image=lowercase__						,						generator=lowercase__						,						output_type='''np'''						,						)
           UpperCamelCase__    :   Optional[Any]					=				output.images[0]

           assert image.shape == (5_1_2, 5_1_2, 3)
           assert np.abs(expected_image - image  ).max() < 1e-3






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   str:




           """simple docstring"""







           UpperCamelCase__    :   str					=				load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
               '''/sd2-upscale/low_res_cat.png'''  )
           UpperCamelCase__    :   Optional[Any]					=				load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'''
               '''/upsampled_cat_fp16.npy'''  )

           UpperCamelCase__    :   int					=				'''stabilityai/stable-diffusion-x4-upscaler'''
           UpperCamelCase__    :   Union[str, Any]					=				StableDiffusionUpscalePipeline.from_pretrained(
               lowercase__						,						torch_dtype=torch.floataa						,						)
           pipe.to(lowercase__  )
           pipe.set_progress_bar_config(disable=lowercase__  )
           pipe.enable_attention_slicing()

           UpperCamelCase__    :   List[str]					=				'''a cat sitting on a park bench'''

           UpperCamelCase__    :   List[str]					=				torch.manual_seed(0  )
           UpperCamelCase__    :   str					=				pipe(
               prompt=lowercase__						,						image=lowercase__						,						generator=lowercase__						,						output_type='''np'''						,						)
           UpperCamelCase__    :   Optional[int]					=				output.images[0]

           assert image.shape == (5_1_2, 5_1_2, 3)
           assert np.abs(expected_image - image  ).max() < 5e-1






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Optional[Any]:




           """simple docstring"""







           torch.cuda.empty_cache()
           torch.cuda.reset_max_memory_allocated()
           torch.cuda.reset_peak_memory_stats()

           UpperCamelCase__    :   List[Any]					=				load_image(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
               '''/sd2-upscale/low_res_cat.png'''  )

           UpperCamelCase__    :   Optional[int]					=				'''stabilityai/stable-diffusion-x4-upscaler'''
           UpperCamelCase__    :   Union[str, Any]					=				StableDiffusionUpscalePipeline.from_pretrained(
               lowercase__						,						torch_dtype=torch.floataa						,						)
           pipe.to(lowercase__  )
           pipe.set_progress_bar_config(disable=lowercase__  )
           pipe.enable_attention_slicing(1  )
           pipe.enable_sequential_cpu_offload()

           UpperCamelCase__    :   str					=				'''a cat sitting on a park bench'''

           UpperCamelCase__    :   Tuple					=				torch.manual_seed(0  )
           UpperCamelCase__    :   List[str]					=				pipe(
               prompt=lowercase__						,						image=lowercase__						,						generator=lowercase__						,						num_inference_steps=5						,						output_type='''np'''						,						)

           UpperCamelCase__    :   Any					=				torch.cuda.max_memory_allocated()
           # make sure that less than 2.9 GB is allocated
           assert mem_bytes < 2.9 * 1_0**9




from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase      =logging.get_logger(__name__)

lowerCamelCase      ={
    "google/realm-cc-news-pretrained-embedder": (
        "https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/config.json"
    ),
    "google/realm-cc-news-pretrained-encoder": (
        "https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/config.json"
    ),
    "google/realm-cc-news-pretrained-scorer": (
        "https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/config.json"
    ),
    "google/realm-cc-news-pretrained-openqa": (
        "https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/config.json"
    ),
    "google/realm-orqa-nq-openqa": "https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/config.json",
    "google/realm-orqa-nq-reader": "https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/config.json",
    "google/realm-orqa-wq-openqa": "https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/config.json",
    "google/realm-orqa-wq-reader": "https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/config.json",
    # See all REALM models at https://huggingface.co/models?filter=realm
}




class 				_lowerCamelCase (      UpperCamelCase_ ):





     """simple docstring"""







     SCREAMING_SNAKE_CASE_		   =  '''realm'''






     def __init__(							self						,						__SCREAMING_SNAKE_CASE=3_0_5_2_2						,						__SCREAMING_SNAKE_CASE=7_6_8						,						__SCREAMING_SNAKE_CASE=1_2_8						,						__SCREAMING_SNAKE_CASE=1_2						,						__SCREAMING_SNAKE_CASE=1_2						,						__SCREAMING_SNAKE_CASE=8						,						__SCREAMING_SNAKE_CASE=3_0_7_2						,						__SCREAMING_SNAKE_CASE="gelu_new"						,						__SCREAMING_SNAKE_CASE=0.1						,						__SCREAMING_SNAKE_CASE=0.1						,						__SCREAMING_SNAKE_CASE=5_1_2						,						__SCREAMING_SNAKE_CASE=2						,						__SCREAMING_SNAKE_CASE=0.02						,						__SCREAMING_SNAKE_CASE=1e-12						,						__SCREAMING_SNAKE_CASE=2_5_6						,						__SCREAMING_SNAKE_CASE=1_0						,						__SCREAMING_SNAKE_CASE=1e-3						,						__SCREAMING_SNAKE_CASE=5						,						__SCREAMING_SNAKE_CASE=3_2_0						,						__SCREAMING_SNAKE_CASE=1_3_3_5_3_7_1_8						,						__SCREAMING_SNAKE_CASE=5_0_0_0						,						__SCREAMING_SNAKE_CASE=1						,						__SCREAMING_SNAKE_CASE=0						,						__SCREAMING_SNAKE_CASE=2						,						**__SCREAMING_SNAKE_CASE						,						)  ->   List[str]:




           """simple docstring"""







           super().__init__(pad_token_id=__SCREAMING_SNAKE_CASE						,						bos_token_id=__SCREAMING_SNAKE_CASE						,						eos_token_id=__SCREAMING_SNAKE_CASE						,						**__SCREAMING_SNAKE_CASE  )

           # Common config
           UpperCamelCase__    :   int					=				vocab_size
           UpperCamelCase__    :   Any					=				max_position_embeddings
           UpperCamelCase__    :   List[str]					=				hidden_size
           UpperCamelCase__    :   Union[str, Any]					=				retriever_proj_size
           UpperCamelCase__    :   int					=				num_hidden_layers
           UpperCamelCase__    :   Union[str, Any]					=				num_attention_heads
           UpperCamelCase__    :   str					=				num_candidates
           UpperCamelCase__    :   int					=				intermediate_size
           UpperCamelCase__    :   Dict					=				hidden_act
           UpperCamelCase__    :   Any					=				hidden_dropout_prob
           UpperCamelCase__    :   Any					=				attention_probs_dropout_prob
           UpperCamelCase__    :   Any					=				initializer_range
           UpperCamelCase__    :   Optional[Any]					=				type_vocab_size
           UpperCamelCase__    :   List[Any]					=				layer_norm_eps

           # Reader config
           UpperCamelCase__    :   List[Any]					=				span_hidden_size
           UpperCamelCase__    :   List[Any]					=				max_span_width
           UpperCamelCase__    :   Optional[Any]					=				reader_layer_norm_eps
           UpperCamelCase__    :   Optional[Any]					=				reader_beam_size
           UpperCamelCase__    :   int					=				reader_seq_len

           # Retrieval config
           UpperCamelCase__    :   List[str]					=				num_block_records
           UpperCamelCase__    :   Union[str, Any]					=				searcher_beam_size


