


'''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__						(		_snake_case  ,  _snake_case  ,  _snake_case  ,  unittest.TestCase      ):





							'''simple docstring'''


							A_							:   int   			=						StableDiffusionModelEditingPipeline
							A_							:   Optional[int]   			=						TEXT_TO_IMAGE_PARAMS
							A_							:   Optional[Any]   			=						TEXT_TO_IMAGE_BATCH_PARAMS
							A_							:   Optional[Any]   			=						TEXT_TO_IMAGE_IMAGE_PARAMS
							A_							:   Tuple   			=						TEXT_TO_IMAGE_IMAGE_PARAMS






							def UpperCAmelCase_			(					self						):
											torch.manual_seed(0						)
											_SCREAMING_SNAKE_CASE				: Tuple  	=   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					,      )
											_SCREAMING_SNAKE_CASE				: Optional[Any]  	=   DDIMScheduler()
											torch.manual_seed(0						)
											_SCREAMING_SNAKE_CASE				: Optional[int]  	=   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						)
											_SCREAMING_SNAKE_CASE				: Dict  	=   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					,      )
											_SCREAMING_SNAKE_CASE				: Union[str, Any]  	=   CLIPTextModel(__snake_case						)
											_SCREAMING_SNAKE_CASE				: Union[str, Any]  	=   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

											_SCREAMING_SNAKE_CASE				: str  	=   {
											    """unet""": unet,
											    """scheduler""": scheduler,
											    """vae""": vae,
											    """text_encoder""": text_encoder,
											    """tokenizer""": tokenizer,
											    """safety_checker""": None,
											    """feature_extractor""": None,
											}
											return components






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

											_SCREAMING_SNAKE_CASE				: List[Any]  	=   self.get_dummy_inputs(__snake_case						)
											_SCREAMING_SNAKE_CASE				: Tuple  	=   sd_pipe(**__snake_case						).images
											_SCREAMING_SNAKE_CASE				: Any  	=   image[0, -3:, -3:, -1]
											assert image.shape == (1, 64, 64, 3)

											_SCREAMING_SNAKE_CASE				: Dict  	=   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 UpperCAmelCase_			(					self						):
											_SCREAMING_SNAKE_CASE				: str  	=   """cpu"""  # ensure determinism for the device-dependent torch.Generator
											_SCREAMING_SNAKE_CASE				: List[str]  	=   self.get_dummy_components()
											_SCREAMING_SNAKE_CASE				: Any  	=   StableDiffusionModelEditingPipeline(**__snake_case						)
											_SCREAMING_SNAKE_CASE				: Any  	=   sd_pipe.to(__snake_case						)
											sd_pipe.set_progress_bar_config(disable=__snake_case						)

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

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

											_SCREAMING_SNAKE_CASE				: Tuple  	=   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 UpperCAmelCase_			(					self						):
											_SCREAMING_SNAKE_CASE				: Tuple  	=   """cpu"""  # ensure determinism for the device-dependent torch.Generator
											_SCREAMING_SNAKE_CASE				: Dict  	=   self.get_dummy_components()
											_SCREAMING_SNAKE_CASE				: Optional[Any]  	=   EulerAncestralDiscreteScheduler(
											    beta_start=0.00085					,      beta_end=0.012					,      beta_schedule="""scaled_linear"""						)
											_SCREAMING_SNAKE_CASE				: Dict  	=   StableDiffusionModelEditingPipeline(**__snake_case						)
											_SCREAMING_SNAKE_CASE				: List[str]  	=   sd_pipe.to(__snake_case						)
											sd_pipe.set_progress_bar_config(disable=__snake_case						)

											_SCREAMING_SNAKE_CASE				: Optional[int]  	=   self.get_dummy_inputs(__snake_case						)
											_SCREAMING_SNAKE_CASE				: List[str]  	=   sd_pipe(**__snake_case						).images
											_SCREAMING_SNAKE_CASE				: int  	=   image[0, -3:, -3:, -1]

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

											_SCREAMING_SNAKE_CASE				: Optional[Any]  	=   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 UpperCAmelCase_			(					self						):
											_SCREAMING_SNAKE_CASE				: Optional[Any]  	=   """cpu"""  # ensure determinism for the device-dependent torch.Generator
											_SCREAMING_SNAKE_CASE				: Dict  	=   self.get_dummy_components()
											_SCREAMING_SNAKE_CASE				: Optional[int]  	=   PNDMScheduler()
											_SCREAMING_SNAKE_CASE				: Any  	=   StableDiffusionModelEditingPipeline(**__snake_case						)
											_SCREAMING_SNAKE_CASE				: Optional[Any]  	=   sd_pipe.to(__snake_case						)
											sd_pipe.set_progress_bar_config(disable=__snake_case						)

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






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






							def UpperCAmelCase_			(					self						):
											super().test_attention_slicing_forward_pass(expected_max_diff=5e-3						)





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





							'''simple docstring'''


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






							def UpperCAmelCase_			(					self					,      __snake_case=0						):
											_SCREAMING_SNAKE_CASE				: Dict  	=   torch.manual_seed(__snake_case						)
											_SCREAMING_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 UpperCAmelCase_			(					self						):
											_SCREAMING_SNAKE_CASE				: List[Any]  	=   """CompVis/stable-diffusion-v1-4"""
											_SCREAMING_SNAKE_CASE				: str  	=   StableDiffusionModelEditingPipeline.from_pretrained(__snake_case					,      safety_checker=__snake_case						)
											pipe.to(__snake_case						)
											pipe.set_progress_bar_config(disable=__snake_case						)
											pipe.enable_attention_slicing()

											_SCREAMING_SNAKE_CASE				: List[Any]  	=   self.get_inputs()
											_SCREAMING_SNAKE_CASE				: Union[str, Any]  	=   pipe(**__snake_case						).images
											_SCREAMING_SNAKE_CASE				: Optional[int]  	=   image[0, -3:, -3:, -1].flatten()

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

											_SCREAMING_SNAKE_CASE				: Optional[Any]  	=   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"""						)

											_SCREAMING_SNAKE_CASE				: str  	=   pipe(**__snake_case						).images
											_SCREAMING_SNAKE_CASE				: Union[str, Any]  	=   image[0, -3:, -3:, -1].flatten()

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

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






							def UpperCAmelCase_			(					self						):
											torch.cuda.empty_cache()
											torch.cuda.reset_max_memory_allocated()
											torch.cuda.reset_peak_memory_stats()

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

											_SCREAMING_SNAKE_CASE				: int  	=   self.get_inputs()
											_SCREAMING_SNAKE_CASE				: Dict  	=   pipe(**__snake_case						)

											_SCREAMING_SNAKE_CASE				: List[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_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCAmelCase_	:     Union[str, Any]							     =	{'configuration_pegasus': ['PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP', 'PegasusConfig']}

try:
				if not is_sentencepiece_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCAmelCase_	:     Dict							     =	['PegasusTokenizer']

try:
				if not is_tokenizers_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCAmelCase_	:     Optional[int]							     =	['PegasusTokenizerFast']

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCAmelCase_	:     Optional[int]							     =	[
				    'PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST',
				    'PegasusForCausalLM',
				    'PegasusForConditionalGeneration',
				    'PegasusModel',
				    'PegasusPreTrainedModel',
				]

try:
				if not is_tf_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCAmelCase_	:     Optional[Any]							     =	[
				    'TFPegasusForConditionalGeneration',
				    'TFPegasusModel',
				    'TFPegasusPreTrainedModel',
				]

try:
				if not is_flax_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCAmelCase_	:     Any							     =	[
				    'FlaxPegasusForConditionalGeneration',
				    'FlaxPegasusModel',
				    'FlaxPegasusPreTrainedModel',
				]


if TYPE_CHECKING:
				from .configuration_pegasus import PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP, PegasusConfig

				try:
								if not is_sentencepiece_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .tokenization_pegasus import PegasusTokenizer

				try:
								if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .tokenization_pegasus_fast import PegasusTokenizerFast

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_pegasus import (
								    PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST,
								    PegasusForCausalLM,
								    PegasusForConditionalGeneration,
								    PegasusModel,
								    PegasusPreTrainedModel,
								)

				try:
								if not is_tf_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_tf_pegasus import TFPegasusForConditionalGeneration, TFPegasusModel, TFPegasusPreTrainedModel

				try:
								if not is_flax_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_flax_pegasus import (
								    FlaxPegasusForConditionalGeneration,
								    FlaxPegasusModel,
								    FlaxPegasusPreTrainedModel,
								)

else:
				import sys

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




