



'''simple docstring'''



import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






@skip_mps
class 			_SCREAMING_SNAKE_CASE				(lowerCAmelCase_,	lowerCAmelCase_,	lowerCAmelCase_,	unittest.TestCase      ):
			A__								=  StableDiffusionModelEditingPipeline
			A__								=  TEXT_TO_IMAGE_PARAMS
			A__								=  TEXT_TO_IMAGE_BATCH_PARAMS
			A__								=  TEXT_TO_IMAGE_IMAGE_PARAMS
			A__								=  TEXT_TO_IMAGE_IMAGE_PARAMS



			def 				lowerCAmelCase (     self		:						Optional[int]			)      ->      str:
							"""simple docstring"""





							torch.manual_seed(0			)
							snake_case__      :   Optional[int]      =    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    ,			)
							snake_case__      :   Any      =    DDIMScheduler()
							torch.manual_seed(0			)
							snake_case__      :   Union[str, Any]      =    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    ,			)
							torch.manual_seed(0			)
							snake_case__      :   List[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    ,			)
							snake_case__      :   List[Any]      =    CLIPTextModel(__UpperCamelCase			)
							snake_case__      :   Tuple      =    CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''			)

							snake_case__      :   Any      =    {
							    '''unet''': unet,
							    '''scheduler''': scheduler,
							    '''vae''': vae,
							    '''text_encoder''': text_encoder,
							    '''tokenizer''': tokenizer,
							    '''safety_checker''': None,
							    '''feature_extractor''': None,
							}
							return components



			def 				lowerCAmelCase (     self		:						Tuple    ,			__UpperCamelCase		:						Optional[int]    ,			__UpperCamelCase		:						List[Any]=0			)      ->      Dict:
							"""simple docstring"""





							snake_case__      :   Union[str, Any]      =    torch.manual_seed(__UpperCamelCase			)
							snake_case__      :   Tuple      =    {
							    '''prompt''': '''A field of roses''',
							    '''generator''': generator,
							    # Setting height and width to None to prevent OOMs on CPU.
							    '''height''': None,
							    '''width''': None,
							    '''num_inference_steps''': 2,
							    '''guidance_scale''': 6.0,
							    '''output_type''': '''numpy''',
							}
							return inputs



			def 				lowerCAmelCase (     self		:						List[Any]			)      ->      int:
							"""simple docstring"""





							snake_case__      :   str      =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							snake_case__      :   int      =    self.get_dummy_components()
							snake_case__      :   Any      =    StableDiffusionModelEditingPipeline(**__UpperCamelCase			)
							snake_case__      :   List[Any]      =    sd_pipe.to(__UpperCamelCase			)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							snake_case__      :   Optional[int]      =    self.get_dummy_inputs(__UpperCamelCase			)
							snake_case__      :   Dict      =    sd_pipe(**__UpperCamelCase			).images
							snake_case__      :   List[Any]      =    image[0, -3:, -3:, -1]
							assert image.shape == (1, 64, 64, 3)

							snake_case__      :   Optional[Any]      =    np.array([0.4755, 0.5132, 0.4976, 0.3904, 0.3554, 0.4765, 0.5139, 0.5158, 0.4889]			)

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



			def 				lowerCAmelCase (     self		:						int			)      ->      List[Any]:
							"""simple docstring"""





							snake_case__      :   Dict      =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							snake_case__      :   List[Any]      =    self.get_dummy_components()
							snake_case__      :   int      =    StableDiffusionModelEditingPipeline(**__UpperCamelCase			)
							snake_case__      :   Optional[int]      =    sd_pipe.to(__UpperCamelCase			)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							snake_case__      :   Optional[int]      =    self.get_dummy_inputs(__UpperCamelCase			)
							snake_case__      :   Union[str, Any]      =    '''french fries'''
							snake_case__      :   Optional[int]      =    sd_pipe(**__UpperCamelCase    ,			negative_prompt=__UpperCamelCase			)
							snake_case__      :   Optional[Any]      =    output.images
							snake_case__      :   str      =    image[0, -3:, -3:, -1]

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

							snake_case__      :   Union[str, Any]      =    np.array([0.4992, 0.5101, 0.5004, 0.3949, 0.3604, 0.4735, 0.5216, 0.5204, 0.4913]			)

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



			def 				lowerCAmelCase (     self		:						List[Any]			)      ->      Tuple:
							"""simple docstring"""





							snake_case__      :   int      =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							snake_case__      :   Dict      =    self.get_dummy_components()
							snake_case__      :   Any      =    EulerAncestralDiscreteScheduler(
							    beta_start=0.00085    ,			beta_end=0.012    ,			beta_schedule='''scaled_linear'''			)
							snake_case__      :   List[Any]      =    StableDiffusionModelEditingPipeline(**__UpperCamelCase			)
							snake_case__      :   Union[str, Any]      =    sd_pipe.to(__UpperCamelCase			)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							snake_case__      :   Optional[Any]      =    self.get_dummy_inputs(__UpperCamelCase			)
							snake_case__      :   int      =    sd_pipe(**__UpperCamelCase			).images
							snake_case__      :   Tuple      =    image[0, -3:, -3:, -1]

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

							snake_case__      :   List[str]      =    np.array([0.4747, 0.5372, 0.4779, 0.4982, 0.5543, 0.4816, 0.5238, 0.4904, 0.5027]			)

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



			def 				lowerCAmelCase (     self		:						List[str]			)      ->      Dict:
							"""simple docstring"""





							snake_case__      :   Any      =    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							snake_case__      :   Union[str, Any]      =    self.get_dummy_components()
							snake_case__      :   str      =    PNDMScheduler()
							snake_case__      :   Dict      =    StableDiffusionModelEditingPipeline(**__UpperCamelCase			)
							snake_case__      :   Optional[Any]      =    sd_pipe.to(__UpperCamelCase			)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							snake_case__      :   Tuple      =    self.get_dummy_inputs(__UpperCamelCase			)
							# the pipeline does not expect pndm so test if it raises error.
							with self.assertRaises(__UpperCamelCase			):
											snake_case__      :   Union[str, Any]      =    sd_pipe(**__UpperCamelCase			).images



			def 				lowerCAmelCase (     self		:						Optional[int]			)      ->      Optional[Any]:
							"""simple docstring"""





							super().test_inference_batch_single_identical(expected_max_diff=5e-3			)







			def 				lowerCAmelCase (     self		:						Any			)      ->      str:
							"""simple docstring"""





							super().test_attention_slicing_forward_pass(expected_max_diff=5e-3			)






