






"""simple docstring"""
from typing import Callable, List, Optional, Union

import PIL
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPSegForImageSegmentation,
    CLIPSegProcessor,
    CLIPTextModel,
    CLIPTokenizer,
)

from diffusers import DiffusionPipeline
from diffusers.configuration_utils import FrozenDict
from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion import StableDiffusionInpaintPipeline
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
from diffusers.utils import deprecate, is_accelerate_available, logging


lowerCamelCase_							      = logging.get_logger(__name__)  # pylint: disable=invalid-name







class      UpperCamelCase_			(__A	):





	def __init__( self   :					Optional[int]      ,      lowerCAmelCase_   :					CLIPSegForImageSegmentation      ,      lowerCAmelCase_   :					CLIPSegProcessor      ,      lowerCAmelCase_   :					AutoencoderKL      ,      lowerCAmelCase_   :					CLIPTextModel      ,      lowerCAmelCase_   :					CLIPTokenizer      ,      lowerCAmelCase_   :					UNetaDConditionModel      ,      lowerCAmelCase_   :					Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]      ,      lowerCAmelCase_   :					StableDiffusionSafetyChecker      ,      lowerCAmelCase_   :					CLIPImageProcessor      ,      )		->     Optional[int]:
				super().__init__()

				if hasattr(scheduler.config      ,      "steps_offset"		) and scheduler.config.steps_offset != 1:
							UpperCAmelCase_	:					Optional[int]     	=   (
							    f"""The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"""
							    f""" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure """
							    "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
							    " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
							    " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
							    " file"
							)
							deprecate("steps_offset!=1"      ,      "1.0.0"      ,      lowerCAmelCase_      ,      standard_warn=lowerCAmelCase_		)
							UpperCAmelCase_	:					str     	=   dict(scheduler.config		)
							UpperCAmelCase_	:					Tuple     	=   1
							UpperCAmelCase_	:					Optional[int]     	=   FrozenDict(lowerCAmelCase_		)

				if hasattr(scheduler.config      ,      "skip_prk_steps"		) and scheduler.config.skip_prk_steps is False:
							UpperCAmelCase_	:					Tuple     	=   (
							    f"""The configuration file of this scheduler: {scheduler} has not set the configuration"""
							    " `skip_prk_steps`. `skip_prk_steps` should be set to True in the configuration file. Please make"
							    " sure to update the config accordingly as not setting `skip_prk_steps` in the config might lead to"
							    " incorrect results in future versions. If you have downloaded this checkpoint from the Hugging Face"
							    " Hub, it would be very nice if you could open a Pull request for the"
							    " `scheduler/scheduler_config.json` file"
							)
							deprecate("skip_prk_steps not set"      ,      "1.0.0"      ,      lowerCAmelCase_      ,      standard_warn=lowerCAmelCase_		)
							UpperCAmelCase_	:					Dict     	=   dict(scheduler.config		)
							UpperCAmelCase_	:					int     	=   True
							UpperCAmelCase_	:					Optional[Any]     	=   FrozenDict(lowerCAmelCase_		)

				if safety_checker is None:
							logger.warning(
							    f"""You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"""
							    " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
							    " results in services or applications open to the public. Both the diffusers team and Hugging Face"
							    " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
							    " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
							    " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."		)

				self.register_modules(
				    segmentation_model=lowerCAmelCase_      ,      segmentation_processor=lowerCAmelCase_      ,      vae=lowerCAmelCase_      ,      text_encoder=lowerCAmelCase_      ,      tokenizer=lowerCAmelCase_      ,      unet=lowerCAmelCase_      ,      scheduler=lowerCAmelCase_      ,      safety_checker=lowerCAmelCase_      ,      feature_extractor=lowerCAmelCase_      ,      )





	def 	_SCREAMING_SNAKE_CASE ( self   :					Dict      ,      lowerCAmelCase_   :					Optional[Union[str, int]] = "auto"		)		->     Tuple:
				if slice_size == "auto":
							# half the attention head size is usually a good trade-off between
							# speed and memory
							UpperCAmelCase_	:					Optional[int]     	=   self.unet.config.attention_head_dim // 2
				self.unet.set_attention_slice(lowerCAmelCase_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Dict		)		->     str:
				self.enable_attention_slicing(lowerCAmelCase_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					int		)		->     Dict:
				if is_accelerate_available():
							from accelerate import cpu_offload
				else:
							raise ImportError("Please install accelerate via `pip install accelerate`"		)

				UpperCAmelCase_	:					Optional[int]     	=   torch.device("cuda"		)

				for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae, self.safety_checker]:
							if cpu_offloaded_model is not None:
										cpu_offload(lowerCAmelCase_      ,      lowerCAmelCase_		)





	@property
	# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
	def 	_SCREAMING_SNAKE_CASE ( self   :					Tuple		)		->     Optional[Any]:
				if self.device != torch.device("meta"		) or not hasattr(self.unet      ,      "_hf_hook"		):
							return self.device
				for module in self.unet.modules():
							if (
							    hasattr(lowerCAmelCase_      ,      "_hf_hook"		)
							    and hasattr(module._hf_hook      ,      "execution_device"		)
							    and module._hf_hook.execution_device is not None
							):
										return torch.device(module._hf_hook.execution_device		)
				return self.device





	@torch.no_grad()
	def __call__( self   :					int      ,      lowerCAmelCase_   :					Union[str, List[str]]      ,      lowerCAmelCase_   :					Union[torch.FloatTensor, PIL.Image.Image]      ,      lowerCAmelCase_   :					str      ,      lowerCAmelCase_   :					int = 512      ,      lowerCAmelCase_   :					int = 512      ,      lowerCAmelCase_   :					int = 50      ,      lowerCAmelCase_   :					float = 7.5      ,      lowerCAmelCase_   :					Optional[Union[str, List[str]]] = None      ,      lowerCAmelCase_   :					Optional[int] = 1      ,      lowerCAmelCase_   :					float = 0.0      ,      lowerCAmelCase_   :					Optional[torch.Generator] = None      ,      lowerCAmelCase_   :					Optional[torch.FloatTensor] = None      ,      lowerCAmelCase_   :					Optional[str] = "pil"      ,      lowerCAmelCase_   :					bool = True      ,      lowerCAmelCase_   :					Optional[Callable[[int, int, torch.FloatTensor], None]] = None      ,      lowerCAmelCase_   :					int = 1      ,      **lowerCAmelCase_   :					Tuple      ,      )		->     Optional[int]:
				UpperCAmelCase_	:					List[str]     	=   self.segmentation_processor(
				    text=[text]      ,      images=[image]      ,      padding="max_length"      ,      return_tensors="pt"		).to(self.device		)
				UpperCAmelCase_	:					Any     	=   self.segmentation_model(**lowerCAmelCase_		)
				UpperCAmelCase_	:					Any     	=   torch.sigmoid(outputs.logits		).cpu().detach().unsqueeze(-1		).numpy()
				UpperCAmelCase_	:					List[Any]     	=   self.numpy_to_pil(lowerCAmelCase_		)[0].resize(image.size		)

				# Run inpainting pipeline with the generated mask
				UpperCAmelCase_	:					str     	=   StableDiffusionInpaintPipeline(
				    vae=self.vae      ,      text_encoder=self.text_encoder      ,      tokenizer=self.tokenizer      ,      unet=self.unet      ,      scheduler=self.scheduler      ,      safety_checker=self.safety_checker      ,      feature_extractor=self.feature_extractor      ,      )
				return inpainting_pipeline(
				    prompt=lowerCAmelCase_      ,      image=lowerCAmelCase_      ,      mask_image=lowerCAmelCase_      ,      height=lowerCAmelCase_      ,      width=lowerCAmelCase_      ,      num_inference_steps=lowerCAmelCase_      ,      guidance_scale=lowerCAmelCase_      ,      negative_prompt=lowerCAmelCase_      ,      num_images_per_prompt=lowerCAmelCase_      ,      eta=lowerCAmelCase_      ,      generator=lowerCAmelCase_      ,      latents=lowerCAmelCase_      ,      output_type=lowerCAmelCase_      ,      return_dict=lowerCAmelCase_      ,      callback=lowerCAmelCase_      ,      callback_steps=lowerCAmelCase_      ,      )





