



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


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







class       __lowerCAmelCase   (      __UpperCamelCase						):



			def __init__(						self       :  Any			,		A       :  CLIPSegForImageSegmentation			,		A       :  CLIPSegProcessor			,		A       :  AutoencoderKL			,		A       :  CLIPTextModel			,		A       :  CLIPTokenizer			,		A       :  UNetaDConditionModel			,		A       :  Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]			,		A       :  StableDiffusionSafetyChecker			,		A       :  CLIPImageProcessor			,		)       ->   Dict:



						"""simple docstring"""






						super().__init__()

						if hasattr(scheduler.config			,		'steps_offset') and scheduler.config.steps_offset != 1:
									_UpperCAmelCase			=       (
									    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'			,		A			,		standard_warn=A)
									_UpperCAmelCase			=       dict(scheduler.config)
									_UpperCAmelCase			=       1
									_UpperCAmelCase			=       FrozenDict(A)

						if hasattr(scheduler.config			,		'skip_prk_steps') and scheduler.config.skip_prk_steps is False:
									_UpperCAmelCase			=       (
									    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'			,		A			,		standard_warn=A)
									_UpperCAmelCase			=       dict(scheduler.config)
									_UpperCAmelCase			=       True
									_UpperCAmelCase			=       FrozenDict(A)

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



			def    _lowerCamelCase	(						self       :  Tuple			,		A       :  Optional[Union[str, int]] = "auto")       ->   List[str]:



						"""simple docstring"""






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



			def    _lowerCamelCase	(						self       :  Tuple)       ->   str:



						"""simple docstring"""






						self.enable_attention_slicing(A)



			def    _lowerCamelCase	(						self       :  str)       ->   Optional[Any]:



						"""simple docstring"""






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

						_UpperCAmelCase			=       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(A			,		A)



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



						"""simple docstring"""






						_UpperCAmelCase			=       self.segmentation_processor(
						    text=[text]			,		images=[image]			,		padding='max_length'			,		return_tensors='pt').to(self.device)
						_UpperCAmelCase			=       self.segmentation_model(**A)
						_UpperCAmelCase			=       torch.sigmoid(outputs.logits).cpu().detach().unsqueeze(-1).numpy()
						_UpperCAmelCase			=       self.numpy_to_pil(A)[0].resize(image.size)

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





import copy

from ...configuration_utils import PretrainedConfig
from ...utils import add_start_docstrings


UpperCAmelCase__       =	r"\n    [`RagConfig`] stores the configuration of a *RagModel*. Configuration objects inherit from [`PretrainedConfig`] and\n    can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information.\n\n    Args:\n        title_sep (`str`, *optional*, defaults to  `\" / \"`):\n            Separator inserted between the title and the text of the retrieved document when calling [`RagRetriever`].\n        doc_sep (`str`, *optional*, defaults to  `\" // \"`):\n            Separator inserted between the text of the retrieved document and the original input when calling\n            [`RagRetriever`].\n        n_docs (`int`, *optional*, defaults to 5):\n            Number of documents to retrieve.\n        max_combined_length (`int`, *optional*, defaults to 300):\n            Max length of contextualized input returned by [`~RagRetriever.__call__`].\n        retrieval_vector_size (`int`, *optional*, defaults to 768):\n            Dimensionality of the document embeddings indexed by [`RagRetriever`].\n        retrieval_batch_size (`int`, *optional*, defaults to 8):\n            Retrieval batch size, defined as the number of queries issues concurrently to the faiss index encapsulated\n            [`RagRetriever`].\n        dataset (`str`, *optional*, defaults to `\"wiki_dpr\"`):\n            A dataset identifier of the indexed dataset in HuggingFace Datasets (list all available datasets and ids\n            using `datasets.list_datasets()`).\n        dataset_split (`str`, *optional*, defaults to `\"train\"`)\n            Which split of the `dataset` to load.\n        index_name (`str`, *optional*, defaults to `\"compressed\"`)\n            The index name of the index associated with the `dataset`. One can choose between `\"legacy\"`, `\"exact\"` and\n            `\"compressed\"`.\n        index_path (`str`, *optional*)\n            The path to the serialized faiss index on disk.\n        passages_path (`str`, *optional*):\n            A path to text passages compatible with the faiss index. Required if using\n            [`~models.rag.retrieval_rag.LegacyIndex`]\n        use_dummy_dataset (`bool`, *optional*, defaults to `False`)\n            Whether to load a \"dummy\" variant of the dataset specified by `dataset`.\n        label_smoothing (`float`, *optional*, defaults to 0.0):\n            Only relevant if `return_loss` is set to `True`. Controls the `epsilon` parameter value for label smoothing\n            in the loss calculation. If set to 0, no label smoothing is performed.\n        do_marginalize (`bool`, *optional*, defaults to `False`):\n            If `True`, the logits are marginalized over all documents by making use of\n            `torch.nn.functional.log_softmax`.\n        reduce_loss (`bool`, *optional*, defaults to `False`):\n            Whether or not to reduce the NLL loss using the `torch.Tensor.sum` operation.\n        do_deduplication (`bool`, *optional*, defaults to `True`):\n            Whether or not to deduplicate the generations from different context documents for a given input. Has to be\n            set to `False` if used while training with distributed backend.\n        exclude_bos_score (`bool`, *optional*, defaults to `False`):\n            Whether or not to disregard the BOS token when computing the loss.\n        output_retrieved(`bool`, *optional*, defaults to `False`):\n            If set to `True`, `retrieved_doc_embeds`, `retrieved_doc_ids`, `context_input_ids` and\n            `context_attention_mask` are returned. See returned tensors for more detail.\n        use_cache (`bool`, *optional*, defaults to `True`):\n            Whether or not the model should return the last key/values attentions (not used by all models).\n        forced_eos_token_id (`int`, *optional*):\n            The id of the token to force as the last generated token when `max_length` is reached. Usually set to\n            `eos_token_id`.\n"







@add_start_docstrings(A						)
class       __lowerCAmelCase   (      A						):
			UpperCamelCase        =					'''rag'''
			UpperCamelCase        =					True



			def __init__(						self       :  str			,		A       :  int=None			,		A       :  Dict=True			,		A       :  int=None			,		A       :  Any=None			,		A       :  Any=None			,		A       :  Any=None			,		A       :  int=None			,		A       :  str=" / "			,		A       :  int=" // "			,		A       :  Any=5			,		A       :  Optional[Any]=3_00			,		A       :  Optional[Any]=7_68			,		A       :  Optional[int]=8			,		A       :  Union[str, Any]="wiki_dpr"			,		A       :  int="train"			,		A       :  Tuple="compressed"			,		A       :  str=None			,		A       :  Any=None			,		A       :  Tuple=False			,		A       :  List[str]=False			,		A       :  Tuple=0.0			,		A       :  Any=True			,		A       :  int=False			,		A       :  Dict=False			,		A       :  List[str]=False			,		A       :  Optional[Any]=True			,		A       :  Tuple=None			,		**A       :  Optional[Any]			,		)       ->   Optional[Any]:



						"""simple docstring"""






						super().__init__(
						    bos_token_id=A			,		pad_token_id=A			,		eos_token_id=A			,		decoder_start_token_id=A			,		forced_eos_token_id=A			,		is_encoder_decoder=A			,		prefix=A			,		vocab_size=A			,		**A			,		)
						assert (
						    "question_encoder" in kwargs and "generator" in kwargs
						), "Config has to be initialized with question_encoder and generator config"
						_UpperCAmelCase			=       kwargs.pop('question_encoder')
						_UpperCAmelCase			=       question_encoder_config.pop('model_type')
						_UpperCAmelCase			=       kwargs.pop('generator')
						_UpperCAmelCase			=       decoder_config.pop('model_type')

						from ..auto.configuration_auto import AutoConfig

						_UpperCAmelCase			=       AutoConfig.for_model(A			,		**A)
						_UpperCAmelCase			=       AutoConfig.for_model(A			,		**A)

						_UpperCAmelCase			=       reduce_loss
						_UpperCAmelCase			=       label_smoothing
						_UpperCAmelCase			=       exclude_bos_score
						_UpperCAmelCase			=       do_marginalize

						_UpperCAmelCase			=       title_sep
						_UpperCAmelCase			=       doc_sep
						_UpperCAmelCase			=       n_docs
						_UpperCAmelCase			=       max_combined_length

						_UpperCAmelCase			=       dataset
						_UpperCAmelCase			=       dataset_split
						_UpperCAmelCase			=       index_name

						_UpperCAmelCase			=       retrieval_vector_size
						_UpperCAmelCase			=       retrieval_batch_size
						_UpperCAmelCase			=       passages_path
						_UpperCAmelCase			=       index_path
						_UpperCAmelCase			=       use_dummy_dataset

						_UpperCAmelCase			=       output_retrieved

						_UpperCAmelCase			=       do_deduplication

						_UpperCAmelCase			=       use_cache

						if self.forced_eos_token_id is None:
									_UpperCAmelCase			=       getattr(self.generator			,		'forced_eos_token_id'			,		A)



			@classmethod
			def    _lowerCamelCase	(						cls       :  Union[str, Any]			,		A       :  PretrainedConfig			,		A       :  PretrainedConfig			,		**A       :  Optional[Any])       ->   PretrainedConfig:



						"""simple docstring"""






						return cls(question_encoder=question_encoder_config.to_dict()			,		generator=generator_config.to_dict()			,		**A)



			def    _lowerCamelCase	(						self       :  Dict)       ->   Optional[Any]:



						"""simple docstring"""






						_UpperCAmelCase			=       copy.deepcopy(self.__dict__)
						_UpperCAmelCase			=       self.question_encoder.to_dict()
						_UpperCAmelCase			=       self.generator.to_dict()
						_UpperCAmelCase			=       self.__class__.model_type
						return output
