






"""simple docstring"""







import copy

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


_lowerCAmelCase            =	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      UpperCamelCase		(A__     ):
						_SCREAMING_SNAKE_CASE  : Tuple												=		"""rag"""
						_SCREAMING_SNAKE_CASE  : Optional[Any]												=		True






						def __init__(		self  :str			, __magic_name__  :Any=None			, __magic_name__  :str=True			, __magic_name__  :Optional[Any]=None			, __magic_name__  :Tuple=None			, __magic_name__  :Dict=None			, __magic_name__  :List[Any]=None			, __magic_name__  :Any=None			, __magic_name__  :str=" / "			, __magic_name__  :Tuple=" // "			, __magic_name__  :Any=5			, __magic_name__  :Optional[int]=300			, __magic_name__  :str=768			, __magic_name__  :int=8			, __magic_name__  :Optional[Any]="wiki_dpr"			, __magic_name__  :Tuple="train"			, __magic_name__  :Optional[int]="compressed"			, __magic_name__  :List[Any]=None			, __magic_name__  :Dict=None			, __magic_name__  :Any=False			, __magic_name__  :str=False			, __magic_name__  :Optional[Any]=0.0			, __magic_name__  :str=True			, __magic_name__  :Dict=False			, __magic_name__  :Dict=False			, __magic_name__  :str=False			, __magic_name__  :Dict=True			, __magic_name__  :int=None			, **__magic_name__  :List[Any]			, )      ->int:
										super().__init__(
										    bos_token_id=snake_case_			, pad_token_id=snake_case_			, eos_token_id=snake_case_			, decoder_start_token_id=snake_case_			, forced_eos_token_id=snake_case_			, is_encoder_decoder=snake_case_			, prefix=snake_case_			, vocab_size=snake_case_			, **snake_case_			, )
										assert (
										    "question_encoder" in kwargs and "generator" in kwargs
										), "Config has to be initialized with question_encoder and generator config"
										lowercase    :      List[str]       = kwargs.pop("""question_encoder"""				)
										lowercase    :      Optional[int]       = question_encoder_config.pop("""model_type"""				)
										lowercase    :      List[Any]       = kwargs.pop("""generator"""				)
										lowercase    :      int       = decoder_config.pop("""model_type"""				)

										from ..auto.configuration_auto import AutoConfig

										lowercase    :      List[str]       = AutoConfig.for_model(snake_case_			, **snake_case_				)
										lowercase    :      Optional[Any]       = AutoConfig.for_model(snake_case_			, **snake_case_				)

										lowercase    :      Dict       = reduce_loss
										lowercase    :      int       = label_smoothing
										lowercase    :      Tuple       = exclude_bos_score
										lowercase    :      Optional[int]       = do_marginalize

										lowercase    :      Union[str, Any]       = title_sep
										lowercase    :      List[str]       = doc_sep
										lowercase    :      Optional[int]       = n_docs
										lowercase    :      Any       = max_combined_length

										lowercase    :      List[Any]       = dataset
										lowercase    :      Optional[int]       = dataset_split
										lowercase    :      Any       = index_name

										lowercase    :      Optional[int]       = retrieval_vector_size
										lowercase    :      Any       = retrieval_batch_size
										lowercase    :      str       = passages_path
										lowercase    :      List[str]       = index_path
										lowercase    :      str       = use_dummy_dataset

										lowercase    :      Any       = output_retrieved

										lowercase    :      int       = do_deduplication

										lowercase    :      Optional[Any]       = use_cache

										if self.forced_eos_token_id is None:
														lowercase    :      str       = getattr(self.generator			, """forced_eos_token_id"""			, snake_case_				)






						@classmethod
						def      __snake_case      (		cls  :List[Any]			, __magic_name__  :int			, __magic_name__  :List[str]			, **__magic_name__  :List[Any]				)      ->PretrainedConfig:
										return cls(question_encoder=question_encoder_config.to_dict()			, generator=generator_config.to_dict()			, **snake_case_				)






						def      __snake_case      (		self  :List[str]				)      ->Union[str, Any]:
										lowercase    :      Optional[Any]       = copy.deepcopy(self.__dict__				)
										lowercase    :      List[str]       = self.question_encoder.to_dict()
										lowercase    :      Dict       = self.generator.to_dict()
										lowercase    :      Any       = self.__class__.model_type
										return output


