


import argparse
import os

import gluonnlp as nlp
import mxnet as mx
import numpy as np
import torch
from gluonnlp.base import get_home_dir
from gluonnlp.model.bert import BERTEncoder
from gluonnlp.model.utils import _load_vocab
from gluonnlp.vocab import Vocab
from packaging import version
from torch import nn

from transformers import BertConfig, BertForMaskedLM, BertModel, RobertaTokenizer
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.utils import logging


if version.parse(nlp.__version__) != version.parse("""0.8.3"""):
     raise Exception("""requires gluonnlp == 0.8.3""")

if version.parse(mx.__version__) != version.parse("""1.5.0"""):
     raise Exception("""requires mxnet == 1.5.0""")

logging.set_verbosity_info()
__lowercase  :   Union[str, Any]						      =logging.get_logger(__name__)

__lowercase  :   Union[str, Any]						      ="""The Nymphenburg Palace is a beautiful palace in Munich!"""




def     a__ (					lowercase__  ,    lowercase__  ):






   '''simple docstring'''






   UpperCAmelCase_											={
       "attention_cell": "multi_head",
       "num_layers": 4,
       "units": 1_0_2_4,
       "hidden_size": 7_6_8,
       "max_length": 5_1_2,
       "num_heads": 8,
       "scaled": True,
       "dropout": 0.1,
       "use_residual": True,
       "embed_size": 1_0_2_4,
       "embed_dropout": 0.1,
       "word_embed": None,
       "layer_norm_eps": 1E-5,
       "token_type_vocab_size": 2,
   }

   UpperCAmelCase_											=bort_4_8_768_1024_hparams

   # Let's construct the original Bort model here
   # Taken from official BERT implementation, see:
   # https://github.com/alexa/bort/blob/master/bort/bort.py
   UpperCAmelCase_											=BERTEncoder(
       attention_cell=predefined_args["attention_cell"]  ,    num_layers=predefined_args["num_layers"]  ,    units=predefined_args["units"]  ,    hidden_size=predefined_args["hidden_size"]  ,    max_length=predefined_args["max_length"]  ,    num_heads=predefined_args["num_heads"]  ,    scaled=predefined_args["scaled"]  ,    dropout=predefined_args["dropout"]  ,    output_attention=lowercase__  ,    output_all_encodings=lowercase__  ,    use_residual=predefined_args["use_residual"]  ,    activation=predefined_args.get("activation"  ,    "gelu"  )  ,    layer_norm_eps=predefined_args.get("layer_norm_eps"  ,    lowercase__  )  ,    )

   # Vocab information needs to be fetched first
   # It's the same as RoBERTa, so RobertaTokenizer can be used later
   UpperCAmelCase_											="openwebtext_ccnews_stories_books_cased"

   # Specify download folder to Gluonnlp's vocab
   UpperCAmelCase_											=os.path.join(get_home_dir()  ,    "models"  )
   UpperCAmelCase_											=_load_vocab(lowercase__  ,    lowercase__  ,    lowercase__  ,    cls=lowercase__  )

   UpperCAmelCase_											=nlp.model.BERTModel(
       lowercase__  ,    len(lowercase__  )  ,    units=predefined_args["units"]  ,    embed_size=predefined_args["embed_size"]  ,    embed_dropout=predefined_args["embed_dropout"]  ,    word_embed=predefined_args["word_embed"]  ,    use_pooler=lowercase__  ,    use_token_type_embed=lowercase__  ,    token_type_vocab_size=predefined_args["token_type_vocab_size"]  ,    use_classifier=lowercase__  ,    use_decoder=lowercase__  ,    )

   original_bort.load_parameters(lowercase__  ,    cast_dtype=lowercase__  ,    ignore_extra=lowercase__  )
   UpperCAmelCase_											=original_bort._collect_params_with_prefix()

   # Build our config 🤗
   UpperCAmelCase_											={
       "architectures": ["BertForMaskedLM"],
       "attention_probs_dropout_prob": predefined_args["dropout"],
       "hidden_act": "gelu",
       "hidden_dropout_prob": predefined_args["dropout"],
       "hidden_size": predefined_args["embed_size"],
       "initializer_range": 0.02,
       "intermediate_size": predefined_args["hidden_size"],
       "layer_norm_eps": predefined_args["layer_norm_eps"],
       "max_position_embeddings": predefined_args["max_length"],
       "model_type": "bort",
       "num_attention_heads": predefined_args["num_heads"],
       "num_hidden_layers": predefined_args["num_layers"],
       "pad_token_id": 1,  # 2 = BERT, 1 = RoBERTa
       "type_vocab_size": 1,  # 2 = BERT, 1 = RoBERTa
       "vocab_size": len(lowercase__  ),
   }

   UpperCAmelCase_											=BertConfig.from_dict(lowercase__  )
   UpperCAmelCase_											=BertForMaskedLM(lowercase__  )
   hf_bort_model.eval()

   # Parameter mapping table (Gluonnlp to Transformers)
   # * denotes layer index
   #
   # | Gluon Parameter                                                | Transformers Parameter
   # | -------------------------------------------------------------- | ----------------------
   # | `encoder.layer_norm.beta`                                      | `bert.embeddings.LayerNorm.bias`
   # | `encoder.layer_norm.gamma`                                     | `bert.embeddings.LayerNorm.weight`
   # | `encoder.position_weight`                                      | `bert.embeddings.position_embeddings.weight`
   # | `word_embed.0.weight`                                          | `bert.embeddings.word_embeddings.weight`
   # | `encoder.transformer_cells.*.attention_cell.proj_key.bias`     | `bert.encoder.layer.*.attention.self.key.bias`
   # | `encoder.transformer_cells.*.attention_cell.proj_key.weight`   | `bert.encoder.layer.*.attention.self.key.weight`
   # | `encoder.transformer_cells.*.attention_cell.proj_query.bias`   | `bert.encoder.layer.*.attention.self.query.bias`
   # | `encoder.transformer_cells.*.attention_cell.proj_query.weight` | `bert.encoder.layer.*.attention.self.query.weight`
   # | `encoder.transformer_cells.*.attention_cell.proj_value.bias`   | `bert.encoder.layer.*.attention.self.value.bias`
   # | `encoder.transformer_cells.*.attention_cell.proj_value.weight` | `bert.encoder.layer.*.attention.self.value.weight`
   # | `encoder.transformer_cells.*.ffn.ffn_2.bias`                   | `bert.encoder.layer.*.attention.output.dense.bias`
   # | `encoder.transformer_cells.*.ffn.ffn_2.weight`                 | `bert.encoder.layer.*.attention.output.dense.weight`
   # | `encoder.transformer_cells.*.layer_norm.beta`                  | `bert.encoder.layer.*.attention.output.LayerNorm.bias`
   # | `encoder.transformer_cells.*.layer_norm.gamma`                 | `bert.encoder.layer.*.attention.output.LayerNorm.weight`
   # | `encoder.transformer_cells.*.ffn.ffn_1.bias`                   | `bert.encoder.layer.*.intermediate.dense.bias`
   # | `encoder.transformer_cells.*.ffn.ffn_1.weight`                 | `bert.encoder.layer.*.intermediate.dense.weight`
   # | `encoder.transformer_cells.*.ffn.layer_norm.beta`              | `bert.encoder.layer.*.output.LayerNorm.bias`
   # | `encoder.transformer_cells.*.ffn.layer_norm.gamma`             | `bert.encoder.layer.*.output.LayerNorm.weight`
   # | `encoder.transformer_cells.*.proj.bias`                        | `bert.encoder.layer.*.output.dense.bias`
   # | `encoder.transformer_cells.*.proj.weight`                      | `bert.encoder.layer.*.output.dense.weight`

   # Helper function to convert MXNET Arrays to PyTorch
   def to_torch(lowercase__  ) -> nn.Parameter:
      return nn.Parameter(torch.FloatTensor(mx_array.data().asnumpy()  )  )

   # Check param shapes and map new HF param back
   def check_and_map_params(lowercase__  ,    lowercase__  ):
      UpperCAmelCase_											=hf_param.shape

      UpperCAmelCase_											=to_torch(params[gluon_param]  )
      UpperCAmelCase_											=gluon_param.shape

      assert (
          shape_hf == shape_gluon
      ), F'The gluon parameter {gluon_param} has shape {shape_gluon}, but expects shape {shape_hf} for Transformers'

      return gluon_param

   UpperCAmelCase_											=check_and_map_params(
       hf_bort_model.bert.embeddings.word_embeddings.weight  ,    "word_embed.0.weight"  )
   UpperCAmelCase_											=check_and_map_params(
       hf_bort_model.bert.embeddings.position_embeddings.weight  ,    "encoder.position_weight"  )
   UpperCAmelCase_											=check_and_map_params(
       hf_bort_model.bert.embeddings.LayerNorm.bias  ,    "encoder.layer_norm.beta"  )
   UpperCAmelCase_											=check_and_map_params(
       hf_bort_model.bert.embeddings.LayerNorm.weight  ,    "encoder.layer_norm.gamma"  )

   # Inspired by RoBERTa conversion script, we just zero them out (Bort does not use them)
   UpperCAmelCase_											=torch.zeros_like(
       hf_bort_model.bert.embeddings.token_type_embeddings.weight.data  )

   for i in range(hf_bort_config.num_hidden_layers  ):
      UpperCAmelCase_											=hf_bort_model.bert.encoder.layer[i]

      # self attention
      UpperCAmelCase_											=layer.attention.self

      UpperCAmelCase_											=check_and_map_params(
          self_attn.key.bias.data  ,    F'encoder.transformer_cells.{i}.attention_cell.proj_key.bias'  )

      UpperCAmelCase_											=check_and_map_params(
          self_attn.key.weight.data  ,    F'encoder.transformer_cells.{i}.attention_cell.proj_key.weight'  )
      UpperCAmelCase_											=check_and_map_params(
          self_attn.query.bias.data  ,    F'encoder.transformer_cells.{i}.attention_cell.proj_query.bias'  )
      UpperCAmelCase_											=check_and_map_params(
          self_attn.query.weight.data  ,    F'encoder.transformer_cells.{i}.attention_cell.proj_query.weight'  )
      UpperCAmelCase_											=check_and_map_params(
          self_attn.value.bias.data  ,    F'encoder.transformer_cells.{i}.attention_cell.proj_value.bias'  )
      UpperCAmelCase_											=check_and_map_params(
          self_attn.value.weight.data  ,    F'encoder.transformer_cells.{i}.attention_cell.proj_value.weight'  )

      # self attention output
      UpperCAmelCase_											=layer.attention.output

      UpperCAmelCase_											=check_and_map_params(
          self_output.dense.bias  ,    F'encoder.transformer_cells.{i}.proj.bias'  )
      UpperCAmelCase_											=check_and_map_params(
          self_output.dense.weight  ,    F'encoder.transformer_cells.{i}.proj.weight'  )
      UpperCAmelCase_											=check_and_map_params(
          self_output.LayerNorm.bias  ,    F'encoder.transformer_cells.{i}.layer_norm.beta'  )
      UpperCAmelCase_											=check_and_map_params(
          self_output.LayerNorm.weight  ,    F'encoder.transformer_cells.{i}.layer_norm.gamma'  )

      # intermediate
      UpperCAmelCase_											=layer.intermediate

      UpperCAmelCase_											=check_and_map_params(
          intermediate.dense.bias  ,    F'encoder.transformer_cells.{i}.ffn.ffn_1.bias'  )
      UpperCAmelCase_											=check_and_map_params(
          intermediate.dense.weight  ,    F'encoder.transformer_cells.{i}.ffn.ffn_1.weight'  )

      # output
      UpperCAmelCase_											=layer.output

      UpperCAmelCase_											=check_and_map_params(
          bert_output.dense.bias  ,    F'encoder.transformer_cells.{i}.ffn.ffn_2.bias'  )
      UpperCAmelCase_											=check_and_map_params(
          bert_output.dense.weight  ,    F'encoder.transformer_cells.{i}.ffn.ffn_2.weight'  )
      UpperCAmelCase_											=check_and_map_params(
          bert_output.LayerNorm.bias  ,    F'encoder.transformer_cells.{i}.ffn.layer_norm.beta'  )
      UpperCAmelCase_											=check_and_map_params(
          bert_output.LayerNorm.weight  ,    F'encoder.transformer_cells.{i}.ffn.layer_norm.gamma'  )

   # Save space and energy 🎄
   hf_bort_model.half()

   # Compare output of both models
   UpperCAmelCase_											=RobertaTokenizer.from_pretrained("roberta-base"  )

   UpperCAmelCase_											=tokenizer.encode_plus(lowercase__  )["input_ids"]

   # Get gluon output
   UpperCAmelCase_											=mx.nd.array([input_ids]  )
   UpperCAmelCase_											=original_bort(inputs=lowercase__  ,    token_types=[]  )

   # Get Transformer output (save and reload model again)
   hf_bort_model.save_pretrained(lowercase__  )
   UpperCAmelCase_											=BertModel.from_pretrained(lowercase__  )
   hf_bort_model.eval()

   UpperCAmelCase_											=tokenizer.encode_plus(lowercase__  ,    return_tensors="pt"  )
   UpperCAmelCase_											=hf_bort_model(**lowercase__  )[0]

   UpperCAmelCase_											=output_gluon[0].asnumpy()
   UpperCAmelCase_											=output_hf[0].detach().numpy()

   UpperCAmelCase_											=np.max(np.abs(hf_layer - gluon_layer  )  ).item()
   UpperCAmelCase_											=np.allclose(lowercase__  ,    lowercase__  ,    atol=1E-3  )

   if success:
      print("✔️ Both model do output the same tensors"  )
   else:
      print("❌ Both model do **NOT** output the same tensors"  )
      print("Absolute difference is:"  ,    lowercase__  )


