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







			def       SCREAMING_SNAKE_CASE_	(      self							:    Tuple       ):
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()







			@property
			def       SCREAMING_SNAKE_CASE_	(      self							:    List[str]       ):
										lowercase_   :     Any									=			1
										lowercase_   :     Dict									=			3
										lowercase_   :     Tuple									=			(32, 32)

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







			@property
			def       SCREAMING_SNAKE_CASE_	(      self							:    Optional[int]       ):
										torch.manual_seed(0       )
										lowercase_   :     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       SCREAMING_SNAKE_CASE_	(      self							:    Dict       ):
										torch.manual_seed(0       )
										lowercase_   :     Optional[int]									=			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							:    List[Any]       ):
										torch.manual_seed(0       )
										lowercase_   :     Optional[Any]									=			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=1000							,     )
										return CLIPTextModel(a_       )







			@property
			def       SCREAMING_SNAKE_CASE_	(      self							:    Dict       ):



										def extract(*lowercase_							:    List[Any]							,     **lowercase_							:    str       ):





																	class        __magic_name__		:







																				def __init__(      self							:    Dict       ):
																											lowercase_   :     Union[str, Any]									=			torch.ones([0]       )







																				def       SCREAMING_SNAKE_CASE_	(      self							:    Optional[int]							,     lowercase_							:    str       ):
																											self.pixel_values.to(a_       )
																											return self




																	return Out()

										return extract







			def       SCREAMING_SNAKE_CASE_	(      self							:    int       ):
										lowercase_   :     List[str]									=			"""cpu"""  # ensure determinism for the device-dependent torch.Generator
										lowercase_   :     Optional[Any]									=			self.dummy_cond_unet
										lowercase_   :     Any									=			DDIMScheduler(
										    beta_start=0.0_00_85							,     beta_end=0.0_12							,     beta_schedule="""scaled_linear"""							,     clip_sample=a_							,     set_alpha_to_one=a_							,     )

										lowercase_   :     str									=			self.dummy_vae
										lowercase_   :     Tuple									=			self.dummy_text_encoder
										lowercase_   :     str									=			CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       )

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

										lowercase_   :     str									=			"""A painting of a squirrel eating a burger"""

										lowercase_   :     int									=			torch.Generator(device=a_       ).manual_seed(0       )
										lowercase_   :     str									=			sd_pipe([prompt]							,     generator=a_							,     guidance_scale=6.0							,     num_inference_steps=2							,     output_type="""np"""       )
										lowercase_   :     List[str]									=			output.images

										lowercase_   :     Dict									=			torch.Generator(device=a_       ).manual_seed(0       )
										lowercase_   :     Any									=			sd_pipe(
										    [prompt]							,     generator=a_							,     guidance_scale=6.0							,     num_inference_steps=2							,     output_type="""np"""							,     return_dict=a_							,     )[0]

										lowercase_   :     Tuple									=			image[0, -3:, -3:, -1]
										lowercase_   :     str									=			image_from_tuple[0, -3:, -3:, -1]

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

										# make sure here that pndm scheduler skips prk
										lowercase_   :     List[str]									=			StableDiffusionPipeline(
										    unet=a_							,     scheduler=a_							,     vae=a_							,     text_encoder=a_							,     tokenizer=a_							,     safety_checker=a_							,     feature_extractor=self.dummy_extractor							,     )
										lowercase_   :     Tuple									=			sd_pipe.to(a_       )
										sd_pipe.set_progress_bar_config(disable=a_       )

										lowercase_   :     str									=			"""A painting of a squirrel eating a burger"""
										lowercase_   :     List[str]									=			torch.Generator(device=a_       ).manual_seed(0       )
										lowercase_   :     Dict									=			sd_pipe([prompt]							,     generator=a_							,     guidance_scale=6.0							,     num_inference_steps=2							,     output_type="""np"""       )

										lowercase_   :     Dict									=			output.images

										lowercase_   :     Any									=			torch.Generator(device=a_       ).manual_seed(0       )
										lowercase_   :     Dict									=			sd_pipe(
										    [prompt]							,     generator=a_							,     guidance_scale=6.0							,     num_inference_steps=2							,     output_type="""np"""							,     return_dict=a_							,     )[0]

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

										assert image.shape == (1, 64, 64, 3)
										lowercase_   :     List[Any]									=			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							:    Tuple       ):
										lowercase_   :     List[str]									=			StableDiffusionPipeline.from_pretrained(
										    """hf-internal-testing/tiny-stable-diffusion-lms-pipe"""							,     safety_checker=a_       )
										assert isinstance(a_							,     a_       )
										assert isinstance(pipe.scheduler							,     a_       )
										assert pipe.safety_checker is None

										lowercase_   :     str									=			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(a_       )
																	lowercase_   :     List[Any]									=			StableDiffusionPipeline.from_pretrained(a_       )

										# sanity check that the pipeline still works
										assert pipe.safety_checker is None
										lowercase_   :     List[str]									=			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							:    int       ):
										lowercase_   :     str									=			self.dummy_cond_unet
										lowercase_   :     List[Any]									=			PNDMScheduler(skip_prk_steps=a_       )
										lowercase_   :     Union[str, Any]									=			self.dummy_vae
										lowercase_   :     Dict									=			self.dummy_text_encoder
										lowercase_   :     Optional[Any]									=			CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       )

										# put models in fp16
										lowercase_   :     int									=			unet.half()
										lowercase_   :     Optional[Any]									=			vae.half()
										lowercase_   :     Dict									=			bert.half()

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

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

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





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







			def       SCREAMING_SNAKE_CASE_	(      self							:    List[Any]       ):
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()







			def       SCREAMING_SNAKE_CASE_	(      self							:    Dict       ):
										lowercase_   :     List[str]									=			"""cuda"""
										lowercase_   :     Any									=			StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""       )
										lowercase_   :     List[str]									=			pipe.to(a_       )
										pipe.set_progress_bar_config(disable=a_       )

										lowercase_   :     Optional[Any]									=			"""a photo of a cat"""
										lowercase_   :     Tuple									=			{
										    """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,
										}

										lowercase_   :     List[Any]									=			3
										lowercase_   :     Optional[int]									=			7

										# no sega enabled
										lowercase_   :     Optional[Any]									=			torch.Generator(a_       )
										generator.manual_seed(a_       )
										lowercase_   :     Dict									=			pipe(
										    [prompt]							,     generator=a_							,     guidance_scale=a_							,     num_inference_steps=50							,     output_type="""np"""							,     width=512							,     height=512							,     )

										lowercase_   :     str									=			output.images
										lowercase_   :     Union[str, Any]									=			image[0, -3:, -3:, -1]
										lowercase_   :     List[str]									=			[
										    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(a_       )
										lowercase_   :     List[str]									=			pipe(
										    [prompt]							,     generator=a_							,     guidance_scale=a_							,     num_inference_steps=50							,     output_type="""np"""							,     width=512							,     height=512							,     **a_							,     )

										lowercase_   :     List[str]									=			output.images
										lowercase_   :     Tuple									=			image[0, -3:, -3:, -1]
										lowercase_   :     Tuple									=			[
										    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							:    int       ):
										lowercase_   :     Optional[int]									=			"""cuda"""
										lowercase_   :     Tuple									=			StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""       )
										lowercase_   :     Any									=			pipe.to(a_       )
										pipe.set_progress_bar_config(disable=a_       )

										lowercase_   :     List[Any]									=			"""an image of a crowded boulevard, realistic, 4k"""
										lowercase_   :     Union[str, Any]									=			{
										    """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,
										}

										lowercase_   :     Optional[int]									=			9
										lowercase_   :     Tuple									=			7

										# no sega enabled
										lowercase_   :     List[Any]									=			torch.Generator(a_       )
										generator.manual_seed(a_       )
										lowercase_   :     Union[str, Any]									=			pipe(
										    [prompt]							,     generator=a_							,     guidance_scale=a_							,     num_inference_steps=50							,     output_type="""np"""							,     width=512							,     height=512							,     )

										lowercase_   :     Union[str, Any]									=			output.images
										lowercase_   :     Optional[int]									=			image[0, -3:, -3:, -1]
										lowercase_   :     Any									=			[
										    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(a_       )
										lowercase_   :     Optional[int]									=			pipe(
										    [prompt]							,     generator=a_							,     guidance_scale=a_							,     num_inference_steps=50							,     output_type="""np"""							,     width=512							,     height=512							,     **a_							,     )

										lowercase_   :     str									=			output.images
										lowercase_   :     int									=			image[0, -3:, -3:, -1]
										lowercase_   :     List[Any]									=			[
										    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							:    Dict       ):
										lowercase_   :     Tuple									=			"""cuda"""
										lowercase_   :     Dict									=			StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""       )
										lowercase_   :     List[Any]									=			pipe.to(a_       )
										pipe.set_progress_bar_config(disable=a_       )

										lowercase_   :     Tuple									=			"""a castle next to a river"""
										lowercase_   :     List[str]									=			{
										    """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,
										}

										lowercase_   :     Dict									=			48
										lowercase_   :     Dict									=			7

										# no sega enabled
										lowercase_   :     int									=			torch.Generator(a_       )
										generator.manual_seed(a_       )
										lowercase_   :     Union[str, Any]									=			pipe(
										    [prompt]							,     generator=a_							,     guidance_scale=a_							,     num_inference_steps=50							,     output_type="""np"""							,     width=512							,     height=512							,     )

										lowercase_   :     Optional[Any]									=			output.images
										lowercase_   :     int									=			image[0, -3:, -3:, -1]
										lowercase_   :     Union[str, Any]									=			[
										    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(a_       )
										lowercase_   :     Tuple									=			pipe(
										    [prompt]							,     generator=a_							,     guidance_scale=a_							,     num_inference_steps=50							,     output_type="""np"""							,     width=512							,     height=512							,     **a_							,     )

										lowercase_   :     Dict									=			output.images
										lowercase_   :     List[str]									=			image[0, -3:, -3:, -1]
										lowercase_   :     List[Any]									=			[
										    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							:    int       ):
										lowercase_   :     List[str]									=			"""cuda"""
										lowercase_   :     int									=			StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""							,     torch_dtype=torch.floataa       )
										lowercase_   :     Tuple									=			pipe.to(a_       )
										pipe.set_progress_bar_config(disable=a_       )

										lowercase_   :     Any									=			"""a photo of a cat"""
										lowercase_   :     Tuple									=			{
										    """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,
										}

										lowercase_   :     List[Any]									=			3
										lowercase_   :     int									=			7

										# no sega enabled
										lowercase_   :     Optional[int]									=			torch.Generator(a_       )
										generator.manual_seed(a_       )
										lowercase_   :     str									=			pipe(
										    [prompt]							,     generator=a_							,     guidance_scale=a_							,     num_inference_steps=50							,     output_type="""np"""							,     width=512							,     height=512							,     )

										lowercase_   :     Any									=			output.images
										lowercase_   :     int									=			image[0, -3:, -3:, -1]
										lowercase_   :     Any									=			[
										    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(a_       )
										lowercase_   :     Dict									=			pipe(
										    [prompt]							,     generator=a_							,     guidance_scale=a_							,     num_inference_steps=50							,     output_type="""np"""							,     width=512							,     height=512							,     **a_							,     )

										lowercase_   :     List[str]									=			output.images
										lowercase_   :     Dict									=			image[0, -3:, -3:, -1]
										lowercase_   :     int									=			[
										    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 typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


__lowercase:		Tuple    						=  {
    "configuration_imagegpt": ["IMAGEGPT_PRETRAINED_CONFIG_ARCHIVE_MAP", "ImageGPTConfig", "ImageGPTOnnxConfig"]
}

try:
       if not is_vision_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __lowercase:		List[str]    						=  ["ImageGPTFeatureExtractor"]
       __lowercase:		Tuple    						=  ["ImageGPTImageProcessor"]

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __lowercase:		str    						=  [
           "IMAGEGPT_PRETRAINED_MODEL_ARCHIVE_LIST",
           "ImageGPTForCausalImageModeling",
           "ImageGPTForImageClassification",
           "ImageGPTModel",
           "ImageGPTPreTrainedModel",
           "load_tf_weights_in_imagegpt",
       ]


if TYPE_CHECKING:
       from .configuration_imagegpt import IMAGEGPT_PRETRAINED_CONFIG_ARCHIVE_MAP, ImageGPTConfig, ImageGPTOnnxConfig

       try:
              if not is_vision_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .feature_extraction_imagegpt import ImageGPTFeatureExtractor
              from .image_processing_imagegpt import ImageGPTImageProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_imagegpt import (
                  IMAGEGPT_PRETRAINED_MODEL_ARCHIVE_LIST,
                  ImageGPTForCausalImageModeling,
                  ImageGPTForImageClassification,
                  ImageGPTModel,
                  ImageGPTPreTrainedModel,
                  load_tf_weights_in_imagegpt,
              )

else:
       import sys

       __lowercase:		List[str]    						=  _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)