




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




					'''simple docstring'''






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






					@property
					def        __SCREAMING_SNAKE_CASE			(					self				):
							_snake_case			 = 1
							_snake_case			 = 3
							_snake_case			 = (32, 32)

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






					@property
					def        __SCREAMING_SNAKE_CASE			(					self				):
							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        __SCREAMING_SNAKE_CASE			(					self				):
							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        __SCREAMING_SNAKE_CASE			(					self				):
							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(__UpperCamelCase				)






					@property
					def        __SCREAMING_SNAKE_CASE			(					self				):
							def extract(*__UpperCamelCase,						**__UpperCamelCase				):





									class   lowerCamelCase_							:




														'''simple docstring'''






														def __init__(					self				):
																_snake_case			 = torch.ones([0]				)






														def        __SCREAMING_SNAKE_CASE			(					self,						__UpperCamelCase				):
																self.pixel_values.to(__UpperCamelCase				)
																return self




									return Out()

							return extract






					def        __SCREAMING_SNAKE_CASE			(					self				):
							_snake_case			 = "cpu"  # ensure determinism for the device-dependent torch.Generator
							_snake_case			 = self.dummy_cond_unet
							_snake_case			 = DDIMScheduler(
							    beta_start=0.0_00_85,						beta_end=0.0_12,						beta_schedule="scaled_linear",						clip_sample=__UpperCamelCase,						set_alpha_to_one=__UpperCamelCase,						)

							_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=__UpperCamelCase,						scheduler=__UpperCamelCase,						vae=__UpperCamelCase,						text_encoder=__UpperCamelCase,						tokenizer=__UpperCamelCase,						safety_checker=__UpperCamelCase,						feature_extractor=self.dummy_extractor,						)
							_snake_case			 = sd_pipe.to(__UpperCamelCase				)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase				)

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

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

							_snake_case			 = torch.Generator(device=__UpperCamelCase				).manual_seed(0				)
							_snake_case			 = sd_pipe(
							    [prompt],						generator=__UpperCamelCase,						guidance_scale=6.0,						num_inference_steps=2,						output_type="np",						return_dict=__UpperCamelCase,						)[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.57_53, 0.61_14, 0.50_01, 0.50_34, 0.54_70, 0.47_29, 0.49_71, 0.48_67, 0.48_67]				)

							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				):
							_snake_case			 = "cpu"  # ensure determinism for the device-dependent torch.Generator
							_snake_case			 = self.dummy_cond_unet
							_snake_case			 = PNDMScheduler(skip_prk_steps=__UpperCamelCase				)
							_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=__UpperCamelCase,						scheduler=__UpperCamelCase,						vae=__UpperCamelCase,						text_encoder=__UpperCamelCase,						tokenizer=__UpperCamelCase,						safety_checker=__UpperCamelCase,						feature_extractor=self.dummy_extractor,						)
							_snake_case			 = sd_pipe.to(__UpperCamelCase				)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase				)

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

							_snake_case			 = output.images

							_snake_case			 = torch.Generator(device=__UpperCamelCase				).manual_seed(0				)
							_snake_case			 = sd_pipe(
							    [prompt],						generator=__UpperCamelCase,						guidance_scale=6.0,						num_inference_steps=2,						output_type="np",						return_dict=__UpperCamelCase,						)[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.51_22, 0.57_12, 0.48_25, 0.50_53, 0.56_46, 0.47_69, 0.51_79, 0.48_94, 0.49_94]				)

							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				):
							_snake_case			 = StableDiffusionPipeline.from_pretrained(
							    "hf-internal-testing/tiny-stable-diffusion-lms-pipe",						safety_checker=__UpperCamelCase				)
							assert isinstance(__UpperCamelCase,						__UpperCamelCase				)
							assert isinstance(pipe.scheduler,						__UpperCamelCase				)
							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(__UpperCamelCase				)
									_snake_case			 = StableDiffusionPipeline.from_pretrained(__UpperCamelCase				)

							# 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        __SCREAMING_SNAKE_CASE			(					self				):
							_snake_case			 = self.dummy_cond_unet
							_snake_case			 = PNDMScheduler(skip_prk_steps=__UpperCamelCase				)
							_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=__UpperCamelCase,						scheduler=__UpperCamelCase,						vae=__UpperCamelCase,						text_encoder=__UpperCamelCase,						tokenizer=__UpperCamelCase,						safety_checker=__UpperCamelCase,						feature_extractor=self.dummy_extractor,						)
							_snake_case			 = sd_pipe.to(__UpperCamelCase				)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase				)

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




					'''simple docstring'''






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






					def        __SCREAMING_SNAKE_CASE			(					self				):
							_snake_case			 = "cuda"
							_snake_case			 = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"				)
							_snake_case			 = pipe.to(__UpperCamelCase				)
							pipe.set_progress_bar_config(disable=__UpperCamelCase				)

							_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(__UpperCamelCase				)
							generator.manual_seed(__UpperCamelCase				)
							_snake_case			 = pipe(
							    [prompt],						generator=__UpperCamelCase,						guidance_scale=__UpperCamelCase,						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.34_67_31_13,
							    0.38_49_27_33,
							    0.37_59_73_52,
							    0.34_08_63_35,
							    0.35_65_07_48,
							    0.35_57_92_05,
							    0.3_38_47_63,
							    0.34_34_02_36,
							    0.3_57_32_71,
							]

							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(__UpperCamelCase				)
							_snake_case			 = pipe(
							    [prompt],						generator=__UpperCamelCase,						guidance_scale=__UpperCamelCase,						num_inference_steps=50,						output_type="np",						width=512,						height=512,						**__UpperCamelCase,						)

							_snake_case			 = output.images
							_snake_case			 = image[0, -3:, -3:, -1]
							_snake_case			 = [
							    0.41_88_78_26,
							    0.37_72_87_66,
							    0.30_13_82_72,
							    0.41_41_63_35,
							    0.41_66_49_85,
							    0.36_28_33_92,
							    0.36_19_12_46,
							    0.43_36_44_65,
							    0.43_00_17_32,
							]

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

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






					def        __SCREAMING_SNAKE_CASE			(					self				):
							_snake_case			 = "cuda"
							_snake_case			 = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"				)
							_snake_case			 = pipe.to(__UpperCamelCase				)
							pipe.set_progress_bar_config(disable=__UpperCamelCase				)

							_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(__UpperCamelCase				)
							generator.manual_seed(__UpperCamelCase				)
							_snake_case			 = pipe(
							    [prompt],						generator=__UpperCamelCase,						guidance_scale=__UpperCamelCase,						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.43_49_79_98,
							    0.91_81_40_65,
							    0.7_54_07_39,
							    0.55_58_02_05,
							    0.8_46_72_65,
							    0.5_38_96_91,
							    0.62_57_45_06,
							    0.58_89_77_63,
							    0.50_92_67_57,
							]

							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(__UpperCamelCase				)
							_snake_case			 = pipe(
							    [prompt],						generator=__UpperCamelCase,						guidance_scale=__UpperCamelCase,						num_inference_steps=50,						output_type="np",						width=512,						height=512,						**__UpperCamelCase,						)

							_snake_case			 = output.images
							_snake_case			 = image[0, -3:, -3:, -1]
							_snake_case			 = [
							    0.3_08_97_19,
							    0.30_50_01_44,
							    0.29_01_60_42,
							    0.30_63_09_64,
							    0.32_56_87,
							    0.29_41_92_25,
							    0.2_90_80_91,
							    0.28_72_35_98,
							    0.27_69_62_94,
							]

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

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






					def        __SCREAMING_SNAKE_CASE			(					self				):
							_snake_case			 = "cuda"
							_snake_case			 = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"				)
							_snake_case			 = pipe.to(__UpperCamelCase				)
							pipe.set_progress_bar_config(disable=__UpperCamelCase				)

							_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(__UpperCamelCase				)
							generator.manual_seed(__UpperCamelCase				)
							_snake_case			 = pipe(
							    [prompt],						generator=__UpperCamelCase,						guidance_scale=__UpperCamelCase,						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.75_16_35_55,
							    0.76_03_71_45,
							    0.6_17_85,
							    0.9_18_96_73,
							    0.8_62_77_01,
							    0.85_18_96_94,
							    0.8_51_28_13,
							    0.87_01_20_76,
							    0.8_31_28_57,
							]

							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(__UpperCamelCase				)
							_snake_case			 = pipe(
							    [prompt],						generator=__UpperCamelCase,						guidance_scale=__UpperCamelCase,						num_inference_steps=50,						output_type="np",						width=512,						height=512,						**__UpperCamelCase,						)

							_snake_case			 = output.images
							_snake_case			 = image[0, -3:, -3:, -1]
							_snake_case			 = [
							    0.73_55_33_65,
							    0.7_53_72_71,
							    0.74_34_19_05,
							    0.66_48_03_56,
							    0.6_47_29_25,
							    0.63_03_94_16,
							    0.64_81_29_05,
							    0.6_74_97_17,
							    0.6_51_71_02,
							]

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

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






					def        __SCREAMING_SNAKE_CASE			(					self				):
							_snake_case			 = "cuda"
							_snake_case			 = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5",						torch_dtype=torch.floataa				)
							_snake_case			 = pipe.to(__UpperCamelCase				)
							pipe.set_progress_bar_config(disable=__UpperCamelCase				)

							_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(__UpperCamelCase				)
							generator.manual_seed(__UpperCamelCase				)
							_snake_case			 = pipe(
							    [prompt],						generator=__UpperCamelCase,						guidance_scale=__UpperCamelCase,						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.34_88_76_95,
							    0.3_87_69_53,
							    0.3_75,
							    0.34_42_38_28,
							    0.3_58_15_43,
							    0.35_71_77_73,
							    0.3_38_37_89,
							    0.34_57_03_12,
							    0.35_93_75,
							]

							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(__UpperCamelCase				)
							_snake_case			 = pipe(
							    [prompt],						generator=__UpperCamelCase,						guidance_scale=__UpperCamelCase,						num_inference_steps=50,						output_type="np",						width=512,						height=512,						**__UpperCamelCase,						)

							_snake_case			 = output.images
							_snake_case			 = image[0, -3:, -3:, -1]
							_snake_case			 = [
							    0.42_28_51_56,
							    0.36_91_40_62,
							    0.29_07_71_48,
							    0.42_04_10_16,
							    0.41_91_89_45,
							    0.35_49_80_47,
							    0.3_61_81_64,
							    0.4_42_38_28,
							    0.43_11_52_34,
							]

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

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











