



"""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 	__snake_case	(     _lowercase):




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






         """simple docstring"""






         super().__init__()

         if hasattr(scheduler.config						, '''steps_offset'''		) and scheduler.config.steps_offset != 1:
                _lowerCamelCase						:    Dict         =			(
                    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		)
                _lowerCamelCase						:    List[str]         =			dict(scheduler.config		)
                _lowerCamelCase						:    Union[str, Any]         =			1
                _lowerCamelCase						:    str         =			FrozenDict(__lowerCAmelCase		)

         if hasattr(scheduler.config						, '''skip_prk_steps'''		) and scheduler.config.skip_prk_steps is False:
                _lowerCamelCase						:    str         =			(
                    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		)
                _lowerCamelCase						:    Optional[Any]         =			dict(scheduler.config		)
                _lowerCamelCase						:    Any         =			True
                _lowerCamelCase						:    List[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					:						Tuple						, __lowerCAmelCase					:						Optional[Union[str, int]] = "auto"		):






         """simple docstring"""






         if slice_size == "auto":
                # half the attention head size is usually a good trade-off between
                # speed and memory
                _lowerCamelCase						:    int         =			self.unet.config.attention_head_dim // 2
         self.unet.set_attention_slice(__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         self.enable_attention_slicing(__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         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(__lowerCAmelCase						, __lowerCAmelCase		)




  @property
  # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]		):






         """simple docstring"""






         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					:						str						, __lowerCAmelCase					:						Union[str, List[str]]						, __lowerCAmelCase					:						Union[torch.FloatTensor, PIL.Image.Image]						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						int = 5_1_2						, __lowerCAmelCase					:						int = 5_1_2						, __lowerCAmelCase					:						int = 5_0						, __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					:						List[str]						, ):






         """simple docstring"""






         _lowerCamelCase						:    List[str]         =			self.segmentation_processor(
             text=[text]						, images=[image]						, padding='''max_length'''						, return_tensors='''pt'''		).to(self.device		)
         _lowerCamelCase						:    Optional[int]         =			self.segmentation_model(**__lowerCAmelCase		)
         _lowerCamelCase						:    Optional[int]         =			torch.sigmoid(outputs.logits		).cpu().detach().unsqueeze(-1		).numpy()
         _lowerCamelCase						:    Optional[Any]         =			self.numpy_to_pil(__lowerCAmelCase		)[0].resize(image.size		)

         # Run inpainting pipeline with the generated mask
         _lowerCamelCase						:    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"""






from dataclasses import dataclass
from typing import Dict, Optional, Union

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .attention_processor import AttentionProcessor, AttnProcessor
from .embeddings import TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin






@dataclass
class 	__snake_case	(     _lowercase):
  snake_case__    :       torch.FloatTensor






