




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




									"""simple docstring"""


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







			@property
			def UpperCamelCase     (			self       :     Tuple ):




									"""simple docstring"""


									_UpperCAmelCase						=	1
									_UpperCAmelCase						=	3
									_UpperCAmelCase						=	(32, 32)

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







			@property
			def UpperCamelCase     (			self       :     str ):




									"""simple docstring"""


									torch.manual_seed(0 )
									_UpperCAmelCase						=	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 UpperCamelCase     (			self       :     Optional[int] ):




									"""simple docstring"""


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




									"""simple docstring"""


									torch.manual_seed(0 )
									_UpperCAmelCase						=	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(_A )







			@property
			def UpperCamelCase     (			self       :     List[Any] ):




									"""simple docstring"""


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






															class        __lowerCAmelCase					:







																		def __init__(			self       :     Dict ):




																								"""simple docstring"""


																								_UpperCAmelCase						=	torch.ones([0] )







																		def UpperCamelCase     (			self       :     Optional[int]       ,       snake_case__       :     Dict ):




																								"""simple docstring"""


																								self.pixel_values.to(_A )
																								return self




															return Out()

									return extract







			def UpperCamelCase     (			self       :     List[Any] ):




									"""simple docstring"""


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

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

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

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

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

									_UpperCAmelCase						=	torch.Generator(device=_A ).manual_seed(0 )
									_UpperCAmelCase						=	sd_pipe(
									    [prompt]       ,       generator=_A       ,       guidance_scale=6.0       ,       num_inference_steps=2       ,       output_type="np"       ,       return_dict=_A       ,       )[0]

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

									assert image.shape == (1, 64, 64, 3)
									_UpperCAmelCase						=	np.array([0.5_753, 0.6_114, 0.5_001, 0.5_034, 0.5_470, 0.4_729, 0.4_971, 0.4_867, 0.4_867] )

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




									"""simple docstring"""


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

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

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

									_UpperCAmelCase						=	output.images

									_UpperCAmelCase						=	torch.Generator(device=_A ).manual_seed(0 )
									_UpperCAmelCase						=	sd_pipe(
									    [prompt]       ,       generator=_A       ,       guidance_scale=6.0       ,       num_inference_steps=2       ,       output_type="np"       ,       return_dict=_A       ,       )[0]

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

									assert image.shape == (1, 64, 64, 3)
									_UpperCAmelCase						=	np.array([0.5_122, 0.5_712, 0.4_825, 0.5_053, 0.5_646, 0.4_769, 0.5_179, 0.4_894, 0.4_994] )

									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 UpperCamelCase     (			self       :     Dict ):




									"""simple docstring"""


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

									_UpperCAmelCase						=	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 )
															_UpperCAmelCase						=	StableDiffusionPipeline.from_pretrained(_A )

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




									"""simple docstring"""


									_UpperCAmelCase						=	self.dummy_cond_unet
									_UpperCAmelCase						=	PNDMScheduler(skip_prk_steps=_A )
									_UpperCAmelCase						=	self.dummy_vae
									_UpperCAmelCase						=	self.dummy_text_encoder
									_UpperCAmelCase						=	CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

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

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

									_UpperCAmelCase						=	'A painting of a squirrel eating a burger'
									_UpperCAmelCase						=	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 UpperCamelCase     (			self       :     Any ):




									"""simple docstring"""


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







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




									"""simple docstring"""


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

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

									_UpperCAmelCase						=	3
									_UpperCAmelCase						=	7

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

									_UpperCAmelCase						=	output.images
									_UpperCAmelCase						=	image[0, -3:, -3:, -1]
									_UpperCAmelCase						=	[
									    0.34_673_113,
									    0.38_492_733,
									    0.37_597_352,
									    0.34_086_335,
									    0.35_650_748,
									    0.35_579_205,
									    0.3_384_763,
									    0.34_340_236,
									    0.3_573_271,
									]

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

									_UpperCAmelCase						=	output.images
									_UpperCAmelCase						=	image[0, -3:, -3:, -1]
									_UpperCAmelCase						=	[
									    0.41_887_826,
									    0.37_728_766,
									    0.30_138_272,
									    0.41_416_335,
									    0.41_664_985,
									    0.36_283_392,
									    0.36_191_246,
									    0.43_364_465,
									    0.43_001_732,
									]

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

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







			def UpperCamelCase     (			self       :     int ):




									"""simple docstring"""


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

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

									_UpperCAmelCase						=	9
									_UpperCAmelCase						=	7

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

									_UpperCAmelCase						=	output.images
									_UpperCAmelCase						=	image[0, -3:, -3:, -1]
									_UpperCAmelCase						=	[
									    0.43_497_998,
									    0.91_814_065,
									    0.7_540_739,
									    0.55_580_205,
									    0.8_467_265,
									    0.5_389_691,
									    0.62_574_506,
									    0.58_897_763,
									    0.50_926_757,
									]

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

									_UpperCAmelCase						=	output.images
									_UpperCAmelCase						=	image[0, -3:, -3:, -1]
									_UpperCAmelCase						=	[
									    0.3_089_719,
									    0.30_500_144,
									    0.29_016_042,
									    0.30_630_964,
									    0.325_687,
									    0.29_419_225,
									    0.2_908_091,
									    0.28_723_598,
									    0.27_696_294,
									]

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

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







			def UpperCamelCase     (			self       :     int ):




									"""simple docstring"""


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

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

									_UpperCAmelCase						=	48
									_UpperCAmelCase						=	7

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

									_UpperCAmelCase						=	output.images
									_UpperCAmelCase						=	image[0, -3:, -3:, -1]
									_UpperCAmelCase						=	[
									    0.75_163_555,
									    0.76_037_145,
									    0.61_785,
									    0.9_189_673,
									    0.8_627_701,
									    0.85_189_694,
									    0.8_512_813,
									    0.87_012_076,
									    0.8_312_857,
									]

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

									_UpperCAmelCase						=	output.images
									_UpperCAmelCase						=	image[0, -3:, -3:, -1]
									_UpperCAmelCase						=	[
									    0.73_553_365,
									    0.7_537_271,
									    0.74_341_905,
									    0.66_480_356,
									    0.6_472_925,
									    0.63_039_416,
									    0.64_812_905,
									    0.6_749_717,
									    0.6_517_102,
									]

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

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







			def UpperCamelCase     (			self       :     List[Any] ):




									"""simple docstring"""


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

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

									_UpperCAmelCase						=	3
									_UpperCAmelCase						=	7

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

									_UpperCAmelCase						=	output.images
									_UpperCAmelCase						=	image[0, -3:, -3:, -1]
									_UpperCAmelCase						=	[
									    0.34_887_695,
									    0.3_876_953,
									    0.375,
									    0.34_423_828,
									    0.3_581_543,
									    0.35_717_773,
									    0.3_383_789,
									    0.34_570_312,
									    0.359_375,
									]

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

									_UpperCAmelCase						=	output.images
									_UpperCAmelCase						=	image[0, -3:, -3:, -1]
									_UpperCAmelCase						=	[
									    0.42_285_156,
									    0.36_914_062,
									    0.29_077_148,
									    0.42_041_016,
									    0.41_918_945,
									    0.35_498_047,
									    0.3_618_164,
									    0.4_423_828,
									    0.43_115_234,
									]

									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_flax_available,
    is_tf_available,
    is_torch_available,
)


