
'''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.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()




class      SCREAMING_SNAKE_CASE__  (			unittest.TestCase      ):





	'''simple docstring'''



	def 				A							(						self			:					Dict			):


				'''simple docstring'''


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



	@property
	def 				A							(						self			:					Optional[Any]			):


				'''simple docstring'''


				_snake_case						=    1
				_snake_case						=    3
				_snake_case						=    (32, 32)

				_snake_case						=    floats_tensor((batch_size, num_channels) + sizes    ,						rng=random.Random(0			)			).to(lowercase			)
				return image



	@property
	def 				A							(						self			:					Dict			):


				'''simple docstring'''


				torch.manual_seed(0			)
				_snake_case						=    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 				A							(						self			:					Optional[Any]			):


				'''simple docstring'''


				torch.manual_seed(0			)
				_snake_case						=    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 				A							(						self			:					List[str]			):


				'''simple docstring'''


				torch.manual_seed(0			)
				_snake_case						=    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=1_000    ,						)
				return CLIPTextModel(lowercase			)



	@property
	def 				A							(						self			:					Any			):


				'''simple docstring'''


				def extract(*lowercase			:					Optional[int]    ,						**lowercase			:					int			):




							class      SCREAMING_SNAKE_CASE__  :





								'''simple docstring'''



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


											'''simple docstring'''


											_snake_case						=    torch.ones([0]			)





								def 				A							(						self			:					str    ,						lowercase			:					Any			):


											'''simple docstring'''


											self.pixel_values.to(lowercase			)
											return self




							return Out()

				return extract



	def 				A							(						self			:					str			):


				'''simple docstring'''


				_snake_case						=    'cpu'  # ensure determinism for the device-dependent torch.Generator
				_snake_case						=    self.dummy_cond_unet
				_snake_case						=    DDIMScheduler(
				    beta_start=0.00085    ,						beta_end=0.012    ,						beta_schedule='scaled_linear'    ,						clip_sample=lowercase    ,						set_alpha_to_one=lowercase    ,						)

				_snake_case						=    self.dummy_vae
				_snake_case						=    self.dummy_text_encoder
				_snake_case						=    CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'			)

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

				_snake_case						=    'A painting of a squirrel eating a burger'

				_snake_case						=    torch.Generator(device=lowercase			).manual_seed(0			)
				_snake_case						=    sd_pipe([prompt]    ,						generator=lowercase    ,						guidance_scale=6.0    ,						num_inference_steps=2    ,						output_type='np'			)
				_snake_case						=    output.images

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

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

				assert image.shape == (1, 64, 64, 3)
				_snake_case						=    np.array([0.5753, 0.6114, 0.5001, 0.5034, 0.5470, 0.4729, 0.4971, 0.4867, 0.4867]			)

				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 				A							(						self			:					List[Any]			):


				'''simple docstring'''


				_snake_case						=    'cpu'  # ensure determinism for the device-dependent torch.Generator
				_snake_case						=    self.dummy_cond_unet
				_snake_case						=    PNDMScheduler(skip_prk_steps=lowercase			)
				_snake_case						=    self.dummy_vae
				_snake_case						=    self.dummy_text_encoder
				_snake_case						=    CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'			)

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

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

				_snake_case						=    output.images

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

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

				assert image.shape == (1, 64, 64, 3)
				_snake_case						=    np.array([0.5122, 0.5712, 0.4825, 0.5053, 0.5646, 0.4769, 0.5179, 0.4894, 0.4994]			)

				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 				A							(						self			:					int			):


				'''simple docstring'''


				_snake_case						=    StableDiffusionPipeline.from_pretrained(
				    'hf-internal-testing/tiny-stable-diffusion-lms-pipe'    ,						safety_checker=lowercase			)
				assert isinstance(lowercase    ,						lowercase			)
				assert isinstance(pipe.scheduler    ,						lowercase			)
				assert pipe.safety_checker is None

				_snake_case						=    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(lowercase			)
							_snake_case						=    StableDiffusionPipeline.from_pretrained(lowercase			)

				# sanity check that the pipeline still works
				assert pipe.safety_checker is None
				_snake_case						=    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 				A							(						self			:					int			):


				'''simple docstring'''


				_snake_case						=    self.dummy_cond_unet
				_snake_case						=    PNDMScheduler(skip_prk_steps=lowercase			)
				_snake_case						=    self.dummy_vae
				_snake_case						=    self.dummy_text_encoder
				_snake_case						=    CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'			)

				# put models in fp16
				_snake_case						=    unet.half()
				_snake_case						=    vae.half()
				_snake_case						=    bert.half()

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

				_snake_case						=    'A painting of a squirrel eating a burger'
				_snake_case						=    sd_pipe([prompt]    ,						num_inference_steps=2    ,						output_type='np'			).images

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




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





	'''simple docstring'''



	def 				A							(						self			:					Any			):


				'''simple docstring'''


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



	def 				A							(						self			:					Dict			):


				'''simple docstring'''


				_snake_case						=    'cuda'
				_snake_case						=    StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'			)
				_snake_case						=    pipe.to(lowercase			)
				pipe.set_progress_bar_config(disable=lowercase			)

				_snake_case						=    'a photo of a cat'
				_snake_case						=    {
				    'editing_prompt': ['sunglasses'],
				    'reverse_editing_direction': [False],
				    'edit_warmup_steps': 10,
				    'edit_guidance_scale': 6,
				    'edit_threshold': 0.95,
				    'edit_momentum_scale': 0.5,
				    'edit_mom_beta': 0.6,
				}

				_snake_case						=    3
				_snake_case						=    7

				# no sega enabled
				_snake_case						=    torch.Generator(lowercase			)
				generator.manual_seed(lowercase			)
				_snake_case						=    pipe(
				    [prompt]    ,						generator=lowercase    ,						guidance_scale=lowercase    ,						num_inference_steps=50    ,						output_type='np'    ,						width=512    ,						height=512    ,						)

				_snake_case						=    output.images
				_snake_case						=    image[0, -3:, -3:, -1]
				_snake_case						=    [
				    0.34673113,
				    0.38492733,
				    0.37597352,
				    0.34086335,
				    0.35650748,
				    0.35579205,
				    0.3384763,
				    0.34340236,
				    0.3573271,
				]

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

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

				# with sega enabled
				# generator = torch.manual_seed(seed)
				generator.manual_seed(lowercase			)
				_snake_case						=    pipe(
				    [prompt]    ,						generator=lowercase    ,						guidance_scale=lowercase    ,						num_inference_steps=50    ,						output_type='np'    ,						width=512    ,						height=512    ,						**lowercase    ,						)

				_snake_case						=    output.images
				_snake_case						=    image[0, -3:, -3:, -1]
				_snake_case						=    [
				    0.41887826,
				    0.37728766,
				    0.30138272,
				    0.41416335,
				    0.41664985,
				    0.36283392,
				    0.36191246,
				    0.43364465,
				    0.43001732,
				]

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

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



	def 				A							(						self			:					Any			):


				'''simple docstring'''


				_snake_case						=    'cuda'
				_snake_case						=    StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'			)
				_snake_case						=    pipe.to(lowercase			)
				pipe.set_progress_bar_config(disable=lowercase			)

				_snake_case						=    'an image of a crowded boulevard, realistic, 4k'
				_snake_case						=    {
				    'editing_prompt': 'crowd, crowded, people',
				    'reverse_editing_direction': True,
				    'edit_warmup_steps': 10,
				    'edit_guidance_scale': 8.3,
				    'edit_threshold': 0.9,
				    'edit_momentum_scale': 0.5,
				    'edit_mom_beta': 0.6,
				}

				_snake_case						=    9
				_snake_case						=    7

				# no sega enabled
				_snake_case						=    torch.Generator(lowercase			)
				generator.manual_seed(lowercase			)
				_snake_case						=    pipe(
				    [prompt]    ,						generator=lowercase    ,						guidance_scale=lowercase    ,						num_inference_steps=50    ,						output_type='np'    ,						width=512    ,						height=512    ,						)

				_snake_case						=    output.images
				_snake_case						=    image[0, -3:, -3:, -1]
				_snake_case						=    [
				    0.43497998,
				    0.91814065,
				    0.7540739,
				    0.55580205,
				    0.8467265,
				    0.5389691,
				    0.62574506,
				    0.58897763,
				    0.50926757,
				]

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

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

				# with sega enabled
				# generator = torch.manual_seed(seed)
				generator.manual_seed(lowercase			)
				_snake_case						=    pipe(
				    [prompt]    ,						generator=lowercase    ,						guidance_scale=lowercase    ,						num_inference_steps=50    ,						output_type='np'    ,						width=512    ,						height=512    ,						**lowercase    ,						)

				_snake_case						=    output.images
				_snake_case						=    image[0, -3:, -3:, -1]
				_snake_case						=    [
				    0.3089719,
				    0.30500144,
				    0.29016042,
				    0.30630964,
				    0.325687,
				    0.29419225,
				    0.2908091,
				    0.28723598,
				    0.27696294,
				]

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

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



	def 				A							(						self			:					Any			):


				'''simple docstring'''


				_snake_case						=    'cuda'
				_snake_case						=    StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'			)
				_snake_case						=    pipe.to(lowercase			)
				pipe.set_progress_bar_config(disable=lowercase			)

				_snake_case						=    'a castle next to a river'
				_snake_case						=    {
				    'editing_prompt': ['boat on a river, boat', 'monet, impression, sunrise'],
				    'reverse_editing_direction': False,
				    'edit_warmup_steps': [15, 18],
				    'edit_guidance_scale': 6,
				    'edit_threshold': [0.9, 0.8],
				    'edit_momentum_scale': 0.5,
				    'edit_mom_beta': 0.6,
				}

				_snake_case						=    48
				_snake_case						=    7

				# no sega enabled
				_snake_case						=    torch.Generator(lowercase			)
				generator.manual_seed(lowercase			)
				_snake_case						=    pipe(
				    [prompt]    ,						generator=lowercase    ,						guidance_scale=lowercase    ,						num_inference_steps=50    ,						output_type='np'    ,						width=512    ,						height=512    ,						)

				_snake_case						=    output.images
				_snake_case						=    image[0, -3:, -3:, -1]
				_snake_case						=    [
				    0.75163555,
				    0.76037145,
				    0.61785,
				    0.9189673,
				    0.8627701,
				    0.85189694,
				    0.8512813,
				    0.87012076,
				    0.8312857,
				]

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

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

				# with sega enabled
				# generator = torch.manual_seed(seed)
				generator.manual_seed(lowercase			)
				_snake_case						=    pipe(
				    [prompt]    ,						generator=lowercase    ,						guidance_scale=lowercase    ,						num_inference_steps=50    ,						output_type='np'    ,						width=512    ,						height=512    ,						**lowercase    ,						)

				_snake_case						=    output.images
				_snake_case						=    image[0, -3:, -3:, -1]
				_snake_case						=    [
				    0.73553365,
				    0.7537271,
				    0.74341905,
				    0.66480356,
				    0.6472925,
				    0.63039416,
				    0.64812905,
				    0.6749717,
				    0.6517102,
				]

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

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





	def 				A							(						self			:					str			):


				'''simple docstring'''


				_snake_case						=    'cuda'
				_snake_case						=    StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'    ,						torch_dtype=torch.floataa			)
				_snake_case						=    pipe.to(lowercase			)
				pipe.set_progress_bar_config(disable=lowercase			)

				_snake_case						=    'a photo of a cat'
				_snake_case						=    {
				    'editing_prompt': ['sunglasses'],
				    'reverse_editing_direction': [False],
				    'edit_warmup_steps': 10,
				    'edit_guidance_scale': 6,
				    'edit_threshold': 0.95,
				    'edit_momentum_scale': 0.5,
				    'edit_mom_beta': 0.6,
				}

				_snake_case						=    3
				_snake_case						=    7

				# no sega enabled
				_snake_case						=    torch.Generator(lowercase			)
				generator.manual_seed(lowercase			)
				_snake_case						=    pipe(
				    [prompt]    ,						generator=lowercase    ,						guidance_scale=lowercase    ,						num_inference_steps=50    ,						output_type='np'    ,						width=512    ,						height=512    ,						)

				_snake_case						=    output.images
				_snake_case						=    image[0, -3:, -3:, -1]
				_snake_case						=    [
				    0.34887695,
				    0.3876953,
				    0.375,
				    0.34423828,
				    0.3581543,
				    0.35717773,
				    0.3383789,
				    0.34570312,
				    0.359375,
				]

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

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

				# with sega enabled
				# generator = torch.manual_seed(seed)
				generator.manual_seed(lowercase			)
				_snake_case						=    pipe(
				    [prompt]    ,						generator=lowercase    ,						guidance_scale=lowercase    ,						num_inference_steps=50    ,						output_type='np'    ,						width=512    ,						height=512    ,						**lowercase    ,						)

				_snake_case						=    output.images
				_snake_case						=    image[0, -3:, -3:, -1]
				_snake_case						=    [
				    0.42285156,
				    0.36914062,
				    0.29077148,
				    0.42041016,
				    0.41918945,
				    0.35498047,
				    0.3618164,
				    0.4423828,
				    0.43115234,
				]

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

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

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding




class      SCREAMING_SNAKE_CASE__  (			UpperCAmelCase      ):





	'''simple docstring'''



	_UpperCAmelCase :	int					     =      ["image_processor", "tokenizer"]
	_UpperCAmelCase :	Tuple					     =      "EfficientNetImageProcessor"
	_UpperCAmelCase :	Optional[Any]					     =      ("BertTokenizer", "BertTokenizerFast")



	def __init__(						self			:					Tuple    ,						lowercase			:					List[Any]    ,						lowercase			:					Optional[int]			):


				'''simple docstring'''


				super().__init__(lowercase    ,						lowercase			)



	def __call__(						self			:					Union[str, Any]    ,						lowercase			:					Dict=None    ,						lowercase			:					Optional[int]=None    ,						lowercase			:					List[Any]="max_length"    ,						lowercase			:					Optional[int]=64    ,						lowercase			:					int=None    ,						**lowercase			:					Tuple			):


				'''simple docstring'''


				if text is None and images is None:
							raise ValueError('You have to specify either text or images. Both cannot be none.'			)

				if text is not None:
							_snake_case						=    self.tokenizer(
							    lowercase    ,						padding=lowercase    ,						max_length=lowercase    ,						return_tensors=lowercase    ,						**lowercase			)

				if images is not None:
							_snake_case						=    self.image_processor(lowercase    ,						return_tensors=lowercase    ,						**lowercase			)

				if text is not None and images is not None:
							_snake_case						=    image_features.pixel_values
							return encoding
				elif text is not None:
							return encoding
				else:
							return BatchEncoding(data=dict(**lowercase			)    ,						tensor_type=lowercase			)



	def 				A							(						self			:					Any    ,						*lowercase			:					List[str]    ,						**lowercase			:					Optional[Any]			):


				'''simple docstring'''


				return self.tokenizer.batch_decode(*lowercase    ,						**lowercase			)



	def 				A							(						self			:					List[Any]    ,						*lowercase			:					Union[str, Any]    ,						**lowercase			:					Tuple			):


				'''simple docstring'''


				return self.tokenizer.decode(*lowercase    ,						**lowercase			)





	@property
	def 				A							(						self			:					List[str]			):


				'''simple docstring'''


				_snake_case						=    self.tokenizer.model_input_names
				_snake_case						=    self.image_processor.model_input_names
				return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names			)			)