class 	__snake_case	(     _lowercase  ,					_lowercase):




  @register_to_config
  def __init__(			self					:						int						, __lowerCAmelCase					:						int = 3_2						, __lowerCAmelCase					:						int = 6_4						, __lowerCAmelCase					:						int = 2_0						, __lowerCAmelCase					:						int = 7_6_8						, __lowerCAmelCase					:						Tuple=7_7						, __lowerCAmelCase					:						Tuple=4						, __lowerCAmelCase					:						float = 0.0						, __lowerCAmelCase					:						str = "silu"						, __lowerCAmelCase					:						Optional[str] = None						, __lowerCAmelCase					:						Optional[str] = None						, __lowerCAmelCase					:						Optional[str] = "linear"						, __lowerCAmelCase					:						Optional[str] = "prd"						, __lowerCAmelCase					:						Optional[int] = None						, __lowerCAmelCase					:						Optional[int] = None						, __lowerCAmelCase					:						Optional[int] = None						, ):






         """simple docstring"""






         super().__init__()
         _lowerCamelCase						:    Any         =			num_attention_heads
         _lowerCamelCase						:    Optional[int]         =			attention_head_dim
         _lowerCamelCase						:    Dict         =			num_attention_heads * attention_head_dim
         _lowerCamelCase						:    Optional[Any]         =			additional_embeddings

         _lowerCamelCase						:    List[Any]         =			time_embed_dim or inner_dim
         _lowerCamelCase						:    Optional[Any]         =			embedding_proj_dim or embedding_dim
         _lowerCamelCase						:    Tuple         =			clip_embed_dim or embedding_dim

         _lowerCamelCase						:    int         =			Timesteps(__lowerCAmelCase						, __lowerCAmelCase						, 0		)
         _lowerCamelCase						:    str         =			TimestepEmbedding(__lowerCAmelCase						, __lowerCAmelCase						, out_dim=__lowerCAmelCase						, act_fn=__lowerCAmelCase		)

         _lowerCamelCase						:    Union[str, Any]         =			nn.Linear(__lowerCAmelCase						, __lowerCAmelCase		)

         if embedding_proj_norm_type is None:
                _lowerCamelCase						:    Dict         =			None
         elif embedding_proj_norm_type == "layer":
                _lowerCamelCase						:    Any         =			nn.LayerNorm(__lowerCAmelCase		)
         else:
                raise ValueError(f'''unsupported embedding_proj_norm_type: {embedding_proj_norm_type}'''		)

         _lowerCamelCase						:    List[str]         =			nn.Linear(__lowerCAmelCase						, __lowerCAmelCase		)

         if encoder_hid_proj_type is None:
                _lowerCamelCase						:    Any         =			None
         elif encoder_hid_proj_type == "linear":
                _lowerCamelCase						:    List[str]         =			nn.Linear(__lowerCAmelCase						, __lowerCAmelCase		)
         else:
                raise ValueError(f'''unsupported encoder_hid_proj_type: {encoder_hid_proj_type}'''		)

         _lowerCamelCase						:    Optional[int]         =			nn.Parameter(torch.zeros(1						, num_embeddings + additional_embeddings						, __lowerCAmelCase		)		)

         if added_emb_type == "prd":
                _lowerCamelCase						:    Optional[Any]         =			nn.Parameter(torch.zeros(1						, 1						, __lowerCAmelCase		)		)
         elif added_emb_type is None:
                _lowerCamelCase						:    Union[str, Any]         =			None
         else:
                raise ValueError(
                    f'''`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `\'prd\'` or `None`.'''		)

         _lowerCamelCase						:    Tuple         =			nn.ModuleList(
             [
                 BasicTransformerBlock(
                     __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, dropout=__lowerCAmelCase						, activation_fn='''gelu'''						, attention_bias=__lowerCAmelCase						, )
                 for d in range(__lowerCAmelCase		)
             ]		)

         if norm_in_type == "layer":
                _lowerCamelCase						:    Any         =			nn.LayerNorm(__lowerCAmelCase		)
         elif norm_in_type is None:
                _lowerCamelCase						:    Any         =			None
         else:
                raise ValueError(f'''Unsupported norm_in_type: {norm_in_type}.'''		)

         _lowerCamelCase						:    str         =			nn.LayerNorm(__lowerCAmelCase		)

         _lowerCamelCase						:    List[Any]         =			nn.Linear(__lowerCAmelCase						, __lowerCAmelCase		)

         _lowerCamelCase						:    Optional[int]         =			torch.full(
             [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]						, -1_00_00.0		)
         causal_attention_mask.triu_(1		)
         _lowerCamelCase						:    Dict         =			causal_attention_mask[None, ...]
         self.register_buffer('''causal_attention_mask'''						, __lowerCAmelCase						, persistent=__lowerCAmelCase		)

         _lowerCamelCase						:    int         =			nn.Parameter(torch.zeros(1						, __lowerCAmelCase		)		)
         _lowerCamelCase						:    Tuple         =			nn.Parameter(torch.zeros(1						, __lowerCAmelCase		)		)




  @property
  # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
  def 						SCREAMING_SNAKE_CASE    (			self					:						int		):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			{}

         def fn_recursive_add_processors(__lowerCAmelCase					:						str						, __lowerCAmelCase					:						torch.nn.Module						, __lowerCAmelCase					:						Dict[str, AttentionProcessor]		):
                if hasattr(__lowerCAmelCase						, '''set_processor'''		):
                       _lowerCamelCase						:    str         =			module.processor

                for sub_name, child in module.named_children():
                       fn_recursive_add_processors(f'''{name}.{sub_name}'''						, __lowerCAmelCase						, __lowerCAmelCase		)

                return processors

         for name, module in self.named_children():
                fn_recursive_add_processors(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)

         return processors




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any						, __lowerCAmelCase					:						Union[AttentionProcessor, Dict[str, AttentionProcessor]]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			len(self.attn_processors.keys()		)

         if isinstance(__lowerCAmelCase						, __lowerCAmelCase		) and len(__lowerCAmelCase		) != count:
                raise ValueError(
                    f'''A dict of processors was passed, but the number of processors {len(__lowerCAmelCase		)} does not match the'''
                    f''' number of attention layers: {count}. Please make sure to pass {count} processor classes.'''		)

         def fn_recursive_attn_processor(__lowerCAmelCase					:						str						, __lowerCAmelCase					:						torch.nn.Module						, __lowerCAmelCase					:						Tuple		):
                if hasattr(__lowerCAmelCase						, '''set_processor'''		):
                       if not isinstance(__lowerCAmelCase						, __lowerCAmelCase		):
                              module.set_processor(__lowerCAmelCase		)
                       else:
                              module.set_processor(processor.pop(f'''{name}.processor'''		)		)

                for sub_name, child in module.named_children():
                       fn_recursive_attn_processor(f'''{name}.{sub_name}'''						, __lowerCAmelCase						, __lowerCAmelCase		)

         for name, module in self.named_children():
                fn_recursive_attn_processor(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         self.set_attn_processor(AttnProcessor()		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						str						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						Union[torch.Tensor, float, int]						, __lowerCAmelCase					:						torch.FloatTensor						, __lowerCAmelCase					:						Optional[torch.FloatTensor] = None						, __lowerCAmelCase					:						Optional[torch.BoolTensor] = None						, __lowerCAmelCase					:						bool = True						, ):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			hidden_states.shape[0]

         _lowerCamelCase						:    Union[str, Any]         =			timestep
         if not torch.is_tensor(__lowerCAmelCase		):
                _lowerCamelCase						:    Dict         =			torch.tensor([timesteps]						, dtype=torch.long						, device=hidden_states.device		)
         elif torch.is_tensor(__lowerCAmelCase		) and len(timesteps.shape		) == 0:
                _lowerCamelCase						:    str         =			timesteps[None].to(hidden_states.device		)

         # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
         _lowerCamelCase						:    Tuple         =			timesteps * torch.ones(__lowerCAmelCase						, dtype=timesteps.dtype						, device=timesteps.device		)

         _lowerCamelCase						:    Tuple         =			self.time_proj(__lowerCAmelCase		)

         # timesteps does not contain any weights and will always return f32 tensors
         # but time_embedding might be fp16, so we need to cast here.
         _lowerCamelCase						:    Union[str, Any]         =			timesteps_projected.to(dtype=self.dtype		)
         _lowerCamelCase						:    Any         =			self.time_embedding(__lowerCAmelCase		)

         if self.embedding_proj_norm is not None:
                _lowerCamelCase						:    str         =			self.embedding_proj_norm(__lowerCAmelCase		)

         _lowerCamelCase						:    Optional[Any]         =			self.embedding_proj(__lowerCAmelCase		)
         if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
                _lowerCamelCase						:    Optional[Any]         =			self.encoder_hidden_states_proj(__lowerCAmelCase		)
         elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None:
                raise ValueError('''`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set'''		)

         _lowerCamelCase						:    str         =			self.proj_in(__lowerCAmelCase		)

         _lowerCamelCase						:    Tuple         =			self.positional_embedding.to(hidden_states.dtype		)

         _lowerCamelCase						:    Tuple         =			[]
         _lowerCamelCase						:    Dict         =			0

         if encoder_hidden_states is not None:
                additional_embeds.append(__lowerCAmelCase		)
                additional_embeddings_len += encoder_hidden_states.shape[1]

         if len(proj_embeddings.shape		) == 2:
                _lowerCamelCase						:    str         =			proj_embeddings[:, None, :]

         if len(hidden_states.shape		) == 2:
                _lowerCamelCase						:    Dict         =			hidden_states[:, None, :]

         _lowerCamelCase						:    int         =			additional_embeds + [
             proj_embeddings,
             time_embeddings[:, None, :],
             hidden_states,
         ]

         if self.prd_embedding is not None:
                _lowerCamelCase						:    int         =			self.prd_embedding.to(hidden_states.dtype		).expand(__lowerCAmelCase						, -1						, -1		)
                additional_embeds.append(__lowerCAmelCase		)

         _lowerCamelCase						:    str         =			torch.cat(
             __lowerCAmelCase						, dim=1						, )

         # Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
         _lowerCamelCase						:    Dict         =			additional_embeddings_len + proj_embeddings.shape[1] + 1
         if positional_embeddings.shape[1] < hidden_states.shape[1]:
                _lowerCamelCase						:    List[Any]         =			F.pad(
                    __lowerCAmelCase						, (
                        0,
                        0,
                        additional_embeddings_len,
                        self.prd_embedding.shape[1] if self.prd_embedding is not None else 0,
                    )						, value=0.0						, )

         _lowerCamelCase						:    Any         =			hidden_states + positional_embeddings

         if attention_mask is not None:
                _lowerCamelCase						:    Union[str, Any]         =			(1 - attention_mask.to(hidden_states.dtype		)) * -1_00_00.0
                _lowerCamelCase						:    Tuple         =			F.pad(__lowerCAmelCase						, (0, self.additional_embeddings)						, value=0.0		)
                _lowerCamelCase						:    Any         =			(attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype		)
                _lowerCamelCase						:    int         =			attention_mask.repeat_interleave(self.config.num_attention_heads						, dim=0		)

         if self.norm_in is not None:
                _lowerCamelCase						:    List[str]         =			self.norm_in(__lowerCAmelCase		)

         for block in self.transformer_blocks:
                _lowerCamelCase						:    List[str]         =			block(__lowerCAmelCase						, attention_mask=__lowerCAmelCase		)

         _lowerCamelCase						:    Optional[Any]         =			self.norm_out(__lowerCAmelCase		)

         if self.prd_embedding is not None:
                _lowerCamelCase						:    Any         =			hidden_states[:, -1]
         else:
                _lowerCamelCase						:    Dict         =			hidden_states[:, additional_embeddings_len:]

         _lowerCamelCase						:    str         =			self.proj_to_clip_embeddings(__lowerCAmelCase		)

         if not return_dict:
                return (predicted_image_embedding,)

         return PriorTransformerOutput(predicted_image_embedding=__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]						, __lowerCAmelCase					:						Tuple		):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			(prior_latents * self.clip_std) + self.clip_mean
         return prior_latents