"""simple docstring"""
import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()
class       lowerCAmelCase_		(				A__				,				A__				,				A__				,				unittest.TestCase ):



							'''simple docstring'''





							_snake_case						      =		StableUnCLIPPipeline
							_snake_case						      =		TEXT_TO_IMAGE_PARAMS
							_snake_case						      =		TEXT_TO_IMAGE_BATCH_PARAMS
							_snake_case						      =		TEXT_TO_IMAGE_IMAGE_PARAMS
							_snake_case						      =		TEXT_TO_IMAGE_IMAGE_PARAMS

							# TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
							_snake_case						      =		False




							def    A__						( self				)       ->  Optional[Any]:
													__lowerCAmelCase     =  32
													__lowerCAmelCase     =  embedder_hidden_size

													# prior components

													torch.manual_seed(0				)
													__lowerCAmelCase     =  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)

													torch.manual_seed(0				)
													__lowerCAmelCase     =  CLIPTextModelWithProjection(
													    CLIPTextConfig(
													        bos_token_id=0    ,				eos_token_id=2    ,				hidden_size=snake_case_    ,				projection_dim=snake_case_    ,				intermediate_size=37    ,				layer_norm_eps=1e-0_5    ,				num_attention_heads=4    ,				num_hidden_layers=5    ,				pad_token_id=1    ,				vocab_size=1_000    ,				)				)

													torch.manual_seed(0				)
													__lowerCAmelCase     =  PriorTransformer(
													    num_attention_heads=2    ,				attention_head_dim=12    ,				embedding_dim=snake_case_    ,				num_layers=1    ,				)

													torch.manual_seed(0				)
													__lowerCAmelCase     =  DDPMScheduler(
													    variance_type="""fixed_small_log"""    ,				prediction_type="""sample"""    ,				num_train_timesteps=1_000    ,				clip_sample=snake_case_    ,				clip_sample_range=5.0    ,				beta_schedule="""squaredcos_cap_v2"""    ,				)

													# regular denoising components

													torch.manual_seed(0				)
													__lowerCAmelCase     =  StableUnCLIPImageNormalizer(embedding_dim=snake_case_				)
													__lowerCAmelCase     =  DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""				)

													torch.manual_seed(0				)
													__lowerCAmelCase     =  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)

													torch.manual_seed(0				)
													__lowerCAmelCase     =  CLIPTextModel(
													    CLIPTextConfig(
													        bos_token_id=0    ,				eos_token_id=2    ,				hidden_size=snake_case_    ,				projection_dim=32    ,				intermediate_size=37    ,				layer_norm_eps=1e-0_5    ,				num_attention_heads=4    ,				num_hidden_layers=5    ,				pad_token_id=1    ,				vocab_size=1_000    ,				)				)

													torch.manual_seed(0				)
													__lowerCAmelCase     =  UNetaDConditionModel(
													    sample_size=32    ,				in_channels=4    ,				out_channels=4    ,				down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")    ,				up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")    ,				block_out_channels=(32, 64)    ,				attention_head_dim=(2, 4)    ,				class_embed_type="""projection"""    ,				projection_class_embeddings_input_dim=embedder_projection_dim * 2    ,				cross_attention_dim=snake_case_    ,				layers_per_block=1    ,				upcast_attention=snake_case_    ,				use_linear_projection=snake_case_    ,				)

													torch.manual_seed(0				)
													__lowerCAmelCase     =  DDIMScheduler(
													    beta_schedule="""scaled_linear"""    ,				beta_start=0.00_085    ,				beta_end=0.012    ,				prediction_type="""v_prediction"""    ,				set_alpha_to_one=snake_case_    ,				steps_offset=1    ,				)

													torch.manual_seed(0				)
													__lowerCAmelCase     =  AutoencoderKL()

													__lowerCAmelCase     =  {
													    # prior components
													    """prior_tokenizer""": prior_tokenizer,
													    """prior_text_encoder""": prior_text_encoder,
													    """prior""": prior,
													    """prior_scheduler""": prior_scheduler,
													    # image noising components
													    """image_normalizer""": image_normalizer,
													    """image_noising_scheduler""": image_noising_scheduler,
													    # regular denoising components
													    """tokenizer""": tokenizer,
													    """text_encoder""": text_encoder,
													    """unet""": unet,
													    """scheduler""": scheduler,
													    """vae""": vae,
													}

													return components




							def    A__						( self    ,				snake_case_    ,				snake_case_=0				)       ->  Optional[Any]:
													if str(snake_case_				).startswith("""mps"""				):
																			__lowerCAmelCase     =  torch.manual_seed(snake_case_				)
													else:
																			__lowerCAmelCase     =  torch.Generator(device=snake_case_				).manual_seed(snake_case_				)
													__lowerCAmelCase     =  {
													    """prompt""": """A painting of a squirrel eating a burger""",
													    """generator""": generator,
													    """num_inference_steps""": 2,
													    """prior_num_inference_steps""": 2,
													    """output_type""": """numpy""",
													}
													return inputs




							def    A__						( self				)       ->  Tuple:
													__lowerCAmelCase     =  torch_device == """cpu"""

													self._test_attention_slicing_forward_pass(test_max_difference=snake_case_				)




							def    A__						( self				)       ->  Tuple:
													__lowerCAmelCase     =  torch_device in ["""cpu""", """mps"""]

													self._test_inference_batch_single_identical(test_max_difference=snake_case_				)