@slow
@require_torch_gpu
class 			_SCREAMING_SNAKE_CASE				(unittest.TestCase      ):



			def 				lowerCAmelCase (     self		:						Any			)      ->      int:
							"""simple docstring"""





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



			def 				lowerCAmelCase (     self		:						Union[str, Any]    ,			__UpperCamelCase		:						Dict=0			)      ->      List[Any]:
							"""simple docstring"""





							snake_case__      :   Any      =    torch.manual_seed(__UpperCamelCase			)
							snake_case__      :   List[Any]      =    {
							    '''prompt''': '''A field of roses''',
							    '''generator''': generator,
							    '''num_inference_steps''': 3,
							    '''guidance_scale''': 7.5,
							    '''output_type''': '''numpy''',
							}
							return inputs



			def 				lowerCAmelCase (     self		:						int			)      ->      Any:
							"""simple docstring"""





							snake_case__      :   Tuple      =    '''CompVis/stable-diffusion-v1-4'''
							snake_case__      :   Tuple      =    StableDiffusionModelEditingPipeline.from_pretrained(__UpperCamelCase    ,			safety_checker=__UpperCamelCase			)
							pipe.to(__UpperCamelCase			)
							pipe.set_progress_bar_config(disable=__UpperCamelCase			)
							pipe.enable_attention_slicing()

							snake_case__      :   Tuple      =    self.get_inputs()
							snake_case__      :   List[str]      =    pipe(**__UpperCamelCase			).images
							snake_case__      :   Optional[int]      =    image[0, -3:, -3:, -1].flatten()

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

							snake_case__      :   Tuple      =    np.array(
							    [0.6749496, 0.6386453, 0.51443267, 0.66094905, 0.61921215, 0.5491332, 0.5744417, 0.58075106, 0.5174658]			)

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

							# make sure image changes after editing
							pipe.edit_model('''A pack of roses'''    ,			'''A pack of blue roses'''			)

							snake_case__      :   Any      =    pipe(**__UpperCamelCase			).images
							snake_case__      :   List[str]      =    image[0, -3:, -3:, -1].flatten()

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

							assert np.abs(expected_slice - image_slice			).max() > 1e-1







			def 				lowerCAmelCase (     self		:						Union[str, Any]			)      ->      Tuple:
							"""simple docstring"""





							torch.cuda.empty_cache()
							torch.cuda.reset_max_memory_allocated()
							torch.cuda.reset_peak_memory_stats()

							snake_case__      :   Any      =    '''CompVis/stable-diffusion-v1-4'''
							snake_case__      :   Tuple      =    DDIMScheduler.from_pretrained(__UpperCamelCase    ,			subfolder='''scheduler'''			)
							snake_case__      :   Tuple      =    StableDiffusionModelEditingPipeline.from_pretrained(
							    __UpperCamelCase    ,			scheduler=__UpperCamelCase    ,			safety_checker=__UpperCamelCase			)
							snake_case__      :   int      =    pipe.to(__UpperCamelCase			)
							pipe.set_progress_bar_config(disable=__UpperCamelCase			)
							pipe.enable_attention_slicing(1			)
							pipe.enable_sequential_cpu_offload()

							snake_case__      :   int      =    self.get_inputs()
							snake_case__      :   Optional[int]      =    pipe(**__UpperCamelCase			)

							snake_case__      :   int      =    torch.cuda.max_memory_allocated()
							# make sure that less than 4.4 GB is allocated
							assert mem_bytes < 4.4 * 10**9









