


'''simple docstring'''



import gc
import random
import tempfile
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion_safe import StableDiffusionPipelineSafe as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import require_torch_gpu





class   UpperCAmelCase_    (unittest.TestCase    ):





    """simple docstring"""






    def  lowercase_							(				self      )							->	List[str]:
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()







    @property
    def  lowercase_							(				self      )							->	Any:
         __lowerCamelCase  :      str														=						1
         __lowerCamelCase  :      int														=						3
         __lowerCamelCase  :      Tuple														=						(32, 32)

         __lowerCamelCase  :      Optional[int]														=						floats_tensor((batch_size, num_channels) + sizes     ,	rng=random.Random(0      )      ).to(_snake_case      )
         return image







    @property
    def  lowercase_							(				self      )							->	Tuple:
         torch.manual_seed(0      )
         __lowerCamelCase  :      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')     ,	cross_attention_dim=32     ,	)
         return model







    @property
    def  lowercase_							(				self      )							->	str:
         torch.manual_seed(0      )
         __lowerCamelCase  :      List[str]														=						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     ,	)
         return model







    @property
    def  lowercase_							(				self      )							->	Dict:
         torch.manual_seed(0      )
         __lowerCamelCase  :      Optional[int]														=						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=10_00     ,	)
         return CLIPTextModel(_snake_case      )







    @property
    def  lowercase_							(				self      )							->	int:
         def extract(*SCREAMING_SNAKE_CASE_     ,	**SCREAMING_SNAKE_CASE_      ):





              class   UpperCAmelCase_    :





                  """simple docstring"""






                  def __init__(				self      )							->	Tuple:
                       __lowerCamelCase  :      Any														=						torch.ones([0]      )







                  def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_      )							->	Optional[int]:
                       self.pixel_values.to(_snake_case      )
                       return self







              return Out()

         return extract







    def  lowercase_							(				self      )							->	Union[str, Any]:
         __lowerCamelCase  :      Any														=						'cpu'  # ensure determinism for the device-dependent torch.Generator
         __lowerCamelCase  :      Any														=						self.dummy_cond_unet
         __lowerCamelCase  :      Optional[int]														=						DDIMScheduler(
             beta_start=0.0_0_0_8_5     ,	beta_end=0.0_1_2     ,	beta_schedule='scaled_linear'     ,	clip_sample=_snake_case     ,	set_alpha_to_one=_snake_case     ,	)

         __lowerCamelCase  :      List[Any]														=						self.dummy_vae
         __lowerCamelCase  :      str														=						self.dummy_text_encoder
         __lowerCamelCase  :      Optional[int]														=						CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'      )

         # make sure here that pndm scheduler skips prk
         __lowerCamelCase  :      str														=						StableDiffusionPipeline(
             unet=_snake_case     ,	scheduler=_snake_case     ,	vae=_snake_case     ,	text_encoder=_snake_case     ,	tokenizer=_snake_case     ,	safety_checker=_snake_case     ,	feature_extractor=self.dummy_extractor     ,	)
         __lowerCamelCase  :      int														=						sd_pipe.to(_snake_case      )
         sd_pipe.set_progress_bar_config(disable=_snake_case      )

         __lowerCamelCase  :      Union[str, Any]														=						'A painting of a squirrel eating a burger'

         __lowerCamelCase  :      List[Any]														=						torch.Generator(device=_snake_case      ).manual_seed(0      )
         __lowerCamelCase  :      Optional[Any]														=						sd_pipe([prompt]     ,	generator=_snake_case     ,	guidance_scale=6.0     ,	num_inference_steps=2     ,	output_type='np'      )
         __lowerCamelCase  :      Tuple														=						output.images

         __lowerCamelCase  :      Optional[Any]														=						torch.Generator(device=_snake_case      ).manual_seed(0      )
         __lowerCamelCase  :      Optional[int]														=						sd_pipe(
             [prompt]     ,	generator=_snake_case     ,	guidance_scale=6.0     ,	num_inference_steps=2     ,	output_type='np'     ,	return_dict=_snake_case     ,	)[0]

         __lowerCamelCase  :      Dict														=						image[0, -3:, -3:, -1]
         __lowerCamelCase  :      Dict														=						image_from_tuple[0, -3:, -3:, -1]

         assert image.shape == (1, 64, 64, 3)
         __lowerCamelCase  :      int														=						np.array([0.5_7_5_6, 0.6_1_1_8, 0.5_0_0_5, 0.5_0_4_1, 0.5_4_7_1, 0.4_7_2_6, 0.4_9_7_6, 0.4_8_6_5, 0.4_8_6_4]      )

         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  lowercase_							(				self      )							->	Any:
         __lowerCamelCase  :      Tuple														=						'cpu'  # ensure determinism for the device-dependent torch.Generator
         __lowerCamelCase  :      str														=						self.dummy_cond_unet
         __lowerCamelCase  :      Optional[Any]														=						PNDMScheduler(skip_prk_steps=_snake_case      )
         __lowerCamelCase  :      str														=						self.dummy_vae
         __lowerCamelCase  :      Dict														=						self.dummy_text_encoder
         __lowerCamelCase  :      Dict														=						CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'      )

         # make sure here that pndm scheduler skips prk
         __lowerCamelCase  :      int														=						StableDiffusionPipeline(
             unet=_snake_case     ,	scheduler=_snake_case     ,	vae=_snake_case     ,	text_encoder=_snake_case     ,	tokenizer=_snake_case     ,	safety_checker=_snake_case     ,	feature_extractor=self.dummy_extractor     ,	)
         __lowerCamelCase  :      Union[str, Any]														=						sd_pipe.to(_snake_case      )
         sd_pipe.set_progress_bar_config(disable=_snake_case      )

         __lowerCamelCase  :      Any														=						'A painting of a squirrel eating a burger'
         __lowerCamelCase  :      Tuple														=						torch.Generator(device=_snake_case      ).manual_seed(0      )
         __lowerCamelCase  :      int														=						sd_pipe([prompt]     ,	generator=_snake_case     ,	guidance_scale=6.0     ,	num_inference_steps=2     ,	output_type='np'      )

         __lowerCamelCase  :      Optional[int]														=						output.images

         __lowerCamelCase  :      Union[str, Any]														=						torch.Generator(device=_snake_case      ).manual_seed(0      )
         __lowerCamelCase  :      Any														=						sd_pipe(
             [prompt]     ,	generator=_snake_case     ,	guidance_scale=6.0     ,	num_inference_steps=2     ,	output_type='np'     ,	return_dict=_snake_case     ,	)[0]

         __lowerCamelCase  :      int														=						image[0, -3:, -3:, -1]
         __lowerCamelCase  :      Optional[Any]														=						image_from_tuple[0, -3:, -3:, -1]

         assert image.shape == (1, 64, 64, 3)
         __lowerCamelCase  :      Optional[Any]														=						np.array([0.5_1_2_5, 0.5_7_1_6, 0.4_8_2_8, 0.5_0_6_0, 0.5_6_5_0, 0.4_7_6_8, 0.5_1_8_5, 0.4_8_9_5, 0.4_9_9_3]      )

         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  lowercase_							(				self      )							->	List[str]:
         __lowerCamelCase  :      int														=						StableDiffusionPipeline.from_pretrained(
             'hf-internal-testing/tiny-stable-diffusion-lms-pipe'     ,	safety_checker=_snake_case      )
         assert isinstance(_snake_case     ,	_snake_case      )
         assert isinstance(pipe.scheduler     ,	_snake_case      )
         assert pipe.safety_checker is None

         __lowerCamelCase  :      int														=						pipe('example prompt'     ,	num_inference_steps=2      ).images[0]
         assert image is not None

         # check that there's no error when saving a pipeline with one of the models being None
         with tempfile.TemporaryDirectory() as tmpdirname:
              pipe.save_pretrained(_snake_case      )
              __lowerCamelCase  :      str														=						StableDiffusionPipeline.from_pretrained(_snake_case      )

         # sanity check that the pipeline still works
         assert pipe.safety_checker is None
         __lowerCamelCase  :      Any														=						pipe('example prompt'     ,	num_inference_steps=2      ).images[0]
         assert image is not None







    @unittest.skipIf(torch_device != 'cuda'     ,	'This test requires a GPU'      )
    def  lowercase_							(				self      )							->	Union[str, Any]:
         __lowerCamelCase  :      Any														=						self.dummy_cond_unet
         __lowerCamelCase  :      Dict														=						PNDMScheduler(skip_prk_steps=_snake_case      )
         __lowerCamelCase  :      str														=						self.dummy_vae
         __lowerCamelCase  :      Optional[int]														=						self.dummy_text_encoder
         __lowerCamelCase  :      str														=						CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'      )

         # put models in fp16
         __lowerCamelCase  :      Tuple														=						unet.half()
         __lowerCamelCase  :      int														=						vae.half()
         __lowerCamelCase  :      Tuple														=						bert.half()

         # make sure here that pndm scheduler skips prk
         __lowerCamelCase  :      Optional[int]														=						StableDiffusionPipeline(
             unet=_snake_case     ,	scheduler=_snake_case     ,	vae=_snake_case     ,	text_encoder=_snake_case     ,	tokenizer=_snake_case     ,	safety_checker=_snake_case     ,	feature_extractor=self.dummy_extractor     ,	)
         __lowerCamelCase  :      Any														=						sd_pipe.to(_snake_case      )
         sd_pipe.set_progress_bar_config(disable=_snake_case      )

         __lowerCamelCase  :      Optional[Any]														=						'A painting of a squirrel eating a burger'
         __lowerCamelCase  :      int														=						sd_pipe([prompt]     ,	num_inference_steps=2     ,	output_type='np'      ).images

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