"""simple docstring"""

import itertools
import os
import random
import tempfile
import unittest

import numpy as np
from datasets import load_dataset

from transformers import is_speech_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
   from transformers import WhisperFeatureExtractor

if is_torch_available():
   import torch

SCREAMING_SNAKE_CASE_            =    random.Random()




def      lowercase__    (       lowerCAmelCase					:       int			,					lowerCAmelCase					:       Union[str, Any]=1.0			,					lowerCAmelCase					:       Any=None			,					lowerCAmelCase					:       Optional[Any]=None			)			->   Dict:






    """simple docstring"""




    if rng is None:
        UpperCAmelCase          =       global_rng

    UpperCAmelCase          =       []
    for batch_idx in range(shape[0]			):
        values.append([]			)
        for _ in range(shape[1]			):
            values[-1].append(rng.random() * scale			)

    return values
@require_torch
@require_torchaudio
class   _UpperCAmelCase				(		unittest.TestCase				):




      def __init__(				self ,							lowercase_ ,							lowercase_=7 ,							lowercase_=4_0_0 ,							lowercase_=2_0_0_0 ,							lowercase_=1_0 ,							lowercase_=1_6_0 ,							lowercase_=8 ,							lowercase_=0.0 ,							lowercase_=4_0_0_0 ,							lowercase_=False ,							lowercase_=True ,							)    ->       Optional[int]:
          UpperCAmelCase          =       parent
          UpperCAmelCase          =       batch_size
          UpperCAmelCase          =       min_seq_length
          UpperCAmelCase          =       max_seq_length
          UpperCAmelCase          =       (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
          UpperCAmelCase          =       padding_value
          UpperCAmelCase          =       sampling_rate
          UpperCAmelCase          =       return_attention_mask
          UpperCAmelCase          =       do_normalize
          UpperCAmelCase          =       feature_size
          UpperCAmelCase          =       chunk_length
          UpperCAmelCase          =       hop_length




      def      a_   (				self				)    ->       Dict:
          return {
              "feature_size": self.feature_size,
              "hop_length": self.hop_length,
              "chunk_length": self.chunk_length,
              "padding_value": self.padding_value,
              "sampling_rate": self.sampling_rate,
              "return_attention_mask": self.return_attention_mask,
              "do_normalize": self.do_normalize,
          }




      def      a_   (				self ,							lowercase_=False ,							lowercase_=False				)    ->       Optional[Any]:
          def _flatten(lowercase_				):
              return list(itertools.chain(*lowercase_				)				)

          if equal_length:
              UpperCAmelCase          =       [floats_list((self.max_seq_length, self.feature_size)				) for _ in range(self.batch_size				)]
          else:
              # make sure that inputs increase in size
              UpperCAmelCase          =       [
                  floats_list((x, self.feature_size)				)
                  for x in range(self.min_seq_length ,							self.max_seq_length ,							self.seq_length_diff				)
              ]
          if numpify:
              UpperCAmelCase          =       [np.asarray(lowercase_				) for x in speech_inputs]
          return speech_inputs

@require_torch
@require_torchaudio
class   _UpperCAmelCase				(		SCREAMING_SNAKE_CASE_  ,      unittest.TestCase				):
      __SCREAMING_SNAKE_CASE :				List[Any]				   =						WhisperFeatureExtractor if is_speech_available() else None




      def      a_   (				self				)    ->       Union[str, Any]:
          UpperCAmelCase          =       WhisperFeatureExtractionTester(self				)




      def      a_   (				self				)    ->       List[str]:
          UpperCAmelCase          =       self.feature_extraction_class(**self.feat_extract_dict				)

          with tempfile.TemporaryDirectory() as tmpdirname:
              UpperCAmelCase          =       feat_extract_first.save_pretrained(lowercase_				)[0]
              check_json_file_has_correct_format(lowercase_				)
              UpperCAmelCase          =       self.feature_extraction_class.from_pretrained(lowercase_				)

          UpperCAmelCase          =       feat_extract_first.to_dict()
          UpperCAmelCase          =       feat_extract_second.to_dict()
          UpperCAmelCase          =       feat_extract_first.mel_filters
          UpperCAmelCase          =       feat_extract_second.mel_filters
          self.assertTrue(np.allclose(lowercase_ ,							lowercase_				)				)
          self.assertEqual(lowercase_ ,							lowercase_				)




      def      a_   (				self				)    ->       List[str]:
          UpperCAmelCase          =       self.feature_extraction_class(**self.feat_extract_dict				)

          with tempfile.TemporaryDirectory() as tmpdirname:
              UpperCAmelCase          =       os.path.join(lowercase_ ,							'feat_extract.json'				)
              feat_extract_first.to_json_file(lowercase_				)
              UpperCAmelCase          =       self.feature_extraction_class.from_json_file(lowercase_				)

          UpperCAmelCase          =       feat_extract_first.to_dict()
          UpperCAmelCase          =       feat_extract_second.to_dict()
          UpperCAmelCase          =       feat_extract_first.mel_filters
          UpperCAmelCase          =       feat_extract_second.mel_filters
          self.assertTrue(np.allclose(lowercase_ ,							lowercase_				)				)
          self.assertEqual(lowercase_ ,							lowercase_				)




      def      a_   (				self				)    ->       Tuple:
          # Tests that all call wrap to encode_plus and batch_encode_plus
          UpperCAmelCase          =       self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()				)
          # create three inputs of length 800, 1000, and 1200
          UpperCAmelCase          =       [floats_list((1, x)				)[0] for x in range(8_0_0 ,							1_4_0_0 ,							2_0_0				)]
          UpperCAmelCase          =       [np.asarray(lowercase_				) for speech_input in speech_inputs]

          # Test feature size
          UpperCAmelCase          =       feature_extractor(lowercase_ ,							padding='max_length' ,							return_tensors='np'				).input_features
          self.assertTrue(input_features.ndim == 3				)
          self.assertTrue(input_features.shape[-1] == feature_extractor.nb_max_frames				)
          self.assertTrue(input_features.shape[-2] == feature_extractor.feature_size				)

          # Test not batched input
          UpperCAmelCase          =       feature_extractor(speech_inputs[0] ,							return_tensors='np'				).input_features
          UpperCAmelCase          =       feature_extractor(np_speech_inputs[0] ,							return_tensors='np'				).input_features
          self.assertTrue(np.allclose(lowercase_ ,							lowercase_ ,							atol=1E-3				)				)

          # Test batched
          UpperCAmelCase          =       feature_extractor(lowercase_ ,							return_tensors='np'				).input_features
          UpperCAmelCase          =       feature_extractor(lowercase_ ,							return_tensors='np'				).input_features
          for enc_seq_a, enc_seq_a in zip(lowercase_ ,							lowercase_				):
              self.assertTrue(np.allclose(lowercase_ ,							lowercase_ ,							atol=1E-3				)				)

          # Test 2-D numpy arrays are batched.
          UpperCAmelCase          =       [floats_list((1, x)				)[0] for x in (8_0_0, 8_0_0, 8_0_0)]
          UpperCAmelCase          =       np.asarray(lowercase_				)
          UpperCAmelCase          =       feature_extractor(lowercase_ ,							return_tensors='np'				).input_features
          UpperCAmelCase          =       feature_extractor(lowercase_ ,							return_tensors='np'				).input_features
          for enc_seq_a, enc_seq_a in zip(lowercase_ ,							lowercase_				):
              self.assertTrue(np.allclose(lowercase_ ,							lowercase_ ,							atol=1E-3				)				)

          # Test truncation required
          UpperCAmelCase          =       [floats_list((1, x)				)[0] for x in range(2_0_0 ,							(feature_extractor.n_samples + 5_0_0) ,							2_0_0				)]
          UpperCAmelCase          =       [np.asarray(lowercase_				) for speech_input in speech_inputs]

          UpperCAmelCase          =       [x[: feature_extractor.n_samples] for x in speech_inputs]
          UpperCAmelCase          =       [np.asarray(lowercase_				) for speech_input in speech_inputs_truncated]

          UpperCAmelCase          =       feature_extractor(lowercase_ ,							return_tensors='np'				).input_features
          UpperCAmelCase          =       feature_extractor(lowercase_ ,							return_tensors='np'				).input_features
          for enc_seq_a, enc_seq_a in zip(lowercase_ ,							lowercase_				):
              self.assertTrue(np.allclose(lowercase_ ,							lowercase_ ,							atol=1E-3				)				)




      def      a_   (				self				)    ->       List[Any]:
          import torch

          UpperCAmelCase          =       self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()				)
          UpperCAmelCase          =       np.random.rand(1_0_0 ,							3_2				).astype(np.floataa				)
          UpperCAmelCase          =       np_speech_inputs.tolist()

          for inputs in [py_speech_inputs, np_speech_inputs]:
              UpperCAmelCase          =       feature_extractor.pad([{'input_features': inputs}] ,							return_tensors='np'				)
              self.assertTrue(np_processed.input_features.dtype == np.floataa				)
              UpperCAmelCase          =       feature_extractor.pad([{'input_features': inputs}] ,							return_tensors='pt'				)
              self.assertTrue(pt_processed.input_features.dtype == torch.floataa				)




      def      a_   (				self ,							lowercase_				)    ->       Optional[Any]:
          UpperCAmelCase          =       load_dataset('hf-internal-testing/librispeech_asr_dummy' ,							'clean' ,							split='validation'				)
          # automatic decoding with librispeech
          UpperCAmelCase          =       ds.sort('id'				).select(range(lowercase_				)				)[:num_samples]['audio']

          return [x["array"] for x in speech_samples]




      def      a_   (				self				)    ->       Dict:
          # fmt: off
          UpperCAmelCase          =       torch.tensor(
              [
                  0.1_1_9_3, -0.0_9_4_6, -0.1_0_9_8, -0.0_1_9_6, 0.0_2_2_5, -0.0_6_9_0, -0.1_7_3_6, 0.0_9_5_1,
                  0.0_9_7_1, -0.0_8_1_7, -0.0_7_0_2, 0.0_1_6_2, 0.0_2_6_0, 0.0_0_1_7, -0.0_1_9_2, -0.1_6_7_8,
                  0.0_7_0_9, -0.1_8_6_7, -0.0_6_5_5, -0.0_2_7_4, -0.0_2_3_4, -0.1_8_8_4, -0.0_5_1_6, -0.0_5_5_4,
                  -0.0_2_7_4, -0.1_4_2_5, -0.1_4_2_3, 0.0_8_3_7, 0.0_3_7_7, -0.0_8_5_4
              ]				)
          # fmt: on

          UpperCAmelCase          =       self._load_datasamples(1				)
          UpperCAmelCase          =       WhisperFeatureExtractor()
          UpperCAmelCase          =       feature_extractor(lowercase_ ,							return_tensors='pt'				).input_features
          self.assertEqual(input_features.shape ,							(1, 8_0, 3_0_0_0)				)
          self.assertTrue(torch.allclose(input_features[0, 0, :3_0] ,							lowercase_ ,							atol=1E-4				)				)




      def      a_   (				self				)    ->       Optional[Any]:
          UpperCAmelCase          =       self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()				)
          UpperCAmelCase          =       self._load_datasamples(1				)[0]
          UpperCAmelCase          =       ((audio - audio.min()) / (audio.max() - audio.min())) * 6_5_5_3_5  # Rescale to [0, 65535] to show issue
          UpperCAmelCase          =       feat_extract.zero_mean_unit_var_norm([audio] ,							attention_mask=lowercase_				)[0]

          self.assertTrue(np.all(np.mean(lowercase_				) < 1E-3				)				)
          self.assertTrue(np.all(np.abs(np.var(lowercase_				) - 1				) < 1E-3				)				)