"""simple docstring"""






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_lowercase					:					str 		=		{
    "configuration_openai": ["OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP", "OpenAIGPTConfig"],
    "tokenization_openai": ["OpenAIGPTTokenizer"],
}

try:
		if not is_tokenizers_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase					:					List[Any] 		=		["OpenAIGPTTokenizerFast"]

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase					:					Union[str, Any] 		=		[
		    "OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST",
		    "OpenAIGPTDoubleHeadsModel",
		    "OpenAIGPTForSequenceClassification",
		    "OpenAIGPTLMHeadModel",
		    "OpenAIGPTModel",
		    "OpenAIGPTPreTrainedModel",
		    "load_tf_weights_in_openai_gpt",
		]

try:
		if not is_tf_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase					:					Optional[Any] 		=		[
		    "TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST",
		    "TFOpenAIGPTDoubleHeadsModel",
		    "TFOpenAIGPTForSequenceClassification",
		    "TFOpenAIGPTLMHeadModel",
		    "TFOpenAIGPTMainLayer",
		    "TFOpenAIGPTModel",
		    "TFOpenAIGPTPreTrainedModel",
		]


if TYPE_CHECKING:
		from .configuration_openai import OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP, OpenAIGPTConfig
		from .tokenization_openai import OpenAIGPTTokenizer

		try:
				if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .tokenization_openai_fast import OpenAIGPTTokenizerFast

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_openai import (
				    OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    OpenAIGPTDoubleHeadsModel,
				    OpenAIGPTForSequenceClassification,
				    OpenAIGPTLMHeadModel,
				    OpenAIGPTModel,
				    OpenAIGPTPreTrainedModel,
				    load_tf_weights_in_openai_gpt,
				)

		try:
				if not is_tf_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_tf_openai import (
				    TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFOpenAIGPTDoubleHeadsModel,
				    TFOpenAIGPTForSequenceClassification,
				    TFOpenAIGPTLMHeadModel,
				    TFOpenAIGPTMainLayer,
				    TFOpenAIGPTModel,
				    TFOpenAIGPTPreTrainedModel,
				)

else:
		import sys

		_lowercase					:					Dict 		=		_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)




