






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

														lowerCamelCase__   :      List[Any]			    =			{
														    """unet""": unet,
														    """scheduler""": scheduler,
														    """vae""": vae,
														    """text_encoder""": text_encoder,
														    """tokenizer""": tokenizer,
														    """safety_checker""": None,
														    """feature_extractor""": None,
														}
														return components




							def 						lowerCamelCase_							(				self:							Optional[Any]     ,				UpperCamelCase__:							Tuple     ,				UpperCamelCase__:							Optional[Any]=0	):
														lowerCamelCase__   :      Dict			    =			torch.manual_seed(UpperCamelCase__	)
														lowerCamelCase__   :      List[Any]			    =			{
														    """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:							Optional[int]	):
														lowerCamelCase__   :      Optional[int]			    =			"""cpu"""  # ensure determinism for the device-dependent torch.Generator
														lowerCamelCase__   :      Optional[Any]			    =			self.get_dummy_components()
														lowerCamelCase__   :      Union[str, Any]			    =			StableDiffusionModelEditingPipeline(**UpperCamelCase__	)
														lowerCamelCase__   :      int			    =			sd_pipe.to(UpperCamelCase__	)
														sd_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

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

														lowerCamelCase__   :      Any			    =			np.array([0.4_755, 0.5_132, 0.4_976, 0.3_904, 0.3_554, 0.4_765, 0.5_139, 0.5_158, 0.4_889]	)

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




							def 						lowerCamelCase_							(				self:							Dict	):
														lowerCamelCase__   :      int			    =			"""cpu"""  # ensure determinism for the device-dependent torch.Generator
														lowerCamelCase__   :      Any			    =			self.get_dummy_components()
														lowerCamelCase__   :      Dict			    =			StableDiffusionModelEditingPipeline(**UpperCamelCase__	)
														lowerCamelCase__   :      Tuple			    =			sd_pipe.to(UpperCamelCase__	)
														sd_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

														lowerCamelCase__   :      List[str]			    =			self.get_dummy_inputs(UpperCamelCase__	)
														lowerCamelCase__   :      Optional[int]			    =			"""french fries"""
														lowerCamelCase__   :      Union[str, Any]			    =			sd_pipe(**UpperCamelCase__     ,				negative_prompt=UpperCamelCase__	)
														lowerCamelCase__   :      int			    =			output.images
														lowerCamelCase__   :      List[str]			    =			image[0, -3:, -3:, -1]

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

														lowerCamelCase__   :      str			    =			np.array([0.4_992, 0.5_101, 0.5_004, 0.3_949, 0.3_604, 0.4_735, 0.5_216, 0.5_204, 0.4_913]	)

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




							def 						lowerCamelCase_							(				self:							Union[str, Any]	):
														lowerCamelCase__   :      Optional[Any]			    =			"""cpu"""  # ensure determinism for the device-dependent torch.Generator
														lowerCamelCase__   :      List[str]			    =			self.get_dummy_components()
														lowerCamelCase__   :      int			    =			EulerAncestralDiscreteScheduler(
														    beta_start=0.00_085     ,				beta_end=0.012     ,				beta_schedule="""scaled_linear"""	)
														lowerCamelCase__   :      Dict			    =			StableDiffusionModelEditingPipeline(**UpperCamelCase__	)
														lowerCamelCase__   :      str			    =			sd_pipe.to(UpperCamelCase__	)
														sd_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

														lowerCamelCase__   :      List[Any]			    =			self.get_dummy_inputs(UpperCamelCase__	)
														lowerCamelCase__   :      Optional[int]			    =			sd_pipe(**UpperCamelCase__	).images
														lowerCamelCase__   :      Any			    =			image[0, -3:, -3:, -1]

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

														lowerCamelCase__   :      Optional[Any]			    =			np.array([0.4_747, 0.5_372, 0.4_779, 0.4_982, 0.5_543, 0.4_816, 0.5_238, 0.4_904, 0.5_027]	)

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




							def 						lowerCamelCase_							(				self:							List[str]	):
														lowerCamelCase__   :      str			    =			"""cpu"""  # ensure determinism for the device-dependent torch.Generator
														lowerCamelCase__   :      Dict			    =			self.get_dummy_components()
														lowerCamelCase__   :      List[Any]			    =			PNDMScheduler()
														lowerCamelCase__   :      int			    =			StableDiffusionModelEditingPipeline(**UpperCamelCase__	)
														lowerCamelCase__   :      Dict			    =			sd_pipe.to(UpperCamelCase__	)
														sd_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

														lowerCamelCase__   :      List[str]			    =			self.get_dummy_inputs(UpperCamelCase__	)
														# the pipeline does not expect pndm so test if it raises error.
														with self.assertRaises(UpperCamelCase__	):
																					lowerCamelCase__   :      Dict			    =			sd_pipe(**UpperCamelCase__	).images




							def 						lowerCamelCase_							(				self:							str	):
														super().test_inference_batch_single_identical(expected_max_diff=5e-3	)




							def 						lowerCamelCase_							(				self:							Union[str, Any]	):
														super().test_attention_slicing_forward_pass(expected_max_diff=5e-3	)



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




							def 						lowerCamelCase_							(				self:							str	):
														super().tearDown()
														gc.collect()
														torch.cuda.empty_cache()




							def 						lowerCamelCase_							(				self:							Dict     ,				UpperCamelCase__:							Union[str, Any]=0	):
														lowerCamelCase__   :      Optional[int]			    =			torch.manual_seed(UpperCamelCase__	)
														lowerCamelCase__   :      Dict			    =			{
														    """prompt""": """A field of roses""",
														    """generator""": generator,
														    """num_inference_steps""": 3,
														    """guidance_scale""": 7.5,
														    """output_type""": """numpy""",
														}
														return inputs




							def 						lowerCamelCase_							(				self:							List[str]	):
														lowerCamelCase__   :      Union[str, Any]			    =			"""CompVis/stable-diffusion-v1-4"""
														lowerCamelCase__   :      str			    =			StableDiffusionModelEditingPipeline.from_pretrained(UpperCamelCase__     ,				safety_checker=UpperCamelCase__	)
														pipe.to(UpperCamelCase__	)
														pipe.set_progress_bar_config(disable=UpperCamelCase__	)
														pipe.enable_attention_slicing()

														lowerCamelCase__   :      str			    =			self.get_inputs()
														lowerCamelCase__   :      Dict			    =			pipe(**UpperCamelCase__	).images
														lowerCamelCase__   :      Optional[Any]			    =			image[0, -3:, -3:, -1].flatten()

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

														lowerCamelCase__   :      Any			    =			np.array(
														    [0.6_749_496, 0.6_386_453, 0.51_443_267, 0.66_094_905, 0.61_921_215, 0.5_491_332, 0.5_744_417, 0.58_075_106, 0.5_174_658]	)

														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"""	)

														lowerCamelCase__   :      Optional[Any]			    =			pipe(**UpperCamelCase__	).images
														lowerCamelCase__   :      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:							Dict	):
														torch.cuda.empty_cache()
														torch.cuda.reset_max_memory_allocated()
														torch.cuda.reset_peak_memory_stats()

														lowerCamelCase__   :      List[Any]			    =			"""CompVis/stable-diffusion-v1-4"""
														lowerCamelCase__   :      List[str]			    =			DDIMScheduler.from_pretrained(UpperCamelCase__     ,				subfolder="""scheduler"""	)
														lowerCamelCase__   :      str			    =			StableDiffusionModelEditingPipeline.from_pretrained(
														    UpperCamelCase__     ,				scheduler=UpperCamelCase__     ,				safety_checker=UpperCamelCase__	)
														lowerCamelCase__   :      int			    =			pipe.to(UpperCamelCase__	)
														pipe.set_progress_bar_config(disable=UpperCamelCase__	)
														pipe.enable_attention_slicing(1	)
														pipe.enable_sequential_cpu_offload()

														lowerCamelCase__   :      Any			    =			self.get_inputs()
														lowerCamelCase__   :      List[str]			    =			pipe(**UpperCamelCase__	)

														lowerCamelCase__   :      Any			    =			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_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__)

