





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      ):
   _UpperCamelCase    :		str 					=       StableDiffusionModelEditingPipeline
   _UpperCamelCase    :		int 					=       TEXT_TO_IMAGE_PARAMS
   _UpperCamelCase    :		List[Any] 					=       TEXT_TO_IMAGE_BATCH_PARAMS
   _UpperCamelCase    :		Tuple 					=       TEXT_TO_IMAGE_IMAGE_PARAMS
   _UpperCamelCase    :		Any 					=       TEXT_TO_IMAGE_IMAGE_PARAMS


   def        SCREAMING_SNAKE_CASE_							(    self			:				List[Any]       )->   Dict:





     """simple docstring"""







     torch.manual_seed(0       )
     lowercase__				=  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			,     )
     lowercase__				=  DDIMScheduler()
     torch.manual_seed(0       )
     lowercase__				=  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       )
     lowercase__				=  CLIPTextConfig(
         bos_token_id=0			,     eos_token_id=2			,     hidden_size=32			,     intermediate_size=37			,     layer_norm_eps=1E-0_5			,     num_attention_heads=4			,     num_hidden_layers=5			,     pad_token_id=1			,     vocab_size=1_000			,     )
     lowercase__				=  CLIPTextModel(_lowerCamelCase       )
     lowercase__				=  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

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


   def        SCREAMING_SNAKE_CASE_							(    self			:				List[str]			,     a			:				Any			,     a			:				Union[str, Any]=0       )->   Any:





     """simple docstring"""







     lowercase__				=  torch.manual_seed(_lowerCamelCase       )
     lowercase__				=  {
         '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        SCREAMING_SNAKE_CASE_							(    self			:				Union[str, Any]       )->   Tuple:





     """simple docstring"""







     lowercase__				=  'cpu'  # ensure determinism for the device-dependent torch.Generator
     lowercase__				=  self.get_dummy_components()
     lowercase__				=  StableDiffusionModelEditingPipeline(**_lowerCamelCase       )
     lowercase__				=  sd_pipe.to(_lowerCamelCase       )
     sd_pipe.set_progress_bar_config(disable=_lowerCamelCase       )

     lowercase__				=  self.get_dummy_inputs(_lowerCamelCase       )
     lowercase__				=  sd_pipe(**_lowerCamelCase       ).images
     lowercase__				=  image[0, -3:, -3:, -1]
     assert image.shape == (1, 64, 64, 3)

     lowercase__				=  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        SCREAMING_SNAKE_CASE_							(    self			:				str       )->   Optional[int]:





     """simple docstring"""







     lowercase__				=  'cpu'  # ensure determinism for the device-dependent torch.Generator
     lowercase__				=  self.get_dummy_components()
     lowercase__				=  StableDiffusionModelEditingPipeline(**_lowerCamelCase       )
     lowercase__				=  sd_pipe.to(_lowerCamelCase       )
     sd_pipe.set_progress_bar_config(disable=_lowerCamelCase       )

     lowercase__				=  self.get_dummy_inputs(_lowerCamelCase       )
     lowercase__				=  'french fries'
     lowercase__				=  sd_pipe(**_lowerCamelCase			,     negative_prompt=_lowerCamelCase       )
     lowercase__				=  output.images
     lowercase__				=  image[0, -3:, -3:, -1]

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

     lowercase__				=  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        SCREAMING_SNAKE_CASE_							(    self			:				Union[str, Any]       )->   Tuple:





     """simple docstring"""







     lowercase__				=  'cpu'  # ensure determinism for the device-dependent torch.Generator
     lowercase__				=  self.get_dummy_components()
     lowercase__				=  EulerAncestralDiscreteScheduler(
         beta_start=0.00085			,     beta_end=0.012			,     beta_schedule='scaled_linear'       )
     lowercase__				=  StableDiffusionModelEditingPipeline(**_lowerCamelCase       )
     lowercase__				=  sd_pipe.to(_lowerCamelCase       )
     sd_pipe.set_progress_bar_config(disable=_lowerCamelCase       )

     lowercase__				=  self.get_dummy_inputs(_lowerCamelCase       )
     lowercase__				=  sd_pipe(**_lowerCamelCase       ).images
     lowercase__				=  image[0, -3:, -3:, -1]

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

     lowercase__				=  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        SCREAMING_SNAKE_CASE_							(    self			:				Optional[int]       )->   Optional[int]:





     """simple docstring"""







     lowercase__				=  'cpu'  # ensure determinism for the device-dependent torch.Generator
     lowercase__				=  self.get_dummy_components()
     lowercase__				=  PNDMScheduler()
     lowercase__				=  StableDiffusionModelEditingPipeline(**_lowerCamelCase       )
     lowercase__				=  sd_pipe.to(_lowerCamelCase       )
     sd_pipe.set_progress_bar_config(disable=_lowerCamelCase       )

     lowercase__				=  self.get_dummy_inputs(_lowerCamelCase       )
     # the pipeline does not expect pndm so test if it raises error.
     with self.assertRaises(_lowerCamelCase       ):
       lowercase__				=  sd_pipe(**_lowerCamelCase       ).images


   def        SCREAMING_SNAKE_CASE_							(    self			:				List[Any]       )->   int:





     """simple docstring"""







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


   def        SCREAMING_SNAKE_CASE_							(    self			:				Dict       )->   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        SCREAMING_SNAKE_CASE_							(    self			:				int       )->   List[str]:





     """simple docstring"""







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


   def        SCREAMING_SNAKE_CASE_							(    self			:				Union[str, Any]			,     a			:				Tuple=0       )->   Any:





     """simple docstring"""







     lowercase__				=  torch.manual_seed(_lowerCamelCase       )
     lowercase__				=  {
         'prompt': 'A field of roses',
         'generator': generator,
         'num_inference_steps': 3,
         'guidance_scale': 7.5,
         'output_type': 'numpy',
     }
     return inputs


   def        SCREAMING_SNAKE_CASE_							(    self			:				Optional[int]       )->   str:





     """simple docstring"""







     lowercase__				=  'CompVis/stable-diffusion-v1-4'
     lowercase__				=  StableDiffusionModelEditingPipeline.from_pretrained(_lowerCamelCase			,     safety_checker=_lowerCamelCase       )
     pipe.to(_lowerCamelCase       )
     pipe.set_progress_bar_config(disable=_lowerCamelCase       )
     pipe.enable_attention_slicing()

     lowercase__				=  self.get_inputs()
     lowercase__				=  pipe(**_lowerCamelCase       ).images
     lowercase__				=  image[0, -3:, -3:, -1].flatten()

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

     lowercase__				=  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'       )

     lowercase__				=  pipe(**_lowerCamelCase       ).images
     lowercase__				=  image[0, -3:, -3:, -1].flatten()

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

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


   def        SCREAMING_SNAKE_CASE_							(    self			:				Tuple       )->   List[str]:





     """simple docstring"""







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

     lowercase__				=  'CompVis/stable-diffusion-v1-4'
     lowercase__				=  DDIMScheduler.from_pretrained(_lowerCamelCase			,     subfolder='scheduler'       )
     lowercase__				=  StableDiffusionModelEditingPipeline.from_pretrained(
         _lowerCamelCase			,     scheduler=_lowerCamelCase			,     safety_checker=_lowerCamelCase       )
     lowercase__				=  pipe.to(_lowerCamelCase       )
     pipe.set_progress_bar_config(disable=_lowerCamelCase       )
     pipe.enable_attention_slicing(1       )
     pipe.enable_sequential_cpu_offload()

     lowercase__				=  self.get_inputs()
     lowercase__				=  pipe(**_lowerCamelCase       )

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









'''simple docstring'''



import importlib
import inspect
import json
import os
from collections import OrderedDict

# Build the list of all feature extractors
from ...configuration_utils import PretrainedConfig
from ...dynamic_module_utils import get_class_from_dynamic_module, resolve_trust_remote_code
from ...feature_extraction_utils import FeatureExtractionMixin
from ...image_processing_utils import ImageProcessingMixin
from ...tokenization_utils import TOKENIZER_CONFIG_FILE
from ...utils import FEATURE_EXTRACTOR_NAME, get_file_from_repo, logging
from .auto_factory import _LazyAutoMapping
from .configuration_auto import (
    CONFIG_MAPPING_NAMES,
    AutoConfig,
    model_type_to_module_name,
    replace_list_option_in_docstrings,
)
from .feature_extraction_auto import AutoFeatureExtractor
from .image_processing_auto import AutoImageProcessor
from .tokenization_auto import AutoTokenizer


_SCREAMING_SNAKE_CASE    =			logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE    =			OrderedDict(
    [
        ('''align''', '''AlignProcessor'''),
        ('''altclip''', '''AltCLIPProcessor'''),
        ('''bark''', '''BarkProcessor'''),
        ('''blip''', '''BlipProcessor'''),
        ('''blip-2''', '''Blip2Processor'''),
        ('''bridgetower''', '''BridgeTowerProcessor'''),
        ('''chinese_clip''', '''ChineseCLIPProcessor'''),
        ('''clap''', '''ClapProcessor'''),
        ('''clip''', '''CLIPProcessor'''),
        ('''clipseg''', '''CLIPSegProcessor'''),
        ('''flava''', '''FlavaProcessor'''),
        ('''git''', '''GitProcessor'''),
        ('''groupvit''', '''CLIPProcessor'''),
        ('''hubert''', '''Wav2Vec2Processor'''),
        ('''instructblip''', '''InstructBlipProcessor'''),
        ('''layoutlmv2''', '''LayoutLMv2Processor'''),
        ('''layoutlmv3''', '''LayoutLMv3Processor'''),
        ('''markuplm''', '''MarkupLMProcessor'''),
        ('''mctct''', '''MCTCTProcessor'''),
        ('''mgp-str''', '''MgpstrProcessor'''),
        ('''oneformer''', '''OneFormerProcessor'''),
        ('''owlvit''', '''OwlViTProcessor'''),
        ('''pix2struct''', '''Pix2StructProcessor'''),
        ('''sam''', '''SamProcessor'''),
        ('''sew''', '''Wav2Vec2Processor'''),
        ('''sew-d''', '''Wav2Vec2Processor'''),
        ('''speech_to_text''', '''Speech2TextProcessor'''),
        ('''speech_to_text_2''', '''Speech2Text2Processor'''),
        ('''speecht5''', '''SpeechT5Processor'''),
        ('''trocr''', '''TrOCRProcessor'''),
        ('''tvlt''', '''TvltProcessor'''),
        ('''unispeech''', '''Wav2Vec2Processor'''),
        ('''unispeech-sat''', '''Wav2Vec2Processor'''),
        ('''vilt''', '''ViltProcessor'''),
        ('''vision-text-dual-encoder''', '''VisionTextDualEncoderProcessor'''),
        ('''wav2vec2''', '''Wav2Vec2Processor'''),
        ('''wav2vec2-conformer''', '''Wav2Vec2Processor'''),
        ('''wavlm''', '''Wav2Vec2Processor'''),
        ('''whisper''', '''WhisperProcessor'''),
        ('''xclip''', '''XCLIPProcessor'''),
    ]
)

_SCREAMING_SNAKE_CASE    =			_LazyAutoMapping(CONFIG_MAPPING_NAMES, PROCESSOR_MAPPING_NAMES)




def       _lowerCAmelCase  ( lowerCamelCase_	:     str ):
			for module_name, processors in PROCESSOR_MAPPING_NAMES.items():
						if class_name in processors:
									__lowercase   =							model_type_to_module_name(lowerCamelCase_ )

									__lowercase   =							importlib.import_module(f".{module_name}"						,				'''transformers.models''' )
									try:
												return getattr(lowerCamelCase_						,				lowerCamelCase_ )
									except AttributeError:
												continue

			for processor in PROCESSOR_MAPPING._extra_content.values():
						if getattr(lowerCamelCase_						,				'''__name__'''						,				lowerCamelCase_ ) == class_name:
									return processor

    # We did not fine the class, but maybe it's because a dep is missing. In that case, the class will be in the main
    # init and we return the proper dummy to get an appropriate error message.
			__lowercase   =							importlib.import_module('''transformers''' )
			if hasattr(lowerCamelCase_						,				lowerCamelCase_ ):
						return getattr(lowerCamelCase_						,				lowerCamelCase_ )

			return None






class 							__lowercase				:







	'''simple docstring'''

	def __init__(self       )							->							List[Any]:







				'''simple docstring'''





				raise EnvironmentError(
				    '''AutoProcessor is designed to be instantiated '''
				    '''using the `AutoProcessor.from_pretrained(pretrained_model_name_or_path)` method.'''       )


	@classmethod
	@replace_list_option_in_docstrings(_lowerCamelCase       )
	def   _UpperCAmelCase   (cls   ,_lowerCamelCase   ,**_lowerCamelCase       )							->							Optional[Any]:







				'''simple docstring'''





				__lowercase   =							kwargs.pop('''config'''   ,_lowerCamelCase       )
				__lowercase   =							kwargs.pop('''trust_remote_code'''   ,_lowerCamelCase       )
				__lowercase   =							True

				__lowercase   =							None
				__lowercase   =							None

				# First, let's see if we have a preprocessor config.
				# Filter the kwargs for `get_file_from_repo`.
				__lowercase   =							{
				    key: kwargs[key] for key in inspect.signature(_lowerCamelCase       ).parameters.keys() if key in kwargs
				}
				# Let's start by checking whether the processor class is saved in an image processor
				__lowercase   =							get_file_from_repo(
				    _lowerCamelCase   ,_lowerCamelCase   ,**_lowerCamelCase       )
				if preprocessor_config_file is not None:
							__lowercase  ,						__lowercase   =							ImageProcessingMixin.get_image_processor_dict(_lowerCamelCase   ,**_lowerCamelCase       )
							__lowercase   =							config_dict.get('''processor_class'''   ,_lowerCamelCase       )
							if "AutoProcessor" in config_dict.get('''auto_map'''   ,{}       ):
										__lowercase   =							config_dict['''auto_map''']['''AutoProcessor''']

        # If not found, let's check whether the processor class is saved in a feature extractor config
				if preprocessor_config_file is not None and processor_class is None:
							__lowercase  ,						__lowercase   =							FeatureExtractionMixin.get_feature_extractor_dict(_lowerCamelCase   ,**_lowerCamelCase       )
							__lowercase   =							config_dict.get('''processor_class'''   ,_lowerCamelCase       )
							if "AutoProcessor" in config_dict.get('''auto_map'''   ,{}       ):
										__lowercase   =							config_dict['''auto_map''']['''AutoProcessor''']

				if processor_class is None:
							# Next, let's check whether the processor class is saved in a tokenizer
							__lowercase   =							get_file_from_repo(
							    _lowerCamelCase   ,_lowerCamelCase   ,**_lowerCamelCase       )
							if tokenizer_config_file is not None:
										with open(_lowerCamelCase   ,encoding='''utf-8'''       ) as reader:
													__lowercase   =							json.load(_lowerCamelCase       )

										__lowercase   =							config_dict.get('''processor_class'''   ,_lowerCamelCase       )
										if "AutoProcessor" in config_dict.get('''auto_map'''   ,{}       ):
													__lowercase   =							config_dict['''auto_map''']['''AutoProcessor''']

				if processor_class is None:
							# Otherwise, load config, if it can be loaded.
							if not isinstance(_lowerCamelCase   ,_lowerCamelCase       ):
										__lowercase   =							AutoConfig.from_pretrained(
										    _lowerCamelCase   ,trust_remote_code=_lowerCamelCase   ,**_lowerCamelCase       )

							# And check if the config contains the processor class.
							__lowercase   =							getattr(_lowerCamelCase   ,'''processor_class'''   ,_lowerCamelCase       )
							if hasattr(_lowerCamelCase   ,'''auto_map'''       ) and "AutoProcessor" in config.auto_map:
										__lowercase   =							config.auto_map['''AutoProcessor''']

				if processor_class is not None:
							__lowercase   =							processor_class_from_name(_lowerCamelCase       )

				__lowercase   =							processor_auto_map is not None
				__lowercase   =							processor_class is not None or type(_lowerCamelCase       ) in PROCESSOR_MAPPING
				__lowercase   =							resolve_trust_remote_code(
				    _lowerCamelCase   ,_lowerCamelCase   ,_lowerCamelCase   ,_lowerCamelCase       )

				if has_remote_code and trust_remote_code:
							__lowercase   =							get_class_from_dynamic_module(
							    _lowerCamelCase   ,_lowerCamelCase   ,**_lowerCamelCase       )
							__lowercase   =							kwargs.pop('''code_revision'''   ,_lowerCamelCase       )
							if os.path.isdir(_lowerCamelCase       ):
										processor_class.register_for_auto_class()
							return processor_class.from_pretrained(
							    _lowerCamelCase   ,trust_remote_code=_lowerCamelCase   ,**_lowerCamelCase       )
				elif processor_class is not None:
							return processor_class.from_pretrained(
							    _lowerCamelCase   ,trust_remote_code=_lowerCamelCase   ,**_lowerCamelCase       )
				# Last try: we use the PROCESSOR_MAPPING.
				elif type(_lowerCamelCase       ) in PROCESSOR_MAPPING:
							return PROCESSOR_MAPPING[type(_lowerCamelCase       )].from_pretrained(_lowerCamelCase   ,**_lowerCamelCase       )

				# At this stage, there doesn't seem to be a `Processor` class available for this model, so let's try a
				# tokenizer.
				try:
							return AutoTokenizer.from_pretrained(
							    _lowerCamelCase   ,trust_remote_code=_lowerCamelCase   ,**_lowerCamelCase       )
				except Exception:
							try:
										return AutoImageProcessor.from_pretrained(
										    _lowerCamelCase   ,trust_remote_code=_lowerCamelCase   ,**_lowerCamelCase       )
							except Exception:
										pass

							try:
										return AutoFeatureExtractor.from_pretrained(
										    _lowerCamelCase   ,trust_remote_code=_lowerCamelCase   ,**_lowerCamelCase       )
							except Exception:
										pass

				raise ValueError(
				    f"Unrecognized processing class in {pretrained_model_name_or_path}. Can't instantiate a processor, a "
				    '''tokenizer, an image processor or a feature extractor for this model. Make sure the repository contains'''
				    '''the files of at least one of those processing classes.'''       )


	@staticmethod
	def   _UpperCAmelCase   (_lowerCamelCase   ,_lowerCamelCase       )							->							int:







				'''simple docstring'''





				PROCESSOR_MAPPING.register(_lowerCamelCase   ,_lowerCamelCase       )





