



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


lowercase__		:				Union[str, Any]          =	logging.get_logger(__name__)  # pylint: disable=invalid-name







class       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      def __init__(      self				:      Tuple    ,			SCREAMING_SNAKE_CASE_				:      CLIPSegForImageSegmentation    ,			SCREAMING_SNAKE_CASE_				:      CLIPSegProcessor    ,			SCREAMING_SNAKE_CASE_				:      AutoencoderKL    ,			SCREAMING_SNAKE_CASE_				:      CLIPTextModel    ,			SCREAMING_SNAKE_CASE_				:      CLIPTokenizer    ,			SCREAMING_SNAKE_CASE_				:      UNetaDConditionModel    ,			SCREAMING_SNAKE_CASE_				:      Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]    ,			SCREAMING_SNAKE_CASE_				:      StableDiffusionSafetyChecker    ,			SCREAMING_SNAKE_CASE_				:      CLIPImageProcessor    ,			):
        super().__init__()

        if hasattr(scheduler.config    ,			'steps_offset'       ) and scheduler.config.steps_offset != 1:
          lowerCAmelCase_   :				str								= (
              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'    ,			SCREAMING_SNAKE_CASE_    ,			standard_warn=SCREAMING_SNAKE_CASE_       )
          lowerCAmelCase_   :				Any								= dict(scheduler.config       )
          lowerCAmelCase_   :				Optional[int]								= 1
          lowerCAmelCase_   :				Any								= FrozenDict(SCREAMING_SNAKE_CASE_       )

        if hasattr(scheduler.config    ,			'skip_prk_steps'       ) and scheduler.config.skip_prk_steps is False:
          lowerCAmelCase_   :				Union[str, Any]								= (
              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'    ,			SCREAMING_SNAKE_CASE_    ,			standard_warn=SCREAMING_SNAKE_CASE_       )
          lowerCAmelCase_   :				Dict								= dict(scheduler.config       )
          lowerCAmelCase_   :				List[Any]								= True
          lowerCAmelCase_   :				Dict								= FrozenDict(SCREAMING_SNAKE_CASE_       )

        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=SCREAMING_SNAKE_CASE_    ,			segmentation_processor=SCREAMING_SNAKE_CASE_    ,			vae=SCREAMING_SNAKE_CASE_    ,			text_encoder=SCREAMING_SNAKE_CASE_    ,			tokenizer=SCREAMING_SNAKE_CASE_    ,			unet=SCREAMING_SNAKE_CASE_    ,			scheduler=SCREAMING_SNAKE_CASE_    ,			safety_checker=SCREAMING_SNAKE_CASE_    ,			feature_extractor=SCREAMING_SNAKE_CASE_    ,			)



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



      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]       ):
        self.enable_attention_slicing(SCREAMING_SNAKE_CASE_       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      List[Any]       ):
        if is_accelerate_available():
          from accelerate import cpu_offload
        else:
          raise ImportError('Please install accelerate via `pip install accelerate`'       )

        lowerCAmelCase_   :				Tuple								= 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(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )



      @property
      # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
      def      SCREAMING_SNAKE_CASE__ (      self				:      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(SCREAMING_SNAKE_CASE_    ,			'_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				:      str    ,			SCREAMING_SNAKE_CASE_				:      Union[str, List[str]]    ,			SCREAMING_SNAKE_CASE_				:      Union[torch.FloatTensor, PIL.Image.Image]    ,			SCREAMING_SNAKE_CASE_				:      str    ,			SCREAMING_SNAKE_CASE_				:      int = 5_1_2    ,			SCREAMING_SNAKE_CASE_				:      int = 5_1_2    ,			SCREAMING_SNAKE_CASE_				:      int = 5_0    ,			SCREAMING_SNAKE_CASE_				:      float = 7.5    ,			SCREAMING_SNAKE_CASE_				:      Optional[Union[str, List[str]]] = None    ,			SCREAMING_SNAKE_CASE_				:      Optional[int] = 1    ,			SCREAMING_SNAKE_CASE_				:      float = 0.0    ,			SCREAMING_SNAKE_CASE_				:      Optional[torch.Generator] = None    ,			SCREAMING_SNAKE_CASE_				:      Optional[torch.FloatTensor] = None    ,			SCREAMING_SNAKE_CASE_				:      Optional[str] = "pil"    ,			SCREAMING_SNAKE_CASE_				:      bool = True    ,			SCREAMING_SNAKE_CASE_				:      Optional[Callable[[int, int, torch.FloatTensor], None]] = None    ,			SCREAMING_SNAKE_CASE_				:      int = 1    ,			**SCREAMING_SNAKE_CASE_				:      int    ,			):
        lowerCAmelCase_   :				str								= self.segmentation_processor(
            text=[text]    ,			images=[image]    ,			padding='max_length'    ,			return_tensors='pt'       ).to(self.device       )
        lowerCAmelCase_   :				Union[str, Any]								= self.segmentation_model(**SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Any								= torch.sigmoid(outputs.logits       ).cpu().detach().unsqueeze(-1       ).numpy()
        lowerCAmelCase_   :				int								= self.numpy_to_pil(SCREAMING_SNAKE_CASE_       )[0].resize(image.size       )

        # Run inpainting pipeline with the generated mask
        lowerCAmelCase_   :				Any								= 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=SCREAMING_SNAKE_CASE_    ,			image=SCREAMING_SNAKE_CASE_    ,			mask_image=SCREAMING_SNAKE_CASE_    ,			height=SCREAMING_SNAKE_CASE_    ,			width=SCREAMING_SNAKE_CASE_    ,			num_inference_steps=SCREAMING_SNAKE_CASE_    ,			guidance_scale=SCREAMING_SNAKE_CASE_    ,			negative_prompt=SCREAMING_SNAKE_CASE_    ,			num_images_per_prompt=SCREAMING_SNAKE_CASE_    ,			eta=SCREAMING_SNAKE_CASE_    ,			generator=SCREAMING_SNAKE_CASE_    ,			latents=SCREAMING_SNAKE_CASE_    ,			output_type=SCREAMING_SNAKE_CASE_    ,			return_dict=SCREAMING_SNAKE_CASE_    ,			callback=SCREAMING_SNAKE_CASE_    ,			callback_steps=SCREAMING_SNAKE_CASE_    ,			)










"""simple docstring"""







from ..utils import DummyObject, requires_backends







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      List[Any]    ,			*SCREAMING_SNAKE_CASE_				:      Optional[int]    ,			**SCREAMING_SNAKE_CASE_				:      Any       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Any    ,			*SCREAMING_SNAKE_CASE_				:      int    ,			**SCREAMING_SNAKE_CASE_				:      List[str]       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      int    ,			*SCREAMING_SNAKE_CASE_				:      Optional[int]    ,			**SCREAMING_SNAKE_CASE_				:      Optional[int]       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Tuple    ,			*SCREAMING_SNAKE_CASE_				:      Tuple    ,			**SCREAMING_SNAKE_CASE_				:      Any       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Union[str, Any]    ,			*SCREAMING_SNAKE_CASE_				:      Dict    ,			**SCREAMING_SNAKE_CASE_				:      Dict       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Dict    ,			*SCREAMING_SNAKE_CASE_				:      List[str]    ,			**SCREAMING_SNAKE_CASE_				:      Optional[Any]       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      int    ,			*SCREAMING_SNAKE_CASE_				:      Optional[int]    ,			**SCREAMING_SNAKE_CASE_				:      List[str]       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      List[Any]    ,			*SCREAMING_SNAKE_CASE_				:      Any    ,			**SCREAMING_SNAKE_CASE_				:      List[Any]       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Union[str, Any]    ,			*SCREAMING_SNAKE_CASE_				:      List[str]    ,			**SCREAMING_SNAKE_CASE_				:      List[str]       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Tuple    ,			*SCREAMING_SNAKE_CASE_				:      List[str]    ,			**SCREAMING_SNAKE_CASE_				:      int       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Tuple    ,			*SCREAMING_SNAKE_CASE_				:      List[str]    ,			**SCREAMING_SNAKE_CASE_				:      str       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      List[Any]    ,			*SCREAMING_SNAKE_CASE_				:      Dict    ,			**SCREAMING_SNAKE_CASE_				:      int       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      List[str]    ,			*SCREAMING_SNAKE_CASE_				:      int    ,			**SCREAMING_SNAKE_CASE_				:      int       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Optional[Any]    ,			*SCREAMING_SNAKE_CASE_				:      str    ,			**SCREAMING_SNAKE_CASE_				:      List[str]       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      List[str]    ,			*SCREAMING_SNAKE_CASE_				:      Optional[int]    ,			**SCREAMING_SNAKE_CASE_				:      Dict       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Any    ,			*SCREAMING_SNAKE_CASE_				:      int    ,			**SCREAMING_SNAKE_CASE_				:      Dict       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      str    ,			*SCREAMING_SNAKE_CASE_				:      Optional[Any]    ,			**SCREAMING_SNAKE_CASE_				:      Any       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Tuple    ,			*SCREAMING_SNAKE_CASE_				:      List[Any]    ,			**SCREAMING_SNAKE_CASE_				:      Union[str, Any]       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Tuple    ,			*SCREAMING_SNAKE_CASE_				:      Any    ,			**SCREAMING_SNAKE_CASE_				:      Optional[Any]       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Optional[int]    ,			*SCREAMING_SNAKE_CASE_				:      int    ,			**SCREAMING_SNAKE_CASE_				:      Dict       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Optional[int]    ,			*SCREAMING_SNAKE_CASE_				:      Tuple    ,			**SCREAMING_SNAKE_CASE_				:      Union[str, Any]       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      int    ,			*SCREAMING_SNAKE_CASE_				:      int    ,			**SCREAMING_SNAKE_CASE_				:      Union[str, Any]       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Union[str, Any]    ,			*SCREAMING_SNAKE_CASE_				:      List[str]    ,			**SCREAMING_SNAKE_CASE_				:      Tuple       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Dict    ,			*SCREAMING_SNAKE_CASE_				:      List[Any]    ,			**SCREAMING_SNAKE_CASE_				:      Tuple       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Union[str, Any]    ,			*SCREAMING_SNAKE_CASE_				:      int    ,			**SCREAMING_SNAKE_CASE_				:      Dict       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Any    ,			*SCREAMING_SNAKE_CASE_				:      Optional[int]    ,			**SCREAMING_SNAKE_CASE_				:      str       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Dict    ,			*SCREAMING_SNAKE_CASE_				:      List[Any]    ,			**SCREAMING_SNAKE_CASE_				:      Tuple       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Dict    ,			*SCREAMING_SNAKE_CASE_				:      List[Any]    ,			**SCREAMING_SNAKE_CASE_				:      List[Any]       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Optional[int]    ,			*SCREAMING_SNAKE_CASE_				:      Tuple    ,			**SCREAMING_SNAKE_CASE_				:      List[Any]       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Dict    ,			*SCREAMING_SNAKE_CASE_				:      Optional[Any]    ,			**SCREAMING_SNAKE_CASE_				:      int       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Optional[int]    ,			*SCREAMING_SNAKE_CASE_				:      List[Any]    ,			**SCREAMING_SNAKE_CASE_				:      Dict       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Optional[Any]    ,			*SCREAMING_SNAKE_CASE_				:      Union[str, Any]    ,			**SCREAMING_SNAKE_CASE_				:      Optional[Any]       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Union[str, Any]    ,			*SCREAMING_SNAKE_CASE_				:      Optional[Any]    ,			**SCREAMING_SNAKE_CASE_				:      int       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Tuple    ,			*SCREAMING_SNAKE_CASE_				:      Any    ,			**SCREAMING_SNAKE_CASE_				:      Any       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Union[str, Any]    ,			*SCREAMING_SNAKE_CASE_				:      Tuple    ,			**SCREAMING_SNAKE_CASE_				:      List[str]       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Tuple    ,			*SCREAMING_SNAKE_CASE_				:      Optional[Any]    ,			**SCREAMING_SNAKE_CASE_				:      str       ):
        requires_backends(cls    ,			['flax']       )







class       UpperCamelCase__   (	metaclass=lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     ["""flax"""]



      def __init__(      self				:      Optional[Any]    ,			*SCREAMING_SNAKE_CASE_				:      Any    ,			**SCREAMING_SNAKE_CASE_				:      int       ):
        requires_backends(self    ,			['flax']       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Tuple    ,			*SCREAMING_SNAKE_CASE_				:      Dict    ,			**SCREAMING_SNAKE_CASE_				:      Optional[int]       ):
        requires_backends(cls    ,			['flax']       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Optional[int]    ,			*SCREAMING_SNAKE_CASE_				:      str    ,			**SCREAMING_SNAKE_CASE_				:      Tuple       ):
        requires_backends(cls    ,			['flax']       )





