

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







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

						'''simple docstring'''





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







	@property
	def 						_SCREAMING_SNAKE_CASE				( self:			int   ):

						'''simple docstring'''





						__magic_name__				=     1
						__magic_name__				=     3
						__magic_name__				=     (32, 32)

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







	@property
	def 						_SCREAMING_SNAKE_CASE				( self:			Union[str, Any]   ):

						'''simple docstring'''





						torch.manual_seed(0   )
						__magic_name__				=     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:			List[Any]   ):

						'''simple docstring'''





						torch.manual_seed(0   )
						__magic_name__				=     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:			int   ):

						'''simple docstring'''





						torch.manual_seed(0   )
						__magic_name__				=     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(UpperCamelCase__   )







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

						'''simple docstring'''





						def extract(*__UpperCamelCase:			int	,			**__UpperCamelCase:			List[Any]   ):

											class 					__UpperCamelCase					:







												def __init__( self:			List[str]   ):

																	'''simple docstring'''





																	__magic_name__				=     torch.ones([0]   )







												def 						_SCREAMING_SNAKE_CASE				( self:			List[str]	,			__UpperCamelCase:			List[str]   ):

																	'''simple docstring'''





																	self.pixel_values.to(UpperCamelCase__   )
																	return self




											return Out()

						return extract







	def 						_SCREAMING_SNAKE_CASE				( self:			Any   ):

						'''simple docstring'''





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

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

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

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

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

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

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

						assert image.shape == (1, 64, 64, 3)
						__magic_name__				=     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 						_SCREAMING_SNAKE_CASE				( self:			Union[str, Any]   ):

						'''simple docstring'''





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

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

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

						__magic_name__				=     output.images

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

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

						assert image.shape == (1, 64, 64, 3)
						__magic_name__				=     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 						_SCREAMING_SNAKE_CASE				( self:			List[Any]   ):

						'''simple docstring'''





						__magic_name__				=     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

						__magic_name__				=     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__   )
											__magic_name__				=     StableDiffusionPipeline.from_pretrained(UpperCamelCase__   )

						# sanity check that the pipeline still works
						assert pipe.safety_checker is None
						__magic_name__				=     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:			Any   ):

						'''simple docstring'''





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

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

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

						__magic_name__				=     """A painting of a squirrel eating a burger"""
						__magic_name__				=     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      ):







	def 						_SCREAMING_SNAKE_CASE				( self:			List[str]   ):

						'''simple docstring'''





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







	def 						_SCREAMING_SNAKE_CASE				( self:			Any   ):

						'''simple docstring'''





						__magic_name__				=     """cuda"""
						__magic_name__				=     StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'   )
						__magic_name__				=     pipe.to(UpperCamelCase__   )
						pipe.set_progress_bar_config(disable=UpperCamelCase__   )

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

						__magic_name__				=     3
						__magic_name__				=     7

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

						__magic_name__				=     output.images
						__magic_name__				=     image[0, -3:, -3:, -1]
						__magic_name__				=     [
						    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(UpperCamelCase__   )
						__magic_name__				=     pipe(
						    [prompt]	,			generator=UpperCamelCase__	,			guidance_scale=UpperCamelCase__	,			num_inference_steps=50	,			output_type='np'	,			width=5_12	,			height=5_12	,			**UpperCamelCase__	,			)

						__magic_name__				=     output.images
						__magic_name__				=     image[0, -3:, -3:, -1]
						__magic_name__				=     [
						    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 						_SCREAMING_SNAKE_CASE				( self:			Any   ):

						'''simple docstring'''





						__magic_name__				=     """cuda"""
						__magic_name__				=     StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'   )
						__magic_name__				=     pipe.to(UpperCamelCase__   )
						pipe.set_progress_bar_config(disable=UpperCamelCase__   )

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

						__magic_name__				=     9
						__magic_name__				=     7

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

						__magic_name__				=     output.images
						__magic_name__				=     image[0, -3:, -3:, -1]
						__magic_name__				=     [
						    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(UpperCamelCase__   )
						__magic_name__				=     pipe(
						    [prompt]	,			generator=UpperCamelCase__	,			guidance_scale=UpperCamelCase__	,			num_inference_steps=50	,			output_type='np'	,			width=5_12	,			height=5_12	,			**UpperCamelCase__	,			)

						__magic_name__				=     output.images
						__magic_name__				=     image[0, -3:, -3:, -1]
						__magic_name__				=     [
						    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 						_SCREAMING_SNAKE_CASE				( self:			List[Any]   ):

						'''simple docstring'''





						__magic_name__				=     """cuda"""
						__magic_name__				=     StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'   )
						__magic_name__				=     pipe.to(UpperCamelCase__   )
						pipe.set_progress_bar_config(disable=UpperCamelCase__   )

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

						__magic_name__				=     48
						__magic_name__				=     7

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

						__magic_name__				=     output.images
						__magic_name__				=     image[0, -3:, -3:, -1]
						__magic_name__				=     [
						    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(UpperCamelCase__   )
						__magic_name__				=     pipe(
						    [prompt]	,			generator=UpperCamelCase__	,			guidance_scale=UpperCamelCase__	,			num_inference_steps=50	,			output_type='np'	,			width=5_12	,			height=5_12	,			**UpperCamelCase__	,			)

						__magic_name__				=     output.images
						__magic_name__				=     image[0, -3:, -3:, -1]
						__magic_name__				=     [
						    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 						_SCREAMING_SNAKE_CASE				( self:			int   ):

						'''simple docstring'''





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

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

						__magic_name__				=     3
						__magic_name__				=     7

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

						__magic_name__				=     output.images
						__magic_name__				=     image[0, -3:, -3:, -1]
						__magic_name__				=     [
						    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(UpperCamelCase__   )
						__magic_name__				=     pipe(
						    [prompt]	,			generator=UpperCamelCase__	,			guidance_scale=UpperCamelCase__	,			num_inference_steps=50	,			output_type='np'	,			width=5_12	,			height=5_12	,			**UpperCamelCase__	,			)

						__magic_name__				=     output.images
						__magic_name__				=     image[0, -3:, -3:, -1]
						__magic_name__				=     [
						    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












'''simple docstring'''







from typing import TYPE_CHECKING

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


_A      :     Tuple       	={
    '''configuration_imagegpt''': ['''IMAGEGPT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ImageGPTConfig''', '''ImageGPTOnnxConfig''']
}

try:
			if not is_vision_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_A      :     Dict       	=['''ImageGPTFeatureExtractor''']
			_A      :     Union[str, Any]       	=['''ImageGPTImageProcessor''']

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_A      :     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

			_A      :     str       	=_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)