_a	:	Optional[Any]=					{"configuration_gptj": ["GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPTJConfig", "GPTJOnnxConfig"]}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_a	:	Dict=					[
	    "GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST",
	    "GPTJForCausalLM",
	    "GPTJForQuestionAnswering",
	    "GPTJForSequenceClassification",
	    "GPTJModel",
	    "GPTJPreTrainedModel",
	]

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_a	:	Optional[Any]=					[
	    "TFGPTJForCausalLM",
	    "TFGPTJForQuestionAnswering",
	    "TFGPTJForSequenceClassification",
	    "TFGPTJModel",
	    "TFGPTJPreTrainedModel",
	]

try:
	if not is_flax_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_a	:	List[str]=					[
	    "FlaxGPTJForCausalLM",
	    "FlaxGPTJModel",
	    "FlaxGPTJPreTrainedModel",
	]


if TYPE_CHECKING:
	from .configuration_gptj import GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTJConfig, GPTJOnnxConfig

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_gptj import (
		    GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST,
		    GPTJForCausalLM,
		    GPTJForQuestionAnswering,
		    GPTJForSequenceClassification,
		    GPTJModel,
		    GPTJPreTrainedModel,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_gptj import (
		    TFGPTJForCausalLM,
		    TFGPTJForQuestionAnswering,
		    TFGPTJForSequenceClassification,
		    TFGPTJModel,
		    TFGPTJPreTrainedModel,
		)

	try:
		if not is_flax_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel, FlaxGPTJPreTrainedModel

else:
	import sys

	_a	:	Any=					_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