@slow
@require_torch_gpu
class       lowerCAmelCase_		(				unittest.TestCase ):



							'''simple docstring'''





							def    A__						( self				)       ->  Optional[int]:
													# clean up the VRAM after each test
													super().tearDown()
													gc.collect()
													torch.cuda.empty_cache()




							def    A__						( self				)       ->  Tuple:
													__lowerCAmelCase     =  load_numpy(
													    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"""				)

													__lowerCAmelCase     =  StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""    ,				torch_dtype=torch.floataa				)
													pipe.to(snake_case_				)
													pipe.set_progress_bar_config(disable=snake_case_				)
													# stable unclip will oom when integration tests are run on a V100,
													# so turn on memory savings
													pipe.enable_attention_slicing()
													pipe.enable_sequential_cpu_offload()

													__lowerCAmelCase     =  torch.Generator(device="""cpu"""				).manual_seed(0				)
													__lowerCAmelCase     =  pipe("""anime turle"""    ,				generator=snake_case_    ,				output_type="""np"""				)

													__lowerCAmelCase     =  output.images[0]

													assert image.shape == (768, 768, 3)

													assert_mean_pixel_difference(snake_case_    ,				snake_case_				)




							def    A__						( self				)       ->  Tuple:
													torch.cuda.empty_cache()
													torch.cuda.reset_max_memory_allocated()
													torch.cuda.reset_peak_memory_stats()

													__lowerCAmelCase     =  StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""    ,				torch_dtype=torch.floataa				)
													__lowerCAmelCase     =  pipe.to(snake_case_				)
													pipe.set_progress_bar_config(disable=snake_case_				)
													pipe.enable_attention_slicing()
													pipe.enable_sequential_cpu_offload()

													__lowerCAmelCase     =  pipe(
													    """anime turtle"""    ,				prior_num_inference_steps=2    ,				num_inference_steps=2    ,				output_type="""np"""    ,				)

													__lowerCAmelCase     =  torch.cuda.max_memory_allocated()
													# make sure that less than 7 GB is allocated
													assert mem_bytes < 7 * 10**9
