


import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.utils.versions import require_version


_snake_case  :	int			      =       logging.getLogger(__name__)

require_version("pytorch_lightning>=1.0.4")

_snake_case  :	str			      =       {
    "base": AutoModel,
    "sequence-classification": AutoModelForSequenceClassification,
    "question-answering": AutoModelForQuestionAnswering,
    "pretraining": AutoModelForPreTraining,
    "token-classification": AutoModelForTokenClassification,
    "language-modeling": AutoModelWithLMHead,
    "summarization": AutoModelForSeqaSeqLM,
    "translation": AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
_snake_case  :	Dict			      =       {
    "linear": get_linear_schedule_with_warmup,
    "cosine": get_cosine_schedule_with_warmup,
    "cosine_w_restarts": get_cosine_with_hard_restarts_schedule_with_warmup,
    "polynomial": get_polynomial_decay_schedule_with_warmup,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
_snake_case  :	int			      =       sorted(arg_to_scheduler.keys())
_snake_case  :	Optional[int]			      =       "{" + ", ".join(arg_to_scheduler_choices) + "}"

class        a	(pl.LightningModule   ):




							"""simple docstring"""


							def __init__(							self :						Dict					,			lowerCamelCase :						argparse.Namespace					,			lowerCamelCase :						List[str]=None					,			lowerCamelCase :						Optional[int]="base"					,			lowerCamelCase :						Any=None					,			lowerCamelCase :						List[Any]=None					,			lowerCamelCase :						List[str]=None					,			**lowerCamelCase :						int					,			)			->					Union[str, Any]:
														super().__init__()
														# TODO: move to self.save_hyperparameters()
														# self.save_hyperparameters()
														# can also expand arguments into trainer signature for easier reading

														self.save_hyperparameters(lowerCamelCase       )
														__snake_case	:    Optional[Any]											=     0
														__snake_case	:    Optional[int]											=     Path(self.hparams.output_dir       )
														__snake_case	:    Optional[int]											=     self.hparams.cache_dir if self.hparams.cache_dir else None
														if config is None:
																					__snake_case	:    List[str]											=     AutoConfig.from_pretrained(
																					    self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path					,			**({"num_labels": num_labels} if num_labels is not None else {})					,			cache_dir=lowerCamelCase					,			**lowerCamelCase					,			)
														else:
																					__snake_case	:    PretrainedConfig											=     config

														__snake_case	:    str											=     ("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
														for p in extra_model_params:
																					if getattr(self.hparams					,			lowerCamelCase					,			lowerCamelCase       ):
																												assert hasattr(self.config					,			lowerCamelCase       ), F'model config doesn\'t have a `{p}` attribute'
																												setattr(self.config					,			lowerCamelCase					,			getattr(self.hparams					,			lowerCamelCase       )       )

														if tokenizer is None:
																					__snake_case	:    List[str]											=     AutoTokenizer.from_pretrained(
																					    self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path					,			cache_dir=lowerCamelCase					,			)
														else:
																					__snake_case	:    PreTrainedTokenizer											=     tokenizer
														__snake_case	:    Tuple											=     MODEL_MODES[mode]
														if model is None:
																					__snake_case	:    Union[str, Any]											=     self.model_type.from_pretrained(
																					    self.hparams.model_name_or_path					,			from_tf=bool(".ckpt" in self.hparams.model_name_or_path       )					,			config=self.config					,			cache_dir=lowerCamelCase					,			)
														else:
																					__snake_case	:    Optional[int]											=     model

							def 					__snake_case      (							self :						Tuple					,			*lowerCamelCase :						Dict					,			**lowerCamelCase :						List[Any]       )			->					Optional[int]:
														__snake_case	:    List[Any]											=     self.model_type.from_pretrained(*lowerCamelCase					,			**lowerCamelCase       )

							def 					__snake_case      (							self :						List[str]       )			->					int:
														__snake_case	:    List[str]											=     arg_to_scheduler[self.hparams.lr_scheduler]
														__snake_case	:    Any											=     get_schedule_func(
														    self.opt					,			num_warmup_steps=self.hparams.warmup_steps					,			num_training_steps=self.total_steps()       )
														__snake_case	:    Union[str, Any]											=     {"scheduler": scheduler, "interval": "step", "frequency": 1}
														return scheduler

							def 					__snake_case      (							self :						List[Any]       )			->					Dict:
														__snake_case	:    Optional[int]											=     self.model
														__snake_case	:    Dict											=     ["bias", "LayerNorm.weight"]
														__snake_case	:    Tuple											=     [
														    {
														        "params": [
														            p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay       )
														        ],  # check this named paramters
														        "weight_decay": self.hparams.weight_decay,
														    },
														    {
														        "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay       )],
														        "weight_decay": 0.0,
														    },
														]
														if self.hparams.adafactor:
																					__snake_case	:    Any											=     Adafactor(
																					    lowerCamelCase					,			lr=self.hparams.learning_rate					,			scale_parameter=lowerCamelCase					,			relative_step=lowerCamelCase       )

														else:
																					__snake_case	:    Dict											=     AdamW(
																					    lowerCamelCase					,			lr=self.hparams.learning_rate					,			eps=self.hparams.adam_epsilon       )
														__snake_case	:    Dict											=     optimizer

														__snake_case	:    int											=     self.get_lr_scheduler()

														return [optimizer], [scheduler]

							def 					__snake_case      (							self :						Any					,			lowerCamelCase :						int					,			lowerCamelCase :						Any       )			->					Any:
														return self.validation_step(lowerCamelCase					,			lowerCamelCase       )

							def 					__snake_case      (							self :						Tuple					,			lowerCamelCase :						int       )			->					List[str]:
														return self.validation_end(lowerCamelCase       )

							def 					__snake_case      (							self :						str       )			->					int:
														__snake_case	:    Dict											=     max(1					,			self.hparams.gpus       )  # TODO: consider num_tpu_cores
														__snake_case	:    Optional[int]											=     self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
														return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs

							def 					__snake_case      (							self :						Union[str, Any]					,			lowerCamelCase :						Union[str, Any]       )			->					Dict:
														if stage == "test":
																					__snake_case	:    Any											=     len(self.test_dataloader().dataset       )
														else:
																					__snake_case	:    Dict											=     self.get_dataloader("train"					,			self.hparams.train_batch_size					,			shuffle=lowerCamelCase       )
																					__snake_case	:    Any											=     len(self.train_dataloader().dataset       )

							def 					__snake_case      (							self :						Tuple					,			lowerCamelCase :						str					,			lowerCamelCase :						int					,			lowerCamelCase :						bool = False       )			->					str:
														raise NotImplementedError("You must implement this for your task"       )

							def 					__snake_case      (							self :						List[Any]       )			->					int:
														return self.train_loader

							def 					__snake_case      (							self :						Any       )			->					List[Any]:
														return self.get_dataloader("dev"					,			self.hparams.eval_batch_size					,			shuffle=lowerCamelCase       )

							def 					__snake_case      (							self :						int       )			->					Tuple:
														return self.get_dataloader("test"					,			self.hparams.eval_batch_size					,			shuffle=lowerCamelCase       )

							def 					__snake_case      (							self :						Any					,			lowerCamelCase :						Optional[int]       )			->					List[str]:
														return os.path.join(
														    self.hparams.data_dir					,			"cached_{}_{}_{}".format(
														        lowerCamelCase					,			list(filter(lowerCamelCase					,			self.hparams.model_name_or_path.split("/"       )       )       ).pop()					,			str(self.hparams.max_seq_length       )					,			)					,			)

							@pl.utilities.rank_zero_only
							def 					__snake_case      (							self :						str					,			lowerCamelCase :						Dict[str, Any]       )			->					None:
														__snake_case	:    Optional[Any]											=     self.output_dir.joinpath("best_tfmr"       )
														__snake_case	:    Tuple											=     self.step_count
														self.model.save_pretrained(lowerCamelCase       )
														self.tokenizer.save_pretrained(lowerCamelCase       )






							@staticmethod
							def 					__snake_case      (							lowerCamelCase :						int					,			lowerCamelCase :						Optional[Any]       )			->					Union[str, Any]:
														parser.add_argument(
														    "--model_name_or_path"					,			default=lowerCamelCase					,			type=lowerCamelCase					,			required=lowerCamelCase					,			help="Path to pretrained model or model identifier from huggingface.co/models"					,			)
														parser.add_argument(
														    "--config_name"					,			default=""					,			type=lowerCamelCase					,			help="Pretrained config name or path if not the same as model_name"       )
														parser.add_argument(
														    "--tokenizer_name"					,			default=lowerCamelCase					,			type=lowerCamelCase					,			help="Pretrained tokenizer name or path if not the same as model_name"					,			)
														parser.add_argument(
														    "--cache_dir"					,			default=str(Path(lowerCamelCase       ).parent / "test_run" / "cache"       )					,			type=lowerCamelCase					,			help="Where do you want to store the pre-trained models downloaded from huggingface.co"					,			)
														parser.add_argument(
														    "--encoder_layerdrop"					,			type=lowerCamelCase					,			help="Encoder layer dropout probability (Optional). Goes into model.config"					,			)
														parser.add_argument(
														    "--decoder_layerdrop"					,			type=lowerCamelCase					,			help="Decoder layer dropout probability (Optional). Goes into model.config"					,			)
														parser.add_argument(
														    "--dropout"					,			type=lowerCamelCase					,			help="Dropout probability (Optional). Goes into model.config"					,			)
														parser.add_argument(
														    "--attention_dropout"					,			type=lowerCamelCase					,			help="Attention dropout probability (Optional). Goes into model.config"					,			)
														parser.add_argument("--learning_rate"					,			default=5E-5					,			type=lowerCamelCase					,			help="The initial learning rate for Adam."       )
														parser.add_argument(
														    "--lr_scheduler"					,			default="linear"					,			choices=lowerCamelCase					,			metavar=lowerCamelCase					,			type=lowerCamelCase					,			help="Learning rate scheduler"					,			)
														parser.add_argument("--weight_decay"					,			default=0.0					,			type=lowerCamelCase					,			help="Weight decay if we apply some."       )
														parser.add_argument("--adam_epsilon"					,			default=1E-8					,			type=lowerCamelCase					,			help="Epsilon for Adam optimizer."       )
														parser.add_argument("--warmup_steps"					,			default=0					,			type=lowerCamelCase					,			help="Linear warmup over warmup_steps."       )
														parser.add_argument("--num_workers"					,			default=4					,			type=lowerCamelCase					,			help="kwarg passed to DataLoader"       )
														parser.add_argument("--num_train_epochs"					,			dest="max_epochs"					,			default=3					,			type=lowerCamelCase       )
														parser.add_argument("--train_batch_size"					,			default=32					,			type=lowerCamelCase       )
														parser.add_argument("--eval_batch_size"					,			default=32					,			type=lowerCamelCase       )
														parser.add_argument("--adafactor"					,			action="store_true"       )

class        a	(pl.Callback   ):




							"""simple docstring"""


							def 					__snake_case      (							self :						int					,			lowerCamelCase :						Optional[Any]					,			lowerCamelCase :						List[str]       )			->					Optional[Any]:
														if (
														    trainer.is_global_zero and trainer.global_rank == 0
														):  # we initialize the retriever only on master worker with RAY. In new pytorch-lightning accelorators are removed.
																					pl_module.model.rag.retriever.init_retrieval()  # better to use hook functions.

class        a	(pl.Callback   ):




							"""simple docstring"""


							def 					__snake_case      (							self :						int					,			lowerCamelCase :						int					,			lowerCamelCase :						Dict       )			->					Optional[Any]:
														# print(pl_module.model.rag)
														for name, param in pl_module.model.rag.named_parameters():
																					if param.grad is None:
																												print(lowerCamelCase       )

class        a	(pl.Callback   ):




							"""simple docstring"""


							def 					__snake_case      (							self :						List[Any]					,			lowerCamelCase :						Union[str, Any]					,			lowerCamelCase :						Dict       )			->					Union[str, Any]:
														__snake_case	:    Optional[int]											=     trainer.lr_schedulers[0]["scheduler"]
														__snake_case	:    str											=     {F'lr_group_{i}': lr for i, lr in enumerate(lr_scheduler.get_lr()       )}
														pl_module.logger.log_metrics(lowerCamelCase       )

							def 					__snake_case      (							self :						List[str]					,			lowerCamelCase :						pl.Trainer					,			lowerCamelCase :						pl.LightningModule       )			->					str:
														rank_zero_info("***** Validation results *****"       )
														__snake_case	:    List[str]											=     trainer.callback_metrics
														# Log results
														for key in sorted(lowerCamelCase       ):
																					if key not in ["log", "progress_bar"]:
																												rank_zero_info("{} = {}\n".format(lowerCamelCase					,			str(metrics[key]       )       )       )






							def 					__snake_case      (							self :						Optional[Any]					,			lowerCamelCase :						pl.Trainer					,			lowerCamelCase :						pl.LightningModule       )			->					Optional[int]:
														rank_zero_info("***** Test results *****"       )
														__snake_case	:    Any											=     trainer.callback_metrics
														# Log and save results to file
														__snake_case	:    Dict											=     os.path.join(pl_module.hparams.output_dir					,			"test_results.txt"       )
														with open(lowerCamelCase					,			"w"       ) as writer:
																					for key in sorted(lowerCamelCase       ):
																												if key not in ["log", "progress_bar"]:
																																			rank_zero_info("{} = {}\n".format(lowerCamelCase					,			str(metrics[key]       )       )       )
																																			writer.write("{} = {}\n".format(lowerCamelCase					,			str(metrics[key]       )       )       )







def        lowerCAmelCase_  (      __lowerCamelCase							,			__lowerCamelCase			):
							#  To allow all pl args uncomment the following line
							#  parser = pl.Trainer.add_argparse_args(parser)
							parser.add_argument(
							    "--output_dir"							,			default=str(Path(__lowerCamelCase			).parent / "test_run" / "model_checkpoints"			)							,			type=__lowerCamelCase							,			help="The output directory where the model predictions and checkpoints will be written."							,			)
							parser.add_argument(
							    "--fp16"							,			action="store_true"							,			help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit"							,			)

							parser.add_argument(
							    "--fp16_opt_level"							,			type=__lowerCamelCase							,			default="O2"							,			help=(
							        "For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
							        "See details at https://nvidia.github.io/apex/amp.html"
							    )							,			)
							parser.add_argument("--n_tpu_cores"							,			dest="tpu_cores"							,			type=__lowerCamelCase			)
							parser.add_argument("--max_grad_norm"							,			dest="gradient_clip_val"							,			default=1.0							,			type=__lowerCamelCase							,			help="Max gradient norm"			)
							parser.add_argument("--do_train"							,			action="store_true"							,			help="Whether to run training."			)
							parser.add_argument("--do_predict"							,			action="store_true"							,			help="Whether to run predictions on the test set."			)
							parser.add_argument(
							    "--gradient_accumulation_steps"							,			dest="accumulate_grad_batches"							,			type=__lowerCamelCase							,			default=1							,			help="Number of updates steps to accumulate before performing a backward/update pass."							,			)
							parser.add_argument("--seed"							,			type=__lowerCamelCase							,			default=4_2							,			help="random seed for initialization"			)
							parser.add_argument(
							    "--data_dir"							,			default=str(Path(__lowerCamelCase			).parent / "test_run" / "dummy-train-data"			)							,			type=__lowerCamelCase							,			help="The input data dir. Should contain the training files for the CoNLL-2003 NER task."							,			)







def        lowerCAmelCase_  (      __lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase=None							,			__lowerCamelCase=True							,			__lowerCamelCase=[]							,			__lowerCamelCase=None							,			__lowerCamelCase=None							,			**__lowerCamelCase							,			):
							pl.seed_everything(args.seed			)

							# init model
							__snake_case	:    Optional[int]											=     Path(model.hparams.output_dir			)
							odir.mkdir(exist_ok=__lowerCamelCase			)

							# add custom checkpoints
							if checkpoint_callback is None:
														__snake_case	:    Optional[int]											=     pl.callbacks.ModelCheckpoint(
														    filepath=args.output_dir							,			prefix="checkpoint"							,			monitor="val_loss"							,			mode="min"							,			save_top_k=1			)
							if early_stopping_callback:
														extra_callbacks.append(__lowerCamelCase			)
							if logging_callback is None:
														__snake_case	:    Tuple											=     LoggingCallback()

							__snake_case	:    Tuple											=     {}

							if args.fpaa:
														__snake_case	:    List[str]											=     1_6

							if args.gpus > 1:
														__snake_case	:    Tuple											=     "auto"
														__snake_case	:    List[Any]											=     "ddp"

							__snake_case	:    Union[str, Any]											=     args.accumulate_grad_batches
							__snake_case	:    str											=     None
							__snake_case	:    Optional[Any]											=     "auto"

							__snake_case	:    str											=     pl.Trainer.from_argparse_args(
							    __lowerCamelCase							,			weights_summary=__lowerCamelCase							,			callbacks=[logging_callback] + extra_callbacks + [InitCallback()] + [checkpoint_callback]							,			logger=__lowerCamelCase							,			val_check_interval=1							,			num_sanity_val_steps=2							,			**__lowerCamelCase							,			)

							if args.do_train:
														trainer.fit(__lowerCamelCase			)

							else:
														print("RAG modeling tests with new set functions successfuly executed!"			)
							return trainer





import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class        a	(_lowerCAmelCase  ,       _lowerCAmelCase  ,       unittest.TestCase   ):




							"""simple docstring"""


							__UpperCAmelCase					:							Optional[Any]  						=						StableDiffusionXLImgaImgPipeline
							__UpperCAmelCase					:							Any  						=						TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
							__UpperCAmelCase					:							Dict  						=						PipelineTesterMixin.required_optional_params - {"latents"}
							__UpperCAmelCase					:							Union[str, Any]  						=						TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
							__UpperCAmelCase					:							Tuple  						=						IMAGE_TO_IMAGE_IMAGE_PARAMS
							__UpperCAmelCase					:							str  						=						IMAGE_TO_IMAGE_IMAGE_PARAMS

							def 					__snake_case      (							self :						Optional[Any]       )			->					Tuple:
														torch.manual_seed(0       )
														__snake_case	:    Optional[int]											=     UNetaDConditionModel(
														    block_out_channels=(32, 64)					,			layers_per_block=2					,			sample_size=32					,			in_channels=4					,			out_channels=4					,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")					,			up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")					,			attention_head_dim=(2, 4)					,			use_linear_projection=lowerCamelCase					,			addition_embed_type="text_time"					,			addition_time_embed_dim=8					,			transformer_layers_per_block=(1, 2)					,			projection_class_embeddings_input_dim=80					,			cross_attention_dim=64					,			)
														__snake_case	:    Tuple											=     EulerDiscreteScheduler(
														    beta_start=0.0_00_85					,			beta_end=0.0_12					,			steps_offset=1					,			beta_schedule="scaled_linear"					,			timestep_spacing="leading"					,			)
														torch.manual_seed(0       )
														__snake_case	:    int											=     AutoencoderKL(
														    block_out_channels=[32, 64]					,			in_channels=3					,			out_channels=3					,			down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]					,			up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]					,			latent_channels=4					,			sample_size=128					,			)
														torch.manual_seed(0       )
														__snake_case	:    Optional[int]											=     CLIPTextConfig(
														    bos_token_id=0					,			eos_token_id=2					,			hidden_size=32					,			intermediate_size=37					,			layer_norm_eps=1E-05					,			num_attention_heads=4					,			num_hidden_layers=5					,			pad_token_id=1					,			vocab_size=1000					,			hidden_act="gelu"					,			projection_dim=32					,			)
														__snake_case	:    List[str]											=     CLIPTextModel(lowerCamelCase       )
														__snake_case	:    Any											=     CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"					,			local_files_only=lowerCamelCase       )

														__snake_case	:    List[str]											=     CLIPTextModelWithProjection(lowerCamelCase       )
														__snake_case	:    List[Any]											=     CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"					,			local_files_only=lowerCamelCase       )

														__snake_case	:    Optional[int]											=     {
														    "unet": unet,
														    "scheduler": scheduler,
														    "vae": vae,
														    "text_encoder": text_encoder,
														    "tokenizer": tokenizer,
														    "text_encoder_2": text_encoder_a,
														    "tokenizer_2": tokenizer_a,
														    # "safety_checker": None,
														    # "feature_extractor": None,
														}
														return components

							def 					__snake_case      (							self :						Optional[Any]					,			lowerCamelCase :						List[Any]					,			lowerCamelCase :						List[Any]=0       )			->					Union[str, Any]:
														__snake_case	:    int											=     floats_tensor((1, 3, 32, 32)					,			rng=random.Random(lowerCamelCase       )       ).to(lowerCamelCase       )
														__snake_case	:    Any											=     image / 2 + 0.5
														if str(lowerCamelCase       ).startswith("mps"       ):
																					__snake_case	:    Dict											=     torch.manual_seed(lowerCamelCase       )
														else:
																					__snake_case	:    int											=     torch.Generator(device=lowerCamelCase       ).manual_seed(lowerCamelCase       )
														__snake_case	:    Optional[Any]											=     {
														    "prompt": "A painting of a squirrel eating a burger",
														    "image": image,
														    "generator": generator,
														    "num_inference_steps": 2,
														    "guidance_scale": 5.0,
														    "output_type": "numpy",
														    "strength": 0.75,
														}
														return inputs

							def 					__snake_case      (							self :						Dict       )			->					Any:
														__snake_case	:    Dict											=     "cpu"  # ensure determinism for the device-dependent torch.Generator
														__snake_case	:    Any											=     self.get_dummy_components()
														__snake_case	:    int											=     StableDiffusionXLImgaImgPipeline(**lowerCamelCase       )
														__snake_case	:    List[str]											=     sd_pipe.to(lowerCamelCase       )
														sd_pipe.set_progress_bar_config(disable=lowerCamelCase       )

														__snake_case	:    Tuple											=     self.get_dummy_inputs(lowerCamelCase       )
														__snake_case	:    Dict											=     sd_pipe(**lowerCamelCase       ).images
														__snake_case	:    Dict											=     image[0, -3:, -3:, -1]

														assert image.shape == (1, 32, 32, 3)

														__snake_case	:    str											=     np.array([0.46_56, 0.48_40, 0.44_39, 0.66_98, 0.55_74, 0.45_24, 0.57_99, 0.59_43, 0.51_65]       )

														assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2

							def 					__snake_case      (							self :						str       )			->					Optional[Any]:
														super().test_attention_slicing_forward_pass(expected_max_diff=3E-3       )

							def 					__snake_case      (							self :						Any       )			->					Dict:
														super().test_inference_batch_single_identical(expected_max_diff=3E-3       )

							def 					__snake_case      (							self :						str       )			->					Optional[int]:
														pass






							def 					__snake_case      (							self :						Tuple       )			->					Union[str, Any]:
														__snake_case	:    str											=     self.get_dummy_components()
														__snake_case	:    List[Any]											=     StableDiffusionXLImgaImgPipeline(**lowerCamelCase       )
														__snake_case	:    Optional[Any]											=     sd_pipe.to(lowerCamelCase       )
														__snake_case	:    int											=     sd_pipe.to(lowerCamelCase       )
														sd_pipe.set_progress_bar_config(disable=lowerCamelCase       )

														# forward without prompt embeds
														__snake_case	:    List[str]											=     self.get_dummy_inputs(lowerCamelCase       )
														__snake_case	:    str											=     3 * ["this is a negative prompt"]
														__snake_case	:    Any											=     negative_prompt
														__snake_case	:    Optional[Any]											=     3 * [inputs["prompt"]]

														__snake_case	:    int											=     sd_pipe(**lowerCamelCase       )
														__snake_case	:    List[Any]											=     output.images[0, -3:, -3:, -1]

														# forward with prompt embeds
														__snake_case	:    List[Any]											=     self.get_dummy_inputs(lowerCamelCase       )
														__snake_case	:    Optional[Any]											=     3 * ["this is a negative prompt"]
														__snake_case	:    int											=     3 * [inputs.pop("prompt"       )]

														(
														    (
														    __snake_case
														)				,    (
														    __snake_case
														)				,    (
														    __snake_case
														)				,    (
														    __snake_case
														)				,    
														)	:    Dict											=     sd_pipe.encode_prompt(lowerCamelCase					,			negative_prompt=lowerCamelCase       )

														__snake_case	:    Tuple											=     sd_pipe(
														    **lowerCamelCase					,			prompt_embeds=lowerCamelCase					,			negative_prompt_embeds=lowerCamelCase					,			pooled_prompt_embeds=lowerCamelCase					,			negative_pooled_prompt_embeds=lowerCamelCase					,			)
														__snake_case	:    List[str]											=     output.images[0, -3:, -3:, -1]

														# make sure that it's equal
														assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()       ).max() < 1E-4