'''simple docstring'''







from __future__ import annotations

from collections.abc import Iterable, Iterator
from typing import Any, Generic, TypeVar

_lowerCamelCase							:		List[str]										=						TypeVar('T', bound=bool)





class   lowerCamelCase_							( Generic[T]							):




					'''simple docstring'''






					def __init__(					self,						__UpperCamelCase				):
							_snake_case			 = value
							_snake_case			 = None
							_snake_case			 = None






					@property
					def        __SCREAMING_SNAKE_CASE			(					self				):
							return self._value






					@staticmethod
					def        __SCREAMING_SNAKE_CASE			(					__UpperCamelCase,						__UpperCamelCase				):
							if not roota:
									return roota

							if not roota:
									return roota

							if roota.value > roota.value:
									_snake_case ,					_snake_case			 = roota, roota

							_snake_case			 = roota
							_snake_case			 = roota.right
							_snake_case			 = roota.left
							_snake_case			 = SkewNode.merge(__UpperCamelCase,						__UpperCamelCase				)

							return result





class   lowerCamelCase_							( Generic[T]							):




					'''simple docstring'''






					def __init__(					self,						__UpperCamelCase = ()				):
							_snake_case			 = None
							if data:
									for item in data:
											self.insert(__UpperCamelCase				)






					def __bool__(					self				):
							return self._root is not None






					def __iter__(					self				):
							_snake_case			 = []
							while self:
									result.append(self.pop()				)

							# Pushing items back to the heap not to clear it.
							for item in result:
									self.insert(__UpperCamelCase				)

							return iter(__UpperCamelCase				)






					def        __SCREAMING_SNAKE_CASE			(					self,						__UpperCamelCase				):
							_snake_case			 = SkewNode.merge(self._root,						SkewNode(__UpperCamelCase				)				)






					def        __SCREAMING_SNAKE_CASE			(					self				):
							_snake_case			 = self.top()
							_snake_case			 = (
							    SkewNode.merge(self._root.left,						self._root.right				) if self._root else None
							)

							return result






					def        __SCREAMING_SNAKE_CASE			(					self				):
							if not self._root:
									raise IndexError("Can't get top element for the empty heap."				)
							return self._root.value






					def        __SCREAMING_SNAKE_CASE			(					self				):
							_snake_case			 = None


if __name__ == "__main__":
						import doctest

						doctest.testmod()





