
"""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


_a									=				logging.getLogger(__name__)

require_version('pytorch_lightning>=1.0.4')

_a									=				{
    '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
_a									=				{
    '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
}
_a									=				sorted(arg_to_scheduler.keys())
_a									=				'{' + ', '.join(arg_to_scheduler_choices) + '}'


class A_      (pl.LightningModule ):




  '''simple docstring'''







  def __init__(						self							,  lowercase_							,  lowercase_=None							,  lowercase_="base"							,  lowercase_=None							,  lowercase_=None							,  lowercase_=None							,  **lowercase_							,  ):

     """simple docstring"""
     super().__init__()
     # TODO: move to self.save_hyperparameters()
     # self.save_hyperparameters()
     # can also expand arguments into trainer signature for easier reading

     self.save_hyperparameters(lowercase_    )
     UpperCAmelCase_   :   Tuple     	=       0
     UpperCAmelCase_   :   Optional[int]     	=       Path(self.hparams.output_dir    )
     UpperCAmelCase_   :   Optional[Any]     	=       self.hparams.cache_dir if self.hparams.cache_dir else None
     if config is None:
        UpperCAmelCase_   :   Any     	=       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=lowercase_							,  **lowercase_							,  )
     else:
        UpperCAmelCase_   :   PretrainedConfig     	=       config

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

     if tokenizer is None:
        UpperCAmelCase_   :   Optional[Any]     	=       AutoTokenizer.from_pretrained(
            self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path							,  cache_dir=lowercase_							,  )
     else:
        UpperCAmelCase_   :   PreTrainedTokenizer     	=       tokenizer
     UpperCAmelCase_   :   str     	=       MODEL_MODES[mode]
     if model is None:
        UpperCAmelCase_   :   List[str]     	=       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=lowercase_							,  )
     else:
        UpperCAmelCase_   :   int     	=       model







  def 			UpperCamelCase__   (						self							,  *lowercase_							,  **lowercase_    ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[int]     	=       self.model_type.from_pretrained(*lowercase_							,  **lowercase_    )







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   List[str]     	=       arg_to_scheduler[self.hparams.lr_scheduler]
     UpperCAmelCase_   :   Optional[Any]     	=       get_schedule_func(
         self.opt							,  num_warmup_steps=self.hparams.warmup_steps							,  num_training_steps=self.total_steps()    )
     UpperCAmelCase_   :   Optional[Any]     	=       {"scheduler": scheduler, "interval": "step", "frequency": 1}
     return scheduler







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Dict     	=       self.model
     UpperCAmelCase_   :   Any     	=       ["bias", "LayerNorm.weight"]
     UpperCAmelCase_   :   Dict     	=       [
         {
             "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay    )],
             "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:
        UpperCAmelCase_   :   Any     	=       Adafactor(
            lowercase_							,  lr=self.hparams.learning_rate							,  scale_parameter=lowercase_							,  relative_step=lowercase_    )

     else:
        UpperCAmelCase_   :   str     	=       AdamW(
            lowercase_							,  lr=self.hparams.learning_rate							,  eps=self.hparams.adam_epsilon    )
     UpperCAmelCase_   :   Union[str, Any]     	=       optimizer

     UpperCAmelCase_   :   Union[str, Any]     	=       self.get_lr_scheduler()

     return [optimizer], [scheduler]







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_    ):

     """simple docstring"""
     return self.validation_step(lowercase_							,  lowercase_    )







  def 			UpperCamelCase__   (						self							,  lowercase_    ):

     """simple docstring"""
     return self.validation_end(lowercase_    )







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[Any]     	=       max(1							,  self.hparams.gpus    )  # TODO: consider num_tpu_cores
     UpperCAmelCase_   :   Union[str, Any]     	=       self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
     return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs







  def 			UpperCamelCase__   (						self							,  lowercase_    ):

     """simple docstring"""
     if mode == "test":
        UpperCAmelCase_   :   Dict     	=       len(self.test_dataloader().dataset    )
     else:
        UpperCAmelCase_   :   Union[str, Any]     	=       self.get_dataloader("train"							,  self.hparams.train_batch_size							,  shuffle=lowercase_    )
        UpperCAmelCase_   :   Union[str, Any]     	=       len(self.train_dataloader().dataset    )







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_							,  lowercase_ = False    ):

     """simple docstring"""
     raise NotImplementedError("You must implement this for your task"    )







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     return self.train_loader







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     return self.get_dataloader("dev"							,  self.hparams.eval_batch_size							,  shuffle=lowercase_    )







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     return self.get_dataloader("test"							,  self.hparams.eval_batch_size							,  shuffle=lowercase_    )







  def 			UpperCamelCase__   (						self							,  lowercase_    ):

     """simple docstring"""
     return os.path.join(
         self.hparams.data_dir							,  "cached_{}_{}_{}".format(
             lowercase_							,  list(filter(lowercase_							,  self.hparams.model_name_or_path.split("/"    )    )    ).pop()							,  str(self.hparams.max_seq_length    )							,  )							,  )







  @pl.utilities.rank_zero_only
  def 			UpperCamelCase__   (						self							,  lowercase_    ):

     """simple docstring"""
     UpperCAmelCase_   :   str     	=       self.output_dir.joinpath("best_tfmr"    )
     UpperCAmelCase_   :   Optional[int]     	=       self.step_count
     self.model.save_pretrained(lowercase_    )
     self.tokenizer.save_pretrained(lowercase_    )







  @staticmethod
  def 			UpperCamelCase__   (						lowercase_							,  lowercase_    ):

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




class A_      (pl.Callback ):




  '''simple docstring'''







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_    ):

     """simple docstring"""
     UpperCAmelCase_   :   List[str]     	=       trainer.lr_schedulers[0]["scheduler"]
     UpperCAmelCase_   :   List[Any]     	=       {F"""lr_group_{i}""": lr for i, lr in enumerate(lr_scheduler.get_lr()    )}
     pl_module.logger.log_metrics(lowercase_    )







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_    ):

     """simple docstring"""
     rank_zero_info("***** Validation results *****"    )
     UpperCAmelCase_   :   Dict     	=       trainer.callback_metrics
     # Log results
     for key in sorted(lowercase_    ):
        if key not in ["log", "progress_bar"]:
           rank_zero_info("{} = {}\n".format(lowercase_							,  str(metrics[key]    )    )    )







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_    ):

     """simple docstring"""
     rank_zero_info("***** Test results *****"    )
     UpperCAmelCase_   :   Any     	=       trainer.callback_metrics
     # Log and save results to file
     UpperCAmelCase_   :   Optional[int]     	=       os.path.join(pl_module.hparams.output_dir							,  "test_results.txt"    )
     with open(lowercase_							,  "w"    ) as writer:
        for key in sorted(lowercase_    ):
           if key not in ["log", "progress_bar"]:
              rank_zero_info("{} = {}\n".format(lowercase_							,  str(metrics[key]    )    )    )
              writer.write("{} = {}\n".format(lowercase_							,  str(metrics[key]    )    )    )







def       __a	(						__lowerCamelCase,						__lowerCamelCase    ):
   #  To allow all pl args uncomment the following line
   #  parser = pl.Trainer.add_argparse_args(parser)
   parser.add_argument(
       "--output_dir",						default=__lowerCamelCase,						type=__lowerCamelCase,						required=__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=42,						help="random seed for initialization"    )
   parser.add_argument(
       "--data_dir",						default=__lowerCamelCase,						type=__lowerCamelCase,						required=__lowerCamelCase,						help="The input data dir. Should contain the training files for the CoNLL-2003 NER task.",						)







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

   # init model
   UpperCAmelCase_   :   Any     	=       Path(model.hparams.output_dir    )
   odir.mkdir(exist_ok=__lowerCamelCase    )

   # add custom checkpoints
   if checkpoint_callback is None:
      UpperCAmelCase_   :   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:
      UpperCAmelCase_   :   Union[str, Any]     	=       LoggingCallback()

   UpperCAmelCase_   :   int     	=       {}

   # TODO: remove with PyTorch 1.6 since pl uses native amp
   if args.fpaa:
      UpperCAmelCase_   :   List[Any]     	=       16
      UpperCAmelCase_   :   str     	=       args.fpaa_opt_level

   if args.gpus > 1:
      UpperCAmelCase_   :   List[str]     	=       "ddp"

   UpperCAmelCase_   :   List[str]     	=       args.accumulate_grad_batches
   UpperCAmelCase_   :   Tuple     	=       extra_train_kwargs.get("accelerator",						__lowerCamelCase    )
   UpperCAmelCase_   :   List[str]     	=       extra_train_kwargs.get("profiler",						__lowerCamelCase    )

   UpperCAmelCase_   :   int     	=       pl.Trainer.from_argparse_args(
       __lowerCamelCase,						weights_summary=__lowerCamelCase,						callbacks=[logging_callback] + extra_callbacks,						logger=__lowerCamelCase,						checkpoint_callback=__lowerCamelCase,						**__lowerCamelCase,						)

   if args.do_train:
      trainer.fit(__lowerCamelCase    )

   return trainer




'''simple docstring'''







import gc
import random
import tempfile
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







class 		snake_case_       (						unittest.TestCase						):



   """simple docstring"""





   def     __UpperCAmelCase      ( self):
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()







   @property
   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      1
          lowerCamelCase__				   =      3
          lowerCamelCase__				   =      (32, 32)

          lowerCamelCase__				   =      floats_tensor((batch_size, num_channels) + sizes			,   rng=random.Random(0)).to(UpperCamelCase)
          return image







   @property
   def     __UpperCAmelCase      ( self):
          torch.manual_seed(0)
          lowerCamelCase__				   =      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			,   )
          return model







   @property
   def     __UpperCAmelCase      ( self):
          torch.manual_seed(0)
          lowerCamelCase__				   =      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			,   )
          return model







   @property
   def     __UpperCAmelCase      ( self):
          torch.manual_seed(0)
          lowerCamelCase__				   =      CLIPTextConfig(
              bos_token_id=0			,   eos_token_id=2			,   hidden_size=32			,   intermediate_size=37			,   layer_norm_eps=1E-0_5			,   num_attention_heads=4			,   num_hidden_layers=5			,   pad_token_id=1			,   vocab_size=10_00			,   )
          return CLIPTextModel(UpperCamelCase)







   @property
   def     __UpperCAmelCase      ( self):
          def extract(*UpperCamelCase			,   **UpperCamelCase):







                 class 		snake_case_       :



                    """simple docstring"""





                    def __init__( self):
                           lowerCamelCase__				   =      torch.ones([0])







                    def     __UpperCAmelCase      ( self			,   UpperCamelCase):
                           self.pixel_values.to(UpperCamelCase)
                           return self






                 return Out()

          return extract







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cpu"  # ensure determinism for the device-dependent torch.Generator
          lowerCamelCase__				   =      self.dummy_cond_unet
          lowerCamelCase__				   =      DDIMScheduler(
              beta_start=0.0_0_0_8_5			,   beta_end=0.0_1_2			,   beta_schedule="scaled_linear"			,   clip_sample=UpperCamelCase			,   set_alpha_to_one=UpperCamelCase			,   )

          lowerCamelCase__				   =      self.dummy_vae
          lowerCamelCase__				   =      self.dummy_text_encoder
          lowerCamelCase__				   =      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

          # make sure here that pndm scheduler skips prk
          lowerCamelCase__				   =      StableDiffusionPipeline(
              unet=UpperCamelCase			,   scheduler=UpperCamelCase			,   vae=UpperCamelCase			,   text_encoder=UpperCamelCase			,   tokenizer=UpperCamelCase			,   safety_checker=UpperCamelCase			,   feature_extractor=self.dummy_extractor			,   )
          lowerCamelCase__				   =      sd_pipe.to(UpperCamelCase)
          sd_pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "A painting of a squirrel eating a burger"

          lowerCamelCase__				   =      torch.Generator(device=UpperCamelCase).manual_seed(0)
          lowerCamelCase__				   =      sd_pipe([prompt]			,   generator=UpperCamelCase			,   guidance_scale=6.0			,   num_inference_steps=2			,   output_type="np")
          lowerCamelCase__				   =      output.images

          lowerCamelCase__				   =      torch.Generator(device=UpperCamelCase).manual_seed(0)
          lowerCamelCase__				   =      sd_pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=6.0			,   num_inference_steps=2			,   output_type="np"			,   return_dict=UpperCamelCase			,   )[0]

          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      image_from_tuple[0, -3:, -3:, -1]

          assert image.shape == (1, 64, 64, 3)
          lowerCamelCase__				   =      np.array([0.5_7_5_3, 0.6_1_1_4, 0.5_0_0_1, 0.5_0_3_4, 0.5_4_7_0, 0.4_7_2_9, 0.4_9_7_1, 0.4_8_6_7, 0.4_8_6_7])

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







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cpu"  # ensure determinism for the device-dependent torch.Generator
          lowerCamelCase__				   =      self.dummy_cond_unet
          lowerCamelCase__				   =      PNDMScheduler(skip_prk_steps=UpperCamelCase)
          lowerCamelCase__				   =      self.dummy_vae
          lowerCamelCase__				   =      self.dummy_text_encoder
          lowerCamelCase__				   =      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

          # make sure here that pndm scheduler skips prk
          lowerCamelCase__				   =      StableDiffusionPipeline(
              unet=UpperCamelCase			,   scheduler=UpperCamelCase			,   vae=UpperCamelCase			,   text_encoder=UpperCamelCase			,   tokenizer=UpperCamelCase			,   safety_checker=UpperCamelCase			,   feature_extractor=self.dummy_extractor			,   )
          lowerCamelCase__				   =      sd_pipe.to(UpperCamelCase)
          sd_pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "A painting of a squirrel eating a burger"
          lowerCamelCase__				   =      torch.Generator(device=UpperCamelCase).manual_seed(0)
          lowerCamelCase__				   =      sd_pipe([prompt]			,   generator=UpperCamelCase			,   guidance_scale=6.0			,   num_inference_steps=2			,   output_type="np")

          lowerCamelCase__				   =      output.images

          lowerCamelCase__				   =      torch.Generator(device=UpperCamelCase).manual_seed(0)
          lowerCamelCase__				   =      sd_pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=6.0			,   num_inference_steps=2			,   output_type="np"			,   return_dict=UpperCamelCase			,   )[0]

          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      image_from_tuple[0, -3:, -3:, -1]

          assert image.shape == (1, 64, 64, 3)
          lowerCamelCase__				   =      np.array([0.5_1_2_2, 0.5_7_1_2, 0.4_8_2_5, 0.5_0_5_3, 0.5_6_4_6, 0.4_7_6_9, 0.5_1_7_9, 0.4_8_9_4, 0.4_9_9_4])

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







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained(
              "hf-internal-testing/tiny-stable-diffusion-lms-pipe"			,   safety_checker=UpperCamelCase)
          assert isinstance(UpperCamelCase			,   UpperCamelCase)
          assert isinstance(pipe.scheduler			,   UpperCamelCase)
          assert pipe.safety_checker is None

          lowerCamelCase__				   =      pipe("example prompt"			,   num_inference_steps=2).images[0]
          assert image is not None

          # check that there's no error when saving a pipeline with one of the models being None
          with tempfile.TemporaryDirectory() as tmpdirname:
                 pipe.save_pretrained(UpperCamelCase)
                 lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained(UpperCamelCase)

          # sanity check that the pipeline still works
          assert pipe.safety_checker is None
          lowerCamelCase__				   =      pipe("example prompt"			,   num_inference_steps=2).images[0]
          assert image is not None







   @unittest.skipIf(torch_device != "cuda"			,   "This test requires a GPU")
   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.dummy_cond_unet
          lowerCamelCase__				   =      PNDMScheduler(skip_prk_steps=UpperCamelCase)
          lowerCamelCase__				   =      self.dummy_vae
          lowerCamelCase__				   =      self.dummy_text_encoder
          lowerCamelCase__				   =      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

          # put models in fp16
          lowerCamelCase__				   =      unet.half()
          lowerCamelCase__				   =      vae.half()
          lowerCamelCase__				   =      bert.half()

          # make sure here that pndm scheduler skips prk
          lowerCamelCase__				   =      StableDiffusionPipeline(
              unet=UpperCamelCase			,   scheduler=UpperCamelCase			,   vae=UpperCamelCase			,   text_encoder=UpperCamelCase			,   tokenizer=UpperCamelCase			,   safety_checker=UpperCamelCase			,   feature_extractor=self.dummy_extractor			,   )
          lowerCamelCase__				   =      sd_pipe.to(UpperCamelCase)
          sd_pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "A painting of a squirrel eating a burger"
          lowerCamelCase__				   =      sd_pipe([prompt]			,   num_inference_steps=2			,   output_type="np").images

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







@nightly
@require_torch_gpu
class 		snake_case_       (						unittest.TestCase						):



   """simple docstring"""





   def     __UpperCAmelCase      ( self):
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cuda"
          lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
          lowerCamelCase__				   =      pipe.to(UpperCamelCase)
          pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "a photo of a cat"
          lowerCamelCase__				   =      {
              "editing_prompt": ["sunglasses"],
              "reverse_editing_direction": [False],
              "edit_warmup_steps": 10,
              "edit_guidance_scale": 6,
              "edit_threshold": 0.9_5,
              "edit_momentum_scale": 0.5,
              "edit_mom_beta": 0.6,
          }

          lowerCamelCase__				   =      3
          lowerCamelCase__				   =      7

          # no sega enabled
          lowerCamelCase__				   =      torch.Generator(UpperCamelCase)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              0.3_4_6_7_3_1_1_3,
              0.3_8_4_9_2_7_3_3,
              0.3_7_5_9_7_3_5_2,
              0.3_4_0_8_6_3_3_5,
              0.3_5_6_5_0_7_4_8,
              0.3_5_5_7_9_2_0_5,
              0.3_3_8_4_7_6_3,
              0.3_4_3_4_0_2_3_6,
              0.3_5_7_3_2_7_1,
          ]

          assert image.shape == (1, 5_12, 5_12, 3)

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

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   **UpperCamelCase			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              0.4_1_8_8_7_8_2_6,
              0.3_7_7_2_8_7_6_6,
              0.3_0_1_3_8_2_7_2,
              0.4_1_4_1_6_3_3_5,
              0.4_1_6_6_4_9_8_5,
              0.3_6_2_8_3_3_9_2,
              0.3_6_1_9_1_2_4_6,
              0.4_3_3_6_4_4_6_5,
              0.4_3_0_0_1_7_3_2,
          ]

          assert image.shape == (1, 5_12, 5_12, 3)

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







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cuda"
          lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
          lowerCamelCase__				   =      pipe.to(UpperCamelCase)
          pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "an image of a crowded boulevard, realistic, 4k"
          lowerCamelCase__				   =      {
              "editing_prompt": "crowd, crowded, people",
              "reverse_editing_direction": True,
              "edit_warmup_steps": 10,
              "edit_guidance_scale": 8.3,
              "edit_threshold": 0.9,
              "edit_momentum_scale": 0.5,
              "edit_mom_beta": 0.6,
          }

          lowerCamelCase__				   =      9
          lowerCamelCase__				   =      7

          # no sega enabled
          lowerCamelCase__				   =      torch.Generator(UpperCamelCase)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              0.4_3_4_9_7_9_9_8,
              0.9_1_8_1_4_0_6_5,
              0.7_5_4_0_7_3_9,
              0.5_5_5_8_0_2_0_5,
              0.8_4_6_7_2_6_5,
              0.5_3_8_9_6_9_1,
              0.6_2_5_7_4_5_0_6,
              0.5_8_8_9_7_7_6_3,
              0.5_0_9_2_6_7_5_7,
          ]

          assert image.shape == (1, 5_12, 5_12, 3)

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

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   **UpperCamelCase			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              0.3_0_8_9_7_1_9,
              0.3_0_5_0_0_1_4_4,
              0.2_9_0_1_6_0_4_2,
              0.3_0_6_3_0_9_6_4,
              0.3_2_5_6_8_7,
              0.2_9_4_1_9_2_2_5,
              0.2_9_0_8_0_9_1,
              0.2_8_7_2_3_5_9_8,
              0.2_7_6_9_6_2_9_4,
          ]

          assert image.shape == (1, 5_12, 5_12, 3)

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







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cuda"
          lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
          lowerCamelCase__				   =      pipe.to(UpperCamelCase)
          pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "a castle next to a river"
          lowerCamelCase__				   =      {
              "editing_prompt": ["boat on a river, boat", "monet, impression, sunrise"],
              "reverse_editing_direction": False,
              "edit_warmup_steps": [15, 18],
              "edit_guidance_scale": 6,
              "edit_threshold": [0.9, 0.8],
              "edit_momentum_scale": 0.5,
              "edit_mom_beta": 0.6,
          }

          lowerCamelCase__				   =      48
          lowerCamelCase__				   =      7

          # no sega enabled
          lowerCamelCase__				   =      torch.Generator(UpperCamelCase)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              0.7_5_1_6_3_5_5_5,
              0.7_6_0_3_7_1_4_5,
              0.6_1_7_8_5,
              0.9_1_8_9_6_7_3,
              0.8_6_2_7_7_0_1,
              0.8_5_1_8_9_6_9_4,
              0.8_5_1_2_8_1_3,
              0.8_7_0_1_2_0_7_6,
              0.8_3_1_2_8_5_7,
          ]

          assert image.shape == (1, 5_12, 5_12, 3)

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

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   **UpperCamelCase			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              0.7_3_5_5_3_3_6_5,
              0.7_5_3_7_2_7_1,
              0.7_4_3_4_1_9_0_5,
              0.6_6_4_8_0_3_5_6,
              0.6_4_7_2_9_2_5,
              0.6_3_0_3_9_4_1_6,
              0.6_4_8_1_2_9_0_5,
              0.6_7_4_9_7_1_7,
              0.6_5_1_7_1_0_2,
          ]

          assert image.shape == (1, 5_12, 5_12, 3)

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







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cuda"
          lowerCamelCase__				   =      StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"			,   torch_dtype=torch.floataa)
          lowerCamelCase__				   =      pipe.to(UpperCamelCase)
          pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "a photo of a cat"
          lowerCamelCase__				   =      {
              "editing_prompt": ["sunglasses"],
              "reverse_editing_direction": [False],
              "edit_warmup_steps": 10,
              "edit_guidance_scale": 6,
              "edit_threshold": 0.9_5,
              "edit_momentum_scale": 0.5,
              "edit_mom_beta": 0.6,
          }

          lowerCamelCase__				   =      3
          lowerCamelCase__				   =      7

          # no sega enabled
          lowerCamelCase__				   =      torch.Generator(UpperCamelCase)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              0.3_4_8_8_7_6_9_5,
              0.3_8_7_6_9_5_3,
              0.3_7_5,
              0.3_4_4_2_3_8_2_8,
              0.3_5_8_1_5_4_3,
              0.3_5_7_1_7_7_7_3,
              0.3_3_8_3_7_8_9,
              0.3_4_5_7_0_3_1_2,
              0.3_5_9_3_7_5,
          ]

          assert image.shape == (1, 5_12, 5_12, 3)

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

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   guidance_scale=UpperCamelCase			,   num_inference_steps=50			,   output_type="np"			,   width=5_12			,   height=5_12			,   **UpperCamelCase			,   )

          lowerCamelCase__				   =      output.images
          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      [
              0.4_2_2_8_5_1_5_6,
              0.3_6_9_1_4_0_6_2,
              0.2_9_0_7_7_1_4_8,
              0.4_2_0_4_1_0_1_6,
              0.4_1_9_1_8_9_4_5,
              0.3_5_4_9_8_0_4_7,
              0.3_6_1_8_1_6_4,
              0.4_4_2_3_8_2_8,
              0.4_3_1_1_5_2_3_4,
          ]

          assert image.shape == (1, 5_12, 5_12, 3)

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