if __name__ == "__main__":
     __lowercase  :   Dict						      =argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         """--bort_checkpoint_path""", default=None, type=str, required=True, help="""Path the official Bort params file."""
     )
     parser.add_argument(
         """--pytorch_dump_folder_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
     )
     __lowercase  :   int						      =parser.parse_args()
     convert_bort_checkpoint_to_pytorch(args.bort_checkpoint_path, args.pytorch_dump_folder_path)











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.stable_diffusion_safe import StableDiffusionPipelineSafe as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import require_torch_gpu

class      A       (	unittest.TestCase						):



      def 	lowerCAmelCase__    (	self:						str       )						->    Optional[Any]:
         '''simple docstring'''

         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()



      @property
      def 	lowerCAmelCase__    (	self:						Optional[int]       )						->    List[str]:
         '''simple docstring'''

         UpperCAmelCase_											=1
         UpperCAmelCase_											=3
         UpperCAmelCase_											=(32, 32)

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



      @property
      def 	lowerCAmelCase__    (	self:						Any       )						->    Union[str, Any]:
         '''simple docstring'''

         torch.manual_seed(0       )
         UpperCAmelCase_											=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 	lowerCAmelCase__    (	self:						List[Any]       )						->    int:
         '''simple docstring'''

         torch.manual_seed(0       )
         UpperCAmelCase_											=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 	lowerCAmelCase__    (	self:						Dict       )						->    str:
         '''simple docstring'''

         torch.manual_seed(0       )
         UpperCAmelCase_											=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		,  )
         return CLIPTextModel(_lowerCAmelCase       )



      @property
      def 	lowerCAmelCase__    (	self:						Optional[int]       )						->    Optional[Any]:
         '''simple docstring'''

         def extract(*_lowerCAmelCase:						List[str]		,  **_lowerCAmelCase:						List[str]       ):

            class      A       :



                  def __init__(	self:						List[str]       )						->    Any:
                     '''simple docstring'''

                     UpperCAmelCase_											=torch.ones([0]       )






                  def 	lowerCAmelCase__    (	self:						List[str]		,  _lowerCAmelCase:						Dict       )						->    List[Any]:
                     '''simple docstring'''

                     self.pixel_values.to(_lowerCAmelCase       )
                     return self

            return Out()

         return extract



      def 	lowerCAmelCase__    (	self:						Tuple       )						->    Union[str, Any]:
         '''simple docstring'''

         UpperCAmelCase_											="cpu"  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase_											=self.dummy_cond_unet
         UpperCAmelCase_											=DDIMScheduler(
             beta_start=0.0_00_85		,  beta_end=0.0_12		,  beta_schedule="scaled_linear"		,  clip_sample=_lowerCAmelCase		,  set_alpha_to_one=_lowerCAmelCase		,  )

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

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

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

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

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

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

         assert image.shape == (1, 64, 64, 3)
         UpperCAmelCase_											=np.array([0.57_56, 0.61_18, 0.50_05, 0.50_41, 0.54_71, 0.47_26, 0.49_76, 0.48_65, 0.48_64]       )

         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 	lowerCAmelCase__    (	self:						Optional[int]       )						->    List[Any]:
         '''simple docstring'''

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

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

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

         UpperCAmelCase_											=output.images

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

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

         assert image.shape == (1, 64, 64, 3)
         UpperCAmelCase_											=np.array([0.51_25, 0.57_16, 0.48_28, 0.50_60, 0.56_50, 0.47_68, 0.51_85, 0.48_95, 0.49_93]       )

         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 	lowerCAmelCase__    (	self:						List[Any]       )						->    str:
         '''simple docstring'''

         UpperCAmelCase_											=StableDiffusionPipeline.from_pretrained(
             "hf-internal-testing/tiny-stable-diffusion-lms-pipe"		,  safety_checker=_lowerCAmelCase       )
         assert isinstance(_lowerCAmelCase		,  _lowerCAmelCase       )
         assert isinstance(pipe.scheduler		,  _lowerCAmelCase       )
         assert pipe.safety_checker is None

         UpperCAmelCase_											=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(_lowerCAmelCase       )
            UpperCAmelCase_											=StableDiffusionPipeline.from_pretrained(_lowerCAmelCase       )

         # sanity check that the pipeline still works
         assert pipe.safety_checker is None
         UpperCAmelCase_											=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 	lowerCAmelCase__    (	self:						int       )						->    Optional[int]:
         '''simple docstring'''

         UpperCAmelCase_											=self.dummy_cond_unet
         UpperCAmelCase_											=PNDMScheduler(skip_prk_steps=_lowerCAmelCase       )
         UpperCAmelCase_											=self.dummy_vae
         UpperCAmelCase_											=self.dummy_text_encoder
         UpperCAmelCase_											=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )

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

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

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

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

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



      def 	lowerCAmelCase__    (	self:						List[Any]       )						->    Tuple:
         '''simple docstring'''

         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()



      def 	lowerCAmelCase__    (	self:						Union[str, Any]       )						->    Union[str, Any]:
         '''simple docstring'''

         UpperCAmelCase_											=StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"		,  safety_checker=_lowerCAmelCase       )
         UpperCAmelCase_											=LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config       )
         UpperCAmelCase_											=sd_pipe.to(_lowerCAmelCase       )
         sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         UpperCAmelCase_											=(
             "portrait of girl with smokey eyes makeup in abandoned hotel, grange clothes, redshift, wide high angle"
             " coloured polaroid photograph with flash, kodak film, hyper real, stunning moody cinematography, with"
             " anamorphic lenses, by maripol, fallen angels by wong kar - wai, style of suspiria and neon demon and"
             " children from bahnhof zoo, detailed "
         )
         UpperCAmelCase_											=40_0366_0346
         UpperCAmelCase_											=7

         # without safety guidance (sld_guidance_scale = 0)
         UpperCAmelCase_											=torch.manual_seed(_lowerCAmelCase       )
         UpperCAmelCase_											=sd_pipe(
             [prompt]		,  generator=_lowerCAmelCase		,  guidance_scale=_lowerCAmelCase		,  num_inference_steps=50		,  output_type="np"		,  width=512		,  height=512		,  sld_guidance_scale=0		,  )

         UpperCAmelCase_											=output.images
         UpperCAmelCase_											=image[0, -3:, -3:, -1]
         UpperCAmelCase_											=[0.22_78, 0.22_31, 0.22_49, 0.23_33, 0.23_03, 0.18_85, 0.22_73, 0.21_44, 0.21_76]

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

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

         # without safety guidance (strong configuration)
         UpperCAmelCase_											=torch.manual_seed(_lowerCAmelCase       )
         UpperCAmelCase_											=sd_pipe(
             [prompt]		,  generator=_lowerCAmelCase		,  guidance_scale=_lowerCAmelCase		,  num_inference_steps=50		,  output_type="np"		,  width=512		,  height=512		,  sld_guidance_scale=2000		,  sld_warmup_steps=7		,  sld_threshold=0.0_25		,  sld_momentum_scale=0.5		,  sld_mom_beta=0.7		,  )

         UpperCAmelCase_											=output.images
         UpperCAmelCase_											=image[0, -3:, -3:, -1]
         UpperCAmelCase_											=[0.23_83, 0.22_76, 0.2_36, 0.21_92, 0.21_86, 0.20_53, 0.19_71, 0.19_01, 0.17_19]

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

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



      def 	lowerCAmelCase__    (	self:						Any       )						->    Union[str, Any]:
         '''simple docstring'''

         UpperCAmelCase_											=StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"		,  safety_checker=_lowerCAmelCase       )
         UpperCAmelCase_											=LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config       )
         UpperCAmelCase_											=sd_pipe.to(_lowerCAmelCase       )
         sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         UpperCAmelCase_											="padme amidala taking a bath artwork, safe for work, no nudity"
         UpperCAmelCase_											=27_3497_1755
         UpperCAmelCase_											=7

         UpperCAmelCase_											=torch.manual_seed(_lowerCAmelCase       )
         UpperCAmelCase_											=sd_pipe(
             [prompt]		,  generator=_lowerCAmelCase		,  guidance_scale=_lowerCAmelCase		,  num_inference_steps=50		,  output_type="np"		,  width=512		,  height=512		,  sld_guidance_scale=0		,  )

         UpperCAmelCase_											=output.images
         UpperCAmelCase_											=image[0, -3:, -3:, -1]
         UpperCAmelCase_											=[0.35_02, 0.36_22, 0.33_96, 0.36_42, 0.34_78, 0.33_18, 0.35, 0.33_48, 0.32_97]

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

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

         UpperCAmelCase_											=torch.manual_seed(_lowerCAmelCase       )
         UpperCAmelCase_											=sd_pipe(
             [prompt]		,  generator=_lowerCAmelCase		,  guidance_scale=_lowerCAmelCase		,  num_inference_steps=50		,  output_type="np"		,  width=512		,  height=512		,  sld_guidance_scale=2000		,  sld_warmup_steps=7		,  sld_threshold=0.0_25		,  sld_momentum_scale=0.5		,  sld_mom_beta=0.7		,  )

         UpperCAmelCase_											=output.images
         UpperCAmelCase_											=image[0, -3:, -3:, -1]
         UpperCAmelCase_											=[0.55_31, 0.52_06, 0.48_95, 0.51_56, 0.51_82, 0.47_51, 0.48_02, 0.48_03, 0.44_43]

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

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






      def 	lowerCAmelCase__    (	self:						str       )						->    List[Any]:
         '''simple docstring'''

         UpperCAmelCase_											=StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"       )
         UpperCAmelCase_											=sd_pipe.to(_lowerCAmelCase       )
         sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase       )

         UpperCAmelCase_											=(
             "the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c."
             " leyendecker"
         )
         UpperCAmelCase_											=10_4435_5234
         UpperCAmelCase_											=12

         UpperCAmelCase_											=torch.manual_seed(_lowerCAmelCase       )
         UpperCAmelCase_											=sd_pipe(
             [prompt]		,  generator=_lowerCAmelCase		,  guidance_scale=_lowerCAmelCase		,  num_inference_steps=50		,  output_type="np"		,  width=512		,  height=512		,  sld_guidance_scale=0		,  )

         UpperCAmelCase_											=output.images
         UpperCAmelCase_											=image[0, -3:, -3:, -1]
         UpperCAmelCase_											=np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]       )

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

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

         UpperCAmelCase_											=torch.manual_seed(_lowerCAmelCase       )
         UpperCAmelCase_											=sd_pipe(
             [prompt]		,  generator=_lowerCAmelCase		,  guidance_scale=_lowerCAmelCase		,  num_inference_steps=50		,  output_type="np"		,  width=512		,  height=512		,  sld_guidance_scale=2000		,  sld_warmup_steps=7		,  sld_threshold=0.0_25		,  sld_momentum_scale=0.5		,  sld_mom_beta=0.7		,  )

         UpperCAmelCase_											=output.images
         UpperCAmelCase_											=image[0, -3:, -3:, -1]
         UpperCAmelCase_											=np.array([0.58_18, 0.62_85, 0.68_35, 0.60_19, 0.6_25, 0.67_54, 0.60_96, 0.63_34, 0.65_61]       )
         assert image.shape == (1, 512, 512, 3)

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




