





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        _lowerCAmelCase    (			__lowerCamelCase							):


  '''simple docstring'''







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


    '''simple docstring'''





    super().__init__()

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

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

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







  def UpperCamelCase_ (      self      :					Dict      ,				UpperCamelCase      :					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
      _snake_case   :				Union[str, Any]														=	self.unet.config.attention_head_dim // 2
    self.unet.set_attention_slice(UpperCamelCase_							)







  def UpperCamelCase_ (      self      :					str							):


    '''simple docstring'''





    self.enable_attention_slicing(UpperCamelCase_							)







  def UpperCamelCase_ (      self      :					int							):


    '''simple docstring'''





    if is_accelerate_available():
      from accelerate import cpu_offload
    else:
      raise ImportError('Please install accelerate via `pip install accelerate`'							)

    _snake_case   :				Union[str, Any]														=	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(UpperCamelCase_      ,				UpperCamelCase_							)







  @property
  # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
  def UpperCamelCase_ (      self      :					List[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(UpperCamelCase_      ,				'_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      :					Union[str, Any]      ,				UpperCamelCase      :					Union[str, List[str]]      ,				UpperCamelCase      :					Union[torch.FloatTensor, PIL.Image.Image]      ,				UpperCamelCase      :					str      ,				UpperCamelCase      :					int = 5_12      ,				UpperCamelCase      :					int = 5_12      ,				UpperCamelCase      :					int = 50      ,				UpperCamelCase      :					float = 7.5      ,				UpperCamelCase      :					Optional[Union[str, List[str]]] = None      ,				UpperCamelCase      :					Optional[int] = 1      ,				UpperCamelCase      :					float = 0.0      ,				UpperCamelCase      :					Optional[torch.Generator] = None      ,				UpperCamelCase      :					Optional[torch.FloatTensor] = None      ,				UpperCamelCase      :					Optional[str] = "pil"      ,				UpperCamelCase      :					bool = True      ,				UpperCamelCase      :					Optional[Callable[[int, int, torch.FloatTensor], None]] = None      ,				UpperCamelCase      :					int = 1      ,				**UpperCamelCase      :					Any      ,				):


    '''simple docstring'''





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

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











from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase_						   =    logging.get_logger(__name__)

lowerCAmelCase_						   =    {
    """MIT/ast-finetuned-audioset-10-10-0.4593""": (
        """https://huggingface.co/MIT/ast-finetuned-audioset-10-10-0.4593/resolve/main/config.json"""
    ),
}



class        _lowerCAmelCase    (			UpperCAmelCase_							):


  '''simple docstring'''
  a_							:       Tuple      					="""audio-spectrogram-transformer"""







  def __init__(      self      :					List[Any]      ,				UpperCamelCase      :					Union[str, Any]=7_68      ,				UpperCamelCase      :					int=12      ,				UpperCamelCase      :					str=12      ,				UpperCamelCase      :					Tuple=30_72      ,				UpperCamelCase      :					Optional[Any]="gelu"      ,				UpperCamelCase      :					Any=0.0      ,				UpperCamelCase      :					Dict=0.0      ,				UpperCamelCase      :					List[Any]=0.02      ,				UpperCamelCase      :					Dict=1e-1_2      ,				UpperCamelCase      :					str=16      ,				UpperCamelCase      :					List[Any]=True      ,				UpperCamelCase      :					Any=10      ,				UpperCamelCase      :					Optional[int]=10      ,				UpperCamelCase      :					int=10_24      ,				UpperCamelCase      :					Optional[Any]=1_28      ,				**UpperCamelCase      :					Optional[Any]      ,				):


    '''simple docstring'''





    super().__init__(**UpperCamelCase							)

    _snake_case   :				Tuple														=	hidden_size
    _snake_case   :				str														=	num_hidden_layers
    _snake_case   :				Optional[Any]														=	num_attention_heads
    _snake_case   :				Optional[Any]														=	intermediate_size
    _snake_case   :				Optional[Any]														=	hidden_act
    _snake_case   :				List[str]														=	hidden_dropout_prob
    _snake_case   :				Union[str, Any]														=	attention_probs_dropout_prob
    _snake_case   :				Any														=	initializer_range
    _snake_case   :				List[str]														=	layer_norm_eps
    _snake_case   :				int														=	patch_size
    _snake_case   :				List[str]														=	qkv_bias
    _snake_case   :				int														=	frequency_stride
    _snake_case   :				List[Any]														=	time_stride
    _snake_case   :				List[Any]														=	max_length
    _snake_case   :				List[str]														=	num_mel_bins




