

'''simple docstring'''





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


_SCREAMING_SNAKE_CASE  :						Any         =    logging.getLogger(__name__)

require_version('''pytorch_lightning>=1.0.4''')

_SCREAMING_SNAKE_CASE  :						int         =    {
    '''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
_SCREAMING_SNAKE_CASE  :						List[str]         =    {
    '''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
}
_SCREAMING_SNAKE_CASE  :						Tuple         =    sorted(arg_to_scheduler.keys())
_SCREAMING_SNAKE_CASE  :						int         =    '''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''






class  UpperCAmelCase__	( pl.LightningModule ):






  """simple docstring"""


  def __init__(							self		:	Optional[int]			,     __lowerCamelCase		:	argparse.Namespace			,     __lowerCamelCase		:	List[str]=None			,     __lowerCamelCase		:	Any="base"			,     __lowerCamelCase		:	List[str]=None			,     __lowerCamelCase		:	Dict=None			,     __lowerCamelCase		:	Optional[int]=None			,     **__lowerCamelCase		:	Optional[Any]			,     )	->							Optional[int]:
       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      )
       SCREAMING_SNAKE_CASE__					=					0
       SCREAMING_SNAKE_CASE__					=					Path(self.hparams.output_dir      )
       SCREAMING_SNAKE_CASE__					=					self.hparams.cache_dir if self.hparams.cache_dir else None
       if config is None:
            SCREAMING_SNAKE_CASE__					=					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:
            SCREAMING_SNAKE_CASE__					=					config

       SCREAMING_SNAKE_CASE__					=					('''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:
            SCREAMING_SNAKE_CASE__					=					AutoTokenizer.from_pretrained(
                self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path			,     cache_dir=__lowerCamelCase			,     )
       else:
            SCREAMING_SNAKE_CASE__					=					tokenizer
       SCREAMING_SNAKE_CASE__					=					MODEL_MODES[mode]
       if model is None:
            SCREAMING_SNAKE_CASE__					=					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:
            SCREAMING_SNAKE_CASE__					=					model






  def    lowercase_	(							self		:	Tuple			,     *__lowerCamelCase		:	List[Any]			,     **__lowerCamelCase		:	List[str]      )	->							Optional[int]:
       SCREAMING_SNAKE_CASE__					=					self.model_type.from_pretrained(*__lowerCamelCase			,     **__lowerCamelCase      )






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






  def    lowercase_	(							self		:	str      )	->							List[Any]:
       SCREAMING_SNAKE_CASE__					=					self.model
       SCREAMING_SNAKE_CASE__					=					['''bias''', '''LayerNorm.weight''']
       SCREAMING_SNAKE_CASE__					=					[
           {
               '''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:
            SCREAMING_SNAKE_CASE__					=					Adafactor(
                __lowerCamelCase			,     lr=self.hparams.learning_rate			,     scale_parameter=__lowerCamelCase			,     relative_step=__lowerCamelCase      )

       else:
            SCREAMING_SNAKE_CASE__					=					AdamW(
                __lowerCamelCase			,     lr=self.hparams.learning_rate			,     eps=self.hparams.adam_epsilon      )
       SCREAMING_SNAKE_CASE__					=					optimizer

       SCREAMING_SNAKE_CASE__					=					self.get_lr_scheduler()

       return [optimizer], [scheduler]






  def    lowercase_	(							self		:	Any			,     __lowerCamelCase		:	Optional[Any]			,     __lowerCamelCase		:	Tuple      )	->							Dict:
       return self.validation_step(__lowerCamelCase			,     __lowerCamelCase      )






  def    lowercase_	(							self		:	Any			,     __lowerCamelCase		:	List[Any]      )	->							Union[str, Any]:
       return self.validation_end(__lowerCamelCase      )






  def    lowercase_	(							self		:	List[Any]      )	->							int:
       SCREAMING_SNAKE_CASE__					=					max(1			,     self.hparams.gpus      )  # TODO: consider num_tpu_cores
       SCREAMING_SNAKE_CASE__					=					self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
       return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs






  def    lowercase_	(							self		:	str			,     __lowerCamelCase		:	Any      )	->							Dict:
       if stage == "test":
            SCREAMING_SNAKE_CASE__					=					len(self.test_dataloader().dataset      )
       else:
            SCREAMING_SNAKE_CASE__					=					self.get_dataloader('''train'''			,     self.hparams.train_batch_size			,     shuffle=__lowerCamelCase      )
            SCREAMING_SNAKE_CASE__					=					len(self.train_dataloader().dataset      )






  def    lowercase_	(							self		:	Optional[Any]			,     __lowerCamelCase		:	str			,     __lowerCamelCase		:	int			,     __lowerCamelCase		:	bool = False      )	->							Dict:
       raise NotImplementedError('''You must implement this for your task'''      )






  def    lowercase_	(							self		:	Any      )	->							int:
       return self.train_loader






  def    lowercase_	(							self		:	Optional[Any]      )	->							int:
       return self.get_dataloader('''dev'''			,     self.hparams.eval_batch_size			,     shuffle=__lowerCamelCase      )






  def    lowercase_	(							self		:	Dict      )	->							Optional[Any]:
       return self.get_dataloader('''test'''			,     self.hparams.eval_batch_size			,     shuffle=__lowerCamelCase      )






  def    lowercase_	(							self		:	Optional[int]			,     __lowerCamelCase		:	List[Any]      )	->							Optional[Any]:
       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    lowercase_	(							self		:	int			,     __lowerCamelCase		:	Dict[str, Any]      )	->							None:
       SCREAMING_SNAKE_CASE__					=					self.output_dir.joinpath('''best_tfmr'''      )
       SCREAMING_SNAKE_CASE__					=					self.step_count
       self.model.save_pretrained(__lowerCamelCase      )
       self.tokenizer.save_pretrained(__lowerCamelCase      )







  @staticmethod
  def    lowercase_	(							__lowerCamelCase		:	Union[str, Any]			,     __lowerCamelCase		:	List[str]      )	->							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  UpperCAmelCase__	( pl.Callback ):






  """simple docstring"""


  def    lowercase_	(							self		:	Union[str, Any]			,     __lowerCamelCase		:	Optional[int]			,     __lowerCamelCase		:	Any      )	->							List[str]:
       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  UpperCAmelCase__	( pl.Callback ):






  """simple docstring"""


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






class  UpperCAmelCase__	( pl.Callback ):






  """simple docstring"""


  def    lowercase_	(							self		:	List[str]			,     __lowerCamelCase		:	Optional[int]			,     __lowerCamelCase		:	Optional[int]      )	->							Dict:
       SCREAMING_SNAKE_CASE__					=					trainer.lr_schedulers[0]['''scheduler''']
       SCREAMING_SNAKE_CASE__					=					{f'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr()      )}
       pl_module.logger.log_metrics(__lowerCamelCase      )






  def    lowercase_	(							self		:	List[Any]			,     __lowerCamelCase		:	pl.Trainer			,     __lowerCamelCase		:	pl.LightningModule      )	->							Tuple:
       rank_zero_info('''***** Validation results *****'''      )
       SCREAMING_SNAKE_CASE__					=					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    lowercase_	(							self		:	List[Any]			,     __lowerCamelCase		:	pl.Trainer			,     __lowerCamelCase		:	pl.LightningModule      )	->							List[str]:
       rank_zero_info('''***** Test results *****'''      )
       SCREAMING_SNAKE_CASE__					=					trainer.callback_metrics
       # Log and save results to file
       SCREAMING_SNAKE_CASE__					=					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        UpperCAmelCase_     (		_A   ,	_A       ):





     '''simple docstring'''




     parser.add_argument(
         '''--output_dir'''   ,	default=str(Path(_A       ).parent / '''test_run''' / '''model_checkpoints'''       )   ,	type=_A   ,	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=_A   ,	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=_A       )
     parser.add_argument('''--max_grad_norm'''   ,	dest='''gradient_clip_val'''   ,	default=1.0   ,	type=_A   ,	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=_A   ,	default=1   ,	help='''Number of updates steps to accumulate before performing a backward/update pass.'''   ,	)
     parser.add_argument('''--seed'''   ,	type=_A   ,	default=42   ,	help='''random seed for initialization'''       )
     parser.add_argument(
         '''--data_dir'''   ,	default=str(Path(_A       ).parent / '''test_run''' / '''dummy-train-data'''       )   ,	type=_A   ,	help='''The input data dir. Should contain the training files for the CoNLL-2003 NER task.'''   ,	)







def        UpperCAmelCase_     (		_A   ,	_A   ,	_A=None   ,	_A=True   ,	_A=[]   ,	_A=None   ,	_A=None   ,	**_A   ,	):





     '''simple docstring'''




     pl.seed_everything(args.seed       )

     # init model
     SCREAMING_SNAKE_CASE__					=					Path(model.hparams.output_dir       )
     odir.mkdir(exist_ok=_A       )

     # add custom checkpoints
     if checkpoint_callback is None:
          SCREAMING_SNAKE_CASE__					=					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(_A       )
     if logging_callback is None:
          SCREAMING_SNAKE_CASE__					=					LoggingCallback()

     SCREAMING_SNAKE_CASE__					=					{}

     if args.fpaa:
          SCREAMING_SNAKE_CASE__					=					16

     if args.gpus > 1:
          SCREAMING_SNAKE_CASE__					=					'''auto'''
          SCREAMING_SNAKE_CASE__					=					'''ddp'''

     SCREAMING_SNAKE_CASE__					=					args.accumulate_grad_batches
     SCREAMING_SNAKE_CASE__					=					None
     SCREAMING_SNAKE_CASE__					=					'''auto'''

     SCREAMING_SNAKE_CASE__					=					pl.Trainer.from_argparse_args(
         _A   ,	weights_summary=_A   ,	callbacks=[logging_callback] + extra_callbacks + [InitCallback()] + [checkpoint_callback]   ,	logger=_A   ,	val_check_interval=1   ,	num_sanity_val_steps=2   ,	**_A   ,	)

     if args.do_train:
          trainer.fit(_A       )

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





import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, XLMRobertaTokenizer

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

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


enable_full_determinism()






class  UpperCAmelCase__	( A__		,  A__		,  A__		,  unittest.TestCase ):






  """simple docstring"""


  a     =      AltDiffusionPipeline
  a     =      TEXT_TO_IMAGE_PARAMS
  a     =      TEXT_TO_IMAGE_BATCH_PARAMS
  a     =      TEXT_TO_IMAGE_IMAGE_PARAMS
  a     =      TEXT_TO_IMAGE_IMAGE_PARAMS






  def    lowercase_	(							self		:	Tuple      )	->							Tuple:
       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					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''')			,     cross_attention_dim=32			,     )
       SCREAMING_SNAKE_CASE__					=					DDIMScheduler(
           beta_start=0.00085			,     beta_end=0.012			,     beta_schedule='''scaled_linear'''			,     clip_sample=__lowerCamelCase			,     set_alpha_to_one=__lowerCamelCase			,     )
       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					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			,     )

       # TODO: address the non-deterministic text encoder (fails for save-load tests)
       # torch.manual_seed(0)
       # text_encoder_config = RobertaSeriesConfig(
       #     hidden_size=32,
       #     project_dim=32,
       #     intermediate_size=37,
       #     layer_norm_eps=1e-05,
       #     num_attention_heads=4,
       #     num_hidden_layers=5,
       #     vocab_size=5002,
       # )
       # text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					CLIPTextConfig(
           bos_token_id=0			,     eos_token_id=2			,     hidden_size=32			,     projection_dim=32			,     intermediate_size=37			,     layer_norm_eps=1e-05			,     num_attention_heads=4			,     num_hidden_layers=5			,     pad_token_id=1			,     vocab_size=5002			,     )
       SCREAMING_SNAKE_CASE__					=					CLIPTextModel(__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''      )
       SCREAMING_SNAKE_CASE__					=					77

       SCREAMING_SNAKE_CASE__					=					{
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''vae''': vae,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''safety_checker''': None,
           '''feature_extractor''': None,
       }
       return components






  def    lowercase_	(							self		:	Any			,     __lowerCamelCase		:	Tuple			,     __lowerCamelCase		:	Optional[int]=0      )	->							Union[str, Any]:
       if str(__lowerCamelCase      ).startswith('''mps'''      ):
            SCREAMING_SNAKE_CASE__					=					torch.manual_seed(__lowerCamelCase      )
       else:
            SCREAMING_SNAKE_CASE__					=					torch.Generator(device=__lowerCamelCase      ).manual_seed(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					{
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''guidance_scale''': 6.0,
           '''output_type''': '''numpy''',
       }
       return inputs






  def    lowercase_	(							self		:	List[Any]      )	->							str:
       super().test_attention_slicing_forward_pass(expected_max_diff=3e-3      )






  def    lowercase_	(							self		:	List[Any]      )	->							int:
       super().test_inference_batch_single_identical(expected_max_diff=3e-3      )






  def    lowercase_	(							self		:	Dict      )	->							Optional[int]:
       SCREAMING_SNAKE_CASE__					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator

       SCREAMING_SNAKE_CASE__					=					self.get_dummy_components()
       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					RobertaSeriesConfig(
           hidden_size=32			,     project_dim=32			,     intermediate_size=37			,     layer_norm_eps=1e-05			,     num_attention_heads=4			,     num_hidden_layers=5			,     vocab_size=5002			,     )
       # TODO: remove after fixing the non-deterministic text encoder
       SCREAMING_SNAKE_CASE__					=					RobertaSeriesModelWithTransformation(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					text_encoder

       SCREAMING_SNAKE_CASE__					=					AltDiffusionPipeline(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe.to(__lowerCamelCase      )
       alt_pipe.set_progress_bar_config(disable=__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					self.get_dummy_inputs(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					'''A photo of an astronaut'''
       SCREAMING_SNAKE_CASE__					=					alt_pipe(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					output.images
       SCREAMING_SNAKE_CASE__					=					image[0, -3:, -3:, -1]

       assert image.shape == (1, 64, 64, 3)
       SCREAMING_SNAKE_CASE__					=					np.array(
           [0.5748162, 0.60447145, 0.48821217, 0.50100636, 0.5431185, 0.45763683, 0.49657696, 0.48132733, 0.47573093]      )

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







  def    lowercase_	(							self		:	List[Any]      )	->							Optional[Any]:
       SCREAMING_SNAKE_CASE__					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator

       SCREAMING_SNAKE_CASE__					=					self.get_dummy_components()
       SCREAMING_SNAKE_CASE__					=					PNDMScheduler(skip_prk_steps=__lowerCamelCase      )
       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					RobertaSeriesConfig(
           hidden_size=32			,     project_dim=32			,     intermediate_size=37			,     layer_norm_eps=1e-05			,     num_attention_heads=4			,     num_hidden_layers=5			,     vocab_size=5002			,     )
       # TODO: remove after fixing the non-deterministic text encoder
       SCREAMING_SNAKE_CASE__					=					RobertaSeriesModelWithTransformation(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					text_encoder
       SCREAMING_SNAKE_CASE__					=					AltDiffusionPipeline(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe.to(__lowerCamelCase      )
       alt_pipe.set_progress_bar_config(disable=__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					self.get_dummy_inputs(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					output.images
       SCREAMING_SNAKE_CASE__					=					image[0, -3:, -3:, -1]

       assert image.shape == (1, 64, 64, 3)
       SCREAMING_SNAKE_CASE__					=					np.array(
           [0.51605093, 0.5707241, 0.47365507, 0.50578886, 0.5633877, 0.4642503, 0.5182081, 0.48763484, 0.49084237]      )

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






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






  """simple docstring"""


  def    lowercase_	(							self		:	Tuple      )	->							str:
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






  def    lowercase_	(							self		:	str      )	->							Any:
       # make sure here that pndm scheduler skips prk
       SCREAMING_SNAKE_CASE__					=					AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''			,     safety_checker=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe.to(__lowerCamelCase      )
       alt_pipe.set_progress_bar_config(disable=__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					'''A painting of a squirrel eating a burger'''
       SCREAMING_SNAKE_CASE__					=					torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe([prompt]			,     generator=__lowerCamelCase			,     guidance_scale=6.0			,     num_inference_steps=20			,     output_type='''np'''      )

       SCREAMING_SNAKE_CASE__					=					output.images

       SCREAMING_SNAKE_CASE__					=					image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       SCREAMING_SNAKE_CASE__					=					np.array([0.1010, 0.0800, 0.0794, 0.0885, 0.0843, 0.0762, 0.0769, 0.0729, 0.0586]      )

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







  def    lowercase_	(							self		:	Optional[int]      )	->							List[str]:
       SCREAMING_SNAKE_CASE__					=					DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''			,     subfolder='''scheduler'''      )

       SCREAMING_SNAKE_CASE__					=					AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''			,     scheduler=__lowerCamelCase			,     safety_checker=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe.to(__lowerCamelCase      )
       alt_pipe.set_progress_bar_config(disable=__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					'''A painting of a squirrel eating a burger'''
       SCREAMING_SNAKE_CASE__					=					torch.manual_seed(0      )

       SCREAMING_SNAKE_CASE__					=					alt_pipe([prompt]			,     generator=__lowerCamelCase			,     num_inference_steps=2			,     output_type='''numpy'''      )
       SCREAMING_SNAKE_CASE__					=					output.images

       SCREAMING_SNAKE_CASE__					=					image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       SCREAMING_SNAKE_CASE__					=					np.array([0.4019, 0.4052, 0.3810, 0.4119, 0.3916, 0.3982, 0.4651, 0.4195, 0.5323]      )

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