@slow
@require_torch_gpu
class        a	(unittest.TestCase   ):




							"""simple docstring"""


							def 					__snake_case      (							self :						Optional[int]       )			->					Dict:
														super().tearDown()
														gc.collect()
														torch.cuda.empty_cache()

							def 					__snake_case      (							self :						Optional[Any]					,			lowerCamelCase :						Any					,			lowerCamelCase :						Optional[Any]="cpu"					,			lowerCamelCase :						str=torch.floataa					,			lowerCamelCase :						int=0       )			->					Dict:
														__snake_case	:    int											=     torch.Generator(device=lowerCamelCase       ).manual_seed(lowerCamelCase       )
														__snake_case	:    Optional[Any]											=     np.random.RandomState(lowerCamelCase       ).standard_normal((1, 4, 64, 64)       )
														__snake_case	:    Optional[Any]											=     torch.from_numpy(lowerCamelCase       ).to(device=lowerCamelCase					,			dtype=lowerCamelCase       )
														__snake_case	:    List[str]											=     {
														    "prompt": "a photograph of an astronaut riding a horse",
														    "latents": latents,
														    "generator": generator,
														    "num_inference_steps": 3,
														    "guidance_scale": 7.5,
														    "output_type": "numpy",
														}
														return inputs






							def 					__snake_case      (							self :						str       )			->					Any:
														__snake_case	:    List[str]											=     DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-base"       )
														pipe.to(lowerCamelCase       )
														pipe.set_progress_bar_config(disable=lowerCamelCase       )

														__snake_case	:    int											=     self.get_inputs(lowerCamelCase       )
														__snake_case	:    Optional[Any]											=     pipe(**lowerCamelCase       ).images
														__snake_case	:    Any											=     image[0, -3:, -3:, -1].flatten()

														assert image.shape == (1, 512, 512, 3)
														__snake_case	:    Optional[int]											=     np.array([0.4_94_93, 0.4_78_96, 0.4_07_98, 0.5_42_14, 0.5_32_12, 0.4_82_02, 0.4_76_56, 0.4_63_29, 0.4_85_06]       )
														assert np.abs(image_slice - expected_slice       ).max() < 7E-3