@nightly
@require_torch_gpu
class   UpperCAmelCase_    (unittest.TestCase    ):





    """simple docstring"""






    def  lowercase_							(				self      )							->	Any:
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()







    def  lowercase_							(				self      )							->	Tuple:
         __lowerCamelCase  :      Optional[int]														=						StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'     ,	safety_checker=_snake_case      )
         __lowerCamelCase  :      Optional[int]														=						LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config      )
         __lowerCamelCase  :      List[str]														=						sd_pipe.to(_snake_case      )
         sd_pipe.set_progress_bar_config(disable=_snake_case      )

         __lowerCamelCase  :      Dict														=						(
             'portrait of girl with smokey eyes makeup in abandoned hotel, grange clothes, redshift, wide high angle'
             ' coloured polaroid photograph with flash, kodak film, hyper real, stunning moody cinematography, with'
             ' anamorphic lenses, by maripol, fallen angels by wong kar - wai, style of suspiria and neon demon and'
             ' children from bahnhof zoo, detailed '
         )
         __lowerCamelCase  :      List[Any]														=						40_03_66_03_46
         __lowerCamelCase  :      str														=						7

         # without safety guidance (sld_guidance_scale = 0)
         __lowerCamelCase  :      Tuple														=						torch.manual_seed(_snake_case      )
         __lowerCamelCase  :      Union[str, Any]														=						sd_pipe(
             [prompt]     ,	generator=_snake_case     ,	guidance_scale=_snake_case     ,	num_inference_steps=50     ,	output_type='np'     ,	width=5_12     ,	height=5_12     ,	sld_guidance_scale=0     ,	)

         __lowerCamelCase  :      Optional[Any]														=						output.images
         __lowerCamelCase  :      Optional[Any]														=						image[0, -3:, -3:, -1]
         __lowerCamelCase  :      int														=						[0.2_2_7_8, 0.2_2_3_1, 0.2_2_4_9, 0.2_3_3_3, 0.2_3_0_3, 0.1_8_8_5, 0.2_2_7_3, 0.2_1_4_4, 0.2_1_7_6]

         assert image.shape == (1, 5_12, 5_12, 3)

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

         # without safety guidance (strong configuration)
         __lowerCamelCase  :      List[str]														=						torch.manual_seed(_snake_case      )
         __lowerCamelCase  :      Tuple														=						sd_pipe(
             [prompt]     ,	generator=_snake_case     ,	guidance_scale=_snake_case     ,	num_inference_steps=50     ,	output_type='np'     ,	width=5_12     ,	height=5_12     ,	sld_guidance_scale=20_00     ,	sld_warmup_steps=7     ,	sld_threshold=0.0_2_5     ,	sld_momentum_scale=0.5     ,	sld_mom_beta=0.7     ,	)

         __lowerCamelCase  :      List[Any]														=						output.images
         __lowerCamelCase  :      Tuple														=						image[0, -3:, -3:, -1]
         __lowerCamelCase  :      Optional[Any]														=						[0.2_3_8_3, 0.2_2_7_6, 0.2_3_6, 0.2_1_9_2, 0.2_1_8_6, 0.2_0_5_3, 0.1_9_7_1, 0.1_9_0_1, 0.1_7_1_9]

         assert image.shape == (1, 5_12, 5_12, 3)

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







    def  lowercase_							(				self      )							->	int:
         __lowerCamelCase  :      Any														=						StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'     ,	safety_checker=_snake_case      )
         __lowerCamelCase  :      List[Any]														=						LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config      )
         __lowerCamelCase  :      Dict														=						sd_pipe.to(_snake_case      )
         sd_pipe.set_progress_bar_config(disable=_snake_case      )

         __lowerCamelCase  :      Union[str, Any]														=						'padme amidala taking a bath artwork, safe for work, no nudity'
         __lowerCamelCase  :      List[Any]														=						27_34_97_17_55
         __lowerCamelCase  :      List[Any]														=						7

         __lowerCamelCase  :      Optional[int]														=						torch.manual_seed(_snake_case      )
         __lowerCamelCase  :      List[str]														=						sd_pipe(
             [prompt]     ,	generator=_snake_case     ,	guidance_scale=_snake_case     ,	num_inference_steps=50     ,	output_type='np'     ,	width=5_12     ,	height=5_12     ,	sld_guidance_scale=0     ,	)

         __lowerCamelCase  :      Dict														=						output.images
         __lowerCamelCase  :      str														=						image[0, -3:, -3:, -1]
         __lowerCamelCase  :      List[Any]														=						[0.3_5_0_2, 0.3_6_2_2, 0.3_3_9_6, 0.3_6_4_2, 0.3_4_7_8, 0.3_3_1_8, 0.3_5, 0.3_3_4_8, 0.3_2_9_7]

         assert image.shape == (1, 5_12, 5_12, 3)

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

         __lowerCamelCase  :      Optional[Any]														=						torch.manual_seed(_snake_case      )
         __lowerCamelCase  :      Union[str, Any]														=						sd_pipe(
             [prompt]     ,	generator=_snake_case     ,	guidance_scale=_snake_case     ,	num_inference_steps=50     ,	output_type='np'     ,	width=5_12     ,	height=5_12     ,	sld_guidance_scale=20_00     ,	sld_warmup_steps=7     ,	sld_threshold=0.0_2_5     ,	sld_momentum_scale=0.5     ,	sld_mom_beta=0.7     ,	)

         __lowerCamelCase  :      Tuple														=						output.images
         __lowerCamelCase  :      List[str]														=						image[0, -3:, -3:, -1]
         __lowerCamelCase  :      Optional[Any]														=						[0.5_5_3_1, 0.5_2_0_6, 0.4_8_9_5, 0.5_1_5_6, 0.5_1_8_2, 0.4_7_5_1, 0.4_8_0_2, 0.4_8_0_3, 0.4_4_4_3]

         assert image.shape == (1, 5_12, 5_12, 3)

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







    def  lowercase_							(				self      )							->	Optional[Any]:
         __lowerCamelCase  :      Any														=						StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'      )
         __lowerCamelCase  :      int														=						sd_pipe.to(_snake_case      )
         sd_pipe.set_progress_bar_config(disable=_snake_case      )

         __lowerCamelCase  :      Optional[Any]														=						(
             'the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c.'
             ' leyendecker'
         )
         __lowerCamelCase  :      List[Any]														=						10_44_35_52_34
         __lowerCamelCase  :      Optional[Any]														=						12

         __lowerCamelCase  :      str														=						torch.manual_seed(_snake_case      )
         __lowerCamelCase  :      Dict														=						sd_pipe(
             [prompt]     ,	generator=_snake_case     ,	guidance_scale=_snake_case     ,	num_inference_steps=50     ,	output_type='np'     ,	width=5_12     ,	height=5_12     ,	sld_guidance_scale=0     ,	)

         __lowerCamelCase  :      str														=						output.images
         __lowerCamelCase  :      Any														=						image[0, -3:, -3:, -1]
         __lowerCamelCase  :      Dict														=						np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]      )

         assert image.shape == (1, 5_12, 5_12, 3)

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

         __lowerCamelCase  :      Tuple														=						torch.manual_seed(_snake_case      )
         __lowerCamelCase  :      List[Any]														=						sd_pipe(
             [prompt]     ,	generator=_snake_case     ,	guidance_scale=_snake_case     ,	num_inference_steps=50     ,	output_type='np'     ,	width=5_12     ,	height=5_12     ,	sld_guidance_scale=20_00     ,	sld_warmup_steps=7     ,	sld_threshold=0.0_2_5     ,	sld_momentum_scale=0.5     ,	sld_mom_beta=0.7     ,	)

         __lowerCamelCase  :      Union[str, Any]														=						output.images
         __lowerCamelCase  :      Tuple														=						image[0, -3:, -3:, -1]
         __lowerCamelCase  :      Optional[Any]														=						np.array([0.5_8_1_8, 0.6_2_8_5, 0.6_8_3_5, 0.6_0_1_9, 0.6_2_5, 0.6_7_5_4, 0.6_0_9_6, 0.6_3_3_4, 0.6_5_6_1]      )
         assert image.shape == (1, 5_12, 5_12, 3)

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








