

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      lowerCAmelCase__  (			unittest.TestCase    ):






		def 				__A (	self	:							List[Any]						)  ->							Dict:
						# clean up the VRAM after each test
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()






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

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






		@property
		def 				__A (	self	:							Tuple						)  ->							str:
						torch.manual_seed(0						)
						__lowerCamelCase					      =			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	:							Tuple						)  ->							Optional[int]:
						torch.manual_seed(0						)
						__lowerCamelCase					      =			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	:							Optional[Any]						)  ->							int:
						torch.manual_seed(0						)
						__lowerCamelCase					      =			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 				__A (	self	:							Optional[int]						)  ->							int:
						def extract(*SCREAMING_SNAKE_CASE__	:							List[Any]							,    **SCREAMING_SNAKE_CASE__	:							List[Any]						):

										class      lowerCAmelCase__  :






												def __init__(	self	:							int						)  ->							Union[str, Any]:
																__lowerCamelCase					      =			torch.ones([0]						)







												def 				__A (	self	:							Optional[Any]							,    SCREAMING_SNAKE_CASE__	:							List[str]						)  ->							Optional[int]:
																self.pixel_values.to(snake_case_						)
																return self





										return Out()

						return extract






		def 				__A (	self	:							str						)  ->							Optional[int]:
						__lowerCamelCase					      =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
						__lowerCamelCase					      =			self.dummy_cond_unet
						__lowerCamelCase					      =			DDIMScheduler(
						    beta_start=0.00085							,    beta_end=0.012							,    beta_schedule='''scaled_linear'''							,    clip_sample=snake_case_							,    set_alpha_to_one=snake_case_							,    )

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

						# make sure here that pndm scheduler skips prk
						__lowerCamelCase					      =			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					      =			sd_pipe.to(snake_case_						)
						sd_pipe.set_progress_bar_config(disable=snake_case_						)

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

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

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

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

						assert image.shape == (1, 64, 64, 3)
						__lowerCamelCase					      =			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	:							Optional[int]						)  ->							Optional[int]:
						__lowerCamelCase					      =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
						__lowerCamelCase					      =			self.dummy_cond_unet
						__lowerCamelCase					      =			PNDMScheduler(skip_prk_steps=snake_case_						)
						__lowerCamelCase					      =			self.dummy_vae
						__lowerCamelCase					      =			self.dummy_text_encoder
						__lowerCamelCase					      =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''						)

						# make sure here that pndm scheduler skips prk
						__lowerCamelCase					      =			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					      =			sd_pipe.to(snake_case_						)
						sd_pipe.set_progress_bar_config(disable=snake_case_						)

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

						__lowerCamelCase					      =			output.images

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

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

						assert image.shape == (1, 64, 64, 3)
						__lowerCamelCase					      =			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	:							Tuple						)  ->							int:
						__lowerCamelCase					      =			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					      =			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					      =			StableDiffusionPipeline.from_pretrained(snake_case_						)

						# sanity check that the pipeline still works
						assert pipe.safety_checker is None
						__lowerCamelCase					      =			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	:							Dict						)  ->							str:
						__lowerCamelCase					      =			self.dummy_cond_unet
						__lowerCamelCase					      =			PNDMScheduler(skip_prk_steps=snake_case_						)
						__lowerCamelCase					      =			self.dummy_vae
						__lowerCamelCase					      =			self.dummy_text_encoder
						__lowerCamelCase					      =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''						)

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

						# make sure here that pndm scheduler skips prk
						__lowerCamelCase					      =			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					      =			sd_pipe.to(snake_case_						)
						sd_pipe.set_progress_bar_config(disable=snake_case_						)

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

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





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






		def 				__A (	self	:							int						)  ->							Union[str, Any]:
						# clean up the VRAM after each test
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()






		def 				__A (	self	:							int						)  ->							Union[str, Any]:
						__lowerCamelCase					      =			'''cuda'''
						__lowerCamelCase					      =			StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''						)
						__lowerCamelCase					      =			pipe.to(snake_case_						)
						pipe.set_progress_bar_config(disable=snake_case_						)

						__lowerCamelCase					      =			'''a photo of a cat'''
						__lowerCamelCase					      =			{
						    '''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,
						}

						__lowerCamelCase					      =			3
						__lowerCamelCase					      =			7

						# no sega enabled
						__lowerCamelCase					      =			torch.Generator(snake_case_						)
						generator.manual_seed(snake_case_						)
						__lowerCamelCase					      =			pipe(
						    [prompt]							,    generator=snake_case_							,    guidance_scale=snake_case_							,    num_inference_steps=50							,    output_type='''np'''							,    width=5_12							,    height=5_12							,    )

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

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

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

						# with sega enabled
						# generator = torch.manual_seed(seed)
						generator.manual_seed(snake_case_						)
						__lowerCamelCase					      =			pipe(
						    [prompt]							,    generator=snake_case_							,    guidance_scale=snake_case_							,    num_inference_steps=50							,    output_type='''np'''							,    width=5_12							,    height=5_12							,    **snake_case_							,    )

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

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

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






		def 				__A (	self	:							List[str]						)  ->							Optional[Any]:
						__lowerCamelCase					      =			'''cuda'''
						__lowerCamelCase					      =			StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''						)
						__lowerCamelCase					      =			pipe.to(snake_case_						)
						pipe.set_progress_bar_config(disable=snake_case_						)

						__lowerCamelCase					      =			'''an image of a crowded boulevard, realistic, 4k'''
						__lowerCamelCase					      =			{
						    '''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,
						}

						__lowerCamelCase					      =			9
						__lowerCamelCase					      =			7

						# no sega enabled
						__lowerCamelCase					      =			torch.Generator(snake_case_						)
						generator.manual_seed(snake_case_						)
						__lowerCamelCase					      =			pipe(
						    [prompt]							,    generator=snake_case_							,    guidance_scale=snake_case_							,    num_inference_steps=50							,    output_type='''np'''							,    width=5_12							,    height=5_12							,    )

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

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

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

						# with sega enabled
						# generator = torch.manual_seed(seed)
						generator.manual_seed(snake_case_						)
						__lowerCamelCase					      =			pipe(
						    [prompt]							,    generator=snake_case_							,    guidance_scale=snake_case_							,    num_inference_steps=50							,    output_type='''np'''							,    width=5_12							,    height=5_12							,    **snake_case_							,    )

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

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

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






		def 				__A (	self	:							Optional[Any]						)  ->							int:
						__lowerCamelCase					      =			'''cuda'''
						__lowerCamelCase					      =			StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''						)
						__lowerCamelCase					      =			pipe.to(snake_case_						)
						pipe.set_progress_bar_config(disable=snake_case_						)

						__lowerCamelCase					      =			'''a castle next to a river'''
						__lowerCamelCase					      =			{
						    '''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,
						}

						__lowerCamelCase					      =			48
						__lowerCamelCase					      =			7

						# no sega enabled
						__lowerCamelCase					      =			torch.Generator(snake_case_						)
						generator.manual_seed(snake_case_						)
						__lowerCamelCase					      =			pipe(
						    [prompt]							,    generator=snake_case_							,    guidance_scale=snake_case_							,    num_inference_steps=50							,    output_type='''np'''							,    width=5_12							,    height=5_12							,    )

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

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

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

						# with sega enabled
						# generator = torch.manual_seed(seed)
						generator.manual_seed(snake_case_						)
						__lowerCamelCase					      =			pipe(
						    [prompt]							,    generator=snake_case_							,    guidance_scale=snake_case_							,    num_inference_steps=50							,    output_type='''np'''							,    width=5_12							,    height=5_12							,    **snake_case_							,    )

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

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

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







		def 				__A (	self	:							str						)  ->							Optional[int]:
						__lowerCamelCase					      =			'''cuda'''
						__lowerCamelCase					      =			StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''							,    torch_dtype=torch.floataa						)
						__lowerCamelCase					      =			pipe.to(snake_case_						)
						pipe.set_progress_bar_config(disable=snake_case_						)

						__lowerCamelCase					      =			'''a photo of a cat'''
						__lowerCamelCase					      =			{
						    '''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,
						}

						__lowerCamelCase					      =			3
						__lowerCamelCase					      =			7

						# no sega enabled
						__lowerCamelCase					      =			torch.Generator(snake_case_						)
						generator.manual_seed(snake_case_						)
						__lowerCamelCase					      =			pipe(
						    [prompt]							,    generator=snake_case_							,    guidance_scale=snake_case_							,    num_inference_steps=50							,    output_type='''np'''							,    width=5_12							,    height=5_12							,    )

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

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

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

						# with sega enabled
						# generator = torch.manual_seed(seed)
						generator.manual_seed(snake_case_						)
						__lowerCamelCase					      =			pipe(
						    [prompt]							,    generator=snake_case_							,    guidance_scale=snake_case_							,    num_inference_steps=50							,    output_type='''np'''							,    width=5_12							,    height=5_12							,    **snake_case_							,    )

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

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

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





from ..utils import requires_backends


SCREAMING_SNAKE_CASE__     :    Union[str, Any]      							=					None

class      lowerCAmelCase__  :






		def __init__(	self	:							int							,    *SCREAMING_SNAKE_CASE__	:							Dict							,    **SCREAMING_SNAKE_CASE__	:							Tuple						)  ->							Optional[Any]:
						requires_backends(self							,    ['''detectron2''']						)







		@classmethod
		def 				__A (	cls	:							Dict							,    *SCREAMING_SNAKE_CASE__	:							Dict							,    **SCREAMING_SNAKE_CASE__	:							Tuple						)  ->							Tuple:
						requires_backends(cls							,    ['''detectron2''']						)



