"""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.stable_diffusion_safe import StableDiffusionPipelineSafe as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import require_torch_gpu







class      UpperCAmelCase_       (						unittest.TestCase):






							def       _UpperCamelCase			(		self :      Dict  )       ->       int:
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()






							@property
							def       _UpperCamelCase			(		self :      str  )       ->       List[str]:
												_UpperCamelCase      						=	1
												_UpperCamelCase      						=	3
												_UpperCamelCase      						=	(32, 32)

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






							@property
							def       _UpperCamelCase			(		self :      List[Any]  )       ->       Dict:
												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 :      str  )       ->       Union[str, Any]:
												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 :      Optional[int]  )       ->       Any:
												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=1000     ,						)
												return CLIPTextModel(__UpperCamelCase  )






							@property
							def       _UpperCamelCase			(		self :      str  )       ->       Dict:
												def extract(*__UpperCamelCase :      str     ,						**__UpperCamelCase :      Dict  ):







																	class      UpperCAmelCase_       :






																								def __init__(		self :      List[str]  )       ->       Optional[Any]:
																													_UpperCamelCase      						=	torch.ones([0]  )






																								def       _UpperCamelCase			(		self :      List[Any]     ,						__UpperCamelCase :      Tuple  )       ->       int:
																													self.pixel_values.to(__UpperCamelCase  )
																													return self


																	return Out()

												return extract






							def       _UpperCamelCase			(		self :      Optional[int]  )       ->       str:
												_UpperCamelCase      						=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
												_UpperCamelCase      						=	self.dummy_cond_unet
												_UpperCamelCase      						=	DDIMScheduler(
												    beta_start=0.0_0_0_8_5     ,						beta_end=0.0_1_2     ,						beta_schedule='''scaled_linear'''     ,						clip_sample=__UpperCamelCase     ,						set_alpha_to_one=__UpperCamelCase     ,						)

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

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

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

												_UpperCamelCase      						=	torch.Generator(device=__UpperCamelCase  ).manual_seed(0  )
												_UpperCamelCase      						=	sd_pipe(
												    [prompt]     ,						generator=__UpperCamelCase     ,						guidance_scale=6.0     ,						num_inference_steps=2     ,						output_type='''np'''     ,						return_dict=__UpperCamelCase     ,						)[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_7_5_6, 0.6_1_1_8, 0.5_0_0_5, 0.5_0_4_1, 0.5_4_7_1, 0.4_7_2_6, 0.4_9_7_6, 0.4_8_6_5, 0.4_8_6_4]  )

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

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

												_UpperCamelCase      						=	output.images

												_UpperCamelCase      						=	torch.Generator(device=__UpperCamelCase  ).manual_seed(0  )
												_UpperCamelCase      						=	sd_pipe(
												    [prompt]     ,						generator=__UpperCamelCase     ,						guidance_scale=6.0     ,						num_inference_steps=2     ,						output_type='''np'''     ,						return_dict=__UpperCamelCase     ,						)[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_1_2_5, 0.5_7_1_6, 0.4_8_2_8, 0.5_0_6_0, 0.5_6_5_0, 0.4_7_6_8, 0.5_1_8_5, 0.4_8_9_5, 0.4_9_9_3]  )

												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  )       ->       Optional[int]:
												_UpperCamelCase      						=	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

												_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(__UpperCamelCase  )
																	_UpperCamelCase      						=	StableDiffusionPipeline.from_pretrained(__UpperCamelCase  )

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

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






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






							def       _UpperCamelCase			(		self :      List[Any]  )       ->       int:
												_UpperCamelCase      						=	StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''     ,						safety_checker=__UpperCamelCase  )
												_UpperCamelCase      						=	LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config  )
												_UpperCamelCase      						=	sd_pipe.to(__UpperCamelCase  )
												sd_pipe.set_progress_bar_config(disable=__UpperCamelCase  )

												_UpperCamelCase      						=	(
												    """portrait of girl with smokey eyes makeup in abandoned hotel, grange clothes, redshift, wide high angle"""
												    """ coloured polaroid photograph with flash, kodak film, hyper real, stunning moody cinematography, with"""
												    """ anamorphic lenses, by maripol, fallen angels by wong kar - wai, style of suspiria and neon demon and"""
												    """ children from bahnhof zoo, detailed """
												)
												_UpperCamelCase      						=	40_0366_0346
												_UpperCamelCase      						=	7

												# without safety guidance (sld_guidance_scale = 0)
												_UpperCamelCase      						=	torch.manual_seed(__UpperCamelCase  )
												_UpperCamelCase      						=	sd_pipe(
												    [prompt]     ,						generator=__UpperCamelCase     ,						guidance_scale=__UpperCamelCase     ,						num_inference_steps=50     ,						output_type='''np'''     ,						width=512     ,						height=512     ,						sld_guidance_scale=0     ,						)

												_UpperCamelCase      						=	output.images
												_UpperCamelCase      						=	image[0, -3:, -3:, -1]
												_UpperCamelCase      						=	[0.2_2_7_8, 0.2_2_3_1, 0.2_2_4_9, 0.2_3_3_3, 0.2_3_0_3, 0.1_8_8_5, 0.2_2_7_3, 0.2_1_4_4, 0.2_1_7_6]

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

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

												# without safety guidance (strong configuration)
												_UpperCamelCase      						=	torch.manual_seed(__UpperCamelCase  )
												_UpperCamelCase      						=	sd_pipe(
												    [prompt]     ,						generator=__UpperCamelCase     ,						guidance_scale=__UpperCamelCase     ,						num_inference_steps=50     ,						output_type='''np'''     ,						width=512     ,						height=512     ,						sld_guidance_scale=2000     ,						sld_warmup_steps=7     ,						sld_threshold=0.0_2_5     ,						sld_momentum_scale=0.5     ,						sld_mom_beta=0.7     ,						)

												_UpperCamelCase      						=	output.images
												_UpperCamelCase      						=	image[0, -3:, -3:, -1]
												_UpperCamelCase      						=	[0.2_3_8_3, 0.2_2_7_6, 0.2_3_6, 0.2_1_9_2, 0.2_1_8_6, 0.2_0_5_3, 0.1_9_7_1, 0.1_9_0_1, 0.1_7_1_9]

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

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






							def       _UpperCamelCase			(		self :      Optional[Any]  )       ->       List[Any]:
												_UpperCamelCase      						=	StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''     ,						safety_checker=__UpperCamelCase  )
												_UpperCamelCase      						=	LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config  )
												_UpperCamelCase      						=	sd_pipe.to(__UpperCamelCase  )
												sd_pipe.set_progress_bar_config(disable=__UpperCamelCase  )

												_UpperCamelCase      						=	"""padme amidala taking a bath artwork, safe for work, no nudity"""
												_UpperCamelCase      						=	27_3497_1755
												_UpperCamelCase      						=	7

												_UpperCamelCase      						=	torch.manual_seed(__UpperCamelCase  )
												_UpperCamelCase      						=	sd_pipe(
												    [prompt]     ,						generator=__UpperCamelCase     ,						guidance_scale=__UpperCamelCase     ,						num_inference_steps=50     ,						output_type='''np'''     ,						width=512     ,						height=512     ,						sld_guidance_scale=0     ,						)

												_UpperCamelCase      						=	output.images
												_UpperCamelCase      						=	image[0, -3:, -3:, -1]
												_UpperCamelCase      						=	[0.3_5_0_2, 0.3_6_2_2, 0.3_3_9_6, 0.3_6_4_2, 0.3_4_7_8, 0.3_3_1_8, 0.3_5, 0.3_3_4_8, 0.3_2_9_7]

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

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

												_UpperCamelCase      						=	torch.manual_seed(__UpperCamelCase  )
												_UpperCamelCase      						=	sd_pipe(
												    [prompt]     ,						generator=__UpperCamelCase     ,						guidance_scale=__UpperCamelCase     ,						num_inference_steps=50     ,						output_type='''np'''     ,						width=512     ,						height=512     ,						sld_guidance_scale=2000     ,						sld_warmup_steps=7     ,						sld_threshold=0.0_2_5     ,						sld_momentum_scale=0.5     ,						sld_mom_beta=0.7     ,						)

												_UpperCamelCase      						=	output.images
												_UpperCamelCase      						=	image[0, -3:, -3:, -1]
												_UpperCamelCase      						=	[0.5_5_3_1, 0.5_2_0_6, 0.4_8_9_5, 0.5_1_5_6, 0.5_1_8_2, 0.4_7_5_1, 0.4_8_0_2, 0.4_8_0_3, 0.4_4_4_3]

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

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






							def       _UpperCamelCase			(		self :      Union[str, Any]  )       ->       Union[str, Any]:
												_UpperCamelCase      						=	StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''  )
												_UpperCamelCase      						=	sd_pipe.to(__UpperCamelCase  )
												sd_pipe.set_progress_bar_config(disable=__UpperCamelCase  )

												_UpperCamelCase      						=	(
												    """the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c."""
												    """ leyendecker"""
												)
												_UpperCamelCase      						=	10_4435_5234
												_UpperCamelCase      						=	12

												_UpperCamelCase      						=	torch.manual_seed(__UpperCamelCase  )
												_UpperCamelCase      						=	sd_pipe(
												    [prompt]     ,						generator=__UpperCamelCase     ,						guidance_scale=__UpperCamelCase     ,						num_inference_steps=50     ,						output_type='''np'''     ,						width=512     ,						height=512     ,						sld_guidance_scale=0     ,						)

												_UpperCamelCase      						=	output.images
												_UpperCamelCase      						=	image[0, -3:, -3:, -1]
												_UpperCamelCase      						=	np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]  )

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

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

												_UpperCamelCase      						=	torch.manual_seed(__UpperCamelCase  )
												_UpperCamelCase      						=	sd_pipe(
												    [prompt]     ,						generator=__UpperCamelCase     ,						guidance_scale=__UpperCamelCase     ,						num_inference_steps=50     ,						output_type='''np'''     ,						width=512     ,						height=512     ,						sld_guidance_scale=2000     ,						sld_warmup_steps=7     ,						sld_threshold=0.0_2_5     ,						sld_momentum_scale=0.5     ,						sld_mom_beta=0.7     ,						)

												_UpperCamelCase      						=	output.images
												_UpperCamelCase      						=	image[0, -3:, -3:, -1]
												_UpperCamelCase      						=	np.array([0.5_8_1_8, 0.6_2_8_5, 0.6_8_3_5, 0.6_0_1_9, 0.6_2_5, 0.6_7_5_4, 0.6_0_9_6, 0.6_3_3_4, 0.6_5_6_1]  )
												assert image.shape == (1, 512, 512, 3)

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





'''simple docstring'''




import unittest

import torch
from torch import nn

from accelerate.test_utils import require_cuda
from accelerate.utils.memory import find_executable_batch_size, release_memory



def       __lowerCamelCase  (					)      ->			Any:
							raise RuntimeError("""CUDA out of memory."""		)

class 			UpperCamelCase_ (							nn.Module			):





	def __init__(  self    )			->			Any:
								super().__init__()
								UpperCAmelCase     :		Tuple							     =		nn.Linear(3		,		4    )
								UpperCAmelCase     :		Tuple							     =		nn.BatchNormad(4    )
								UpperCAmelCase     :		int							     =		nn.Linear(4		,		5    )






	def  _lowercase(  self		,		A    )			->			Any:
								return self.lineara(self.batchnorm(self.lineara(A    )    )    )




class 			UpperCamelCase_ (							unittest.TestCase			):





	def  _lowercase(  self    )			->			Optional[Any]:
								UpperCAmelCase     :		Optional[int]							     =		[]

								@find_executable_batch_size(starting_batch_size=128    )
								def mock_training_loop_function(A    ):
															nonlocal batch_sizes
															batch_sizes.append(A    )
															if batch_size != 8:
																						raise_fake_out_of_memory()

								mock_training_loop_function()
								self.assertListEqual(A		,		[128, 64, 32, 16, 8]    )





	def  _lowercase(  self    )			->			Any:
								UpperCAmelCase     :		Optional[Any]							     =		[]

								@find_executable_batch_size(starting_batch_size=128    )
								def mock_training_loop_function(A		,		A    ):
															nonlocal batch_sizes
															batch_sizes.append(A    )
															if batch_size != 8:
																						raise_fake_out_of_memory()
															return batch_size, arga

								UpperCAmelCase     ,     UpperCAmelCase     :		Optional[int]							     =		mock_training_loop_function("""hello"""    )
								self.assertListEqual(A		,		[128, 64, 32, 16, 8]    )
								self.assertListEqual([bs, arga]		,		[8, """hello"""]    )





	def  _lowercase(  self    )			->			Any:
								@find_executable_batch_size(starting_batch_size=0    )
								def mock_training_loop_function(A    ):
															pass

								with self.assertRaises(A    ) as cm:
															mock_training_loop_function()
															self.assertIn("""No executable batch size found, reached zero."""		,		cm.exception.args[0]    )





	def  _lowercase(  self    )			->			Optional[int]:
								@find_executable_batch_size(starting_batch_size=16    )
								def mock_training_loop_function(A    ):
															if batch_size > 0:
																						raise_fake_out_of_memory()
															pass

								with self.assertRaises(A    ) as cm:
															mock_training_loop_function()
															self.assertIn("""No executable batch size found, reached zero."""		,		cm.exception.args[0]    )





	def  _lowercase(  self    )			->			Optional[Any]:
								@find_executable_batch_size(starting_batch_size=128    )
								def mock_training_loop_function(A		,		A		,		A    ):
															if batch_size != 8:
																						raise raise_fake_out_of_memory()

								with self.assertRaises(A    ) as cm:
															mock_training_loop_function(128		,		"""hello"""		,		"""world"""    )
															self.assertIn("""Batch size was passed into `f`"""		,		cm.exception.args[0]    )
															self.assertIn("""`f(arg1='hello', arg2='world')"""		,		cm.exception.args[0]    )





	def  _lowercase(  self    )			->			int:
								@find_executable_batch_size(starting_batch_size=16    )
								def mock_training_loop_function(A    ):
															raise ValueError("""Oops, we had an error!"""    )

								with self.assertRaises(A    ) as cm:
															mock_training_loop_function()
															self.assertIn("""Oops, we had an error!"""		,		cm.exception.args[0]    )






	@require_cuda
	def  _lowercase(  self    )			->			Optional[int]:
								UpperCAmelCase     :		Optional[Any]							     =		torch.cuda.memory_allocated()
								UpperCAmelCase     :		List[str]							     =		ModelForTest()
								model.cuda()
								self.assertGreater(torch.cuda.memory_allocated()		,		A    )
								UpperCAmelCase     :		Tuple							     =		release_memory(A    )
								self.assertEqual(torch.cuda.memory_allocated()		,		A    )