"""simple docstring"""





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


UpperCamelCase_       :    Any						      =			logging.get_logger(__name__)

# TODO: upload to AWS
UpperCamelCase_       :    List[str]						      =			{
    '''yjernite/retribert-base-uncased''': (
        '''https://huggingface.co/yjernite/retribert-base-uncased/resolve/main/config.json'''
    ),
}



class __lowerCAmelCase					(       _lowercase   ):




			"""simple docstring"""







			snake_case						   =					"retribert"




			def __init__(					self    :       str							, _snake_case    :       List[str]=30_522							, _snake_case    :       List[Any]=768							, _snake_case    :       List[Any]=8							, _snake_case    :       Union[str, Any]=12							, _snake_case    :       Optional[Any]=3_072							, _snake_case    :       List[Any]="gelu"							, _snake_case    :       Union[str, Any]=0.1							, _snake_case    :       Dict=0.1							, _snake_case    :       str=512							, _snake_case    :       Tuple=2							, _snake_case    :       Dict=0.0_2							, _snake_case    :       Optional[Any]=1e-12							, _snake_case    :       Tuple=True							, _snake_case    :       Union[str, Any]=128							, _snake_case    :       Union[str, Any]=0							, **_snake_case    :       Any							, )   ->			Optional[int]:







								"""simple docstring"""
								super().__init__(pad_token_id=_snake_case							, **_snake_case						)

								A_								=    vocab_size
								A_								=    hidden_size
								A_								=    num_hidden_layers
								A_								=    num_attention_heads
								A_								=    hidden_act
								A_								=    intermediate_size
								A_								=    hidden_dropout_prob
								A_								=    attention_probs_dropout_prob
								A_								=    max_position_embeddings
								A_								=    type_vocab_size
								A_								=    initializer_range
								A_								=    layer_norm_eps
								A_								=    share_encoders
								A_								=    projection_dim




