


import argparse
import json

import torch

from diffusers import DDPMScheduler, LDMPipeline, UNetaDModel, VQModel




def     a__ (					lowercase__  ,    lowercase__=1  ):






   '''simple docstring'''






   if n_shave_prefix_segments >= 0:
      return ".".join(path.split("."  )[n_shave_prefix_segments:]  )
   else:
      return ".".join(path.split("."  )[:n_shave_prefix_segments]  )




def     a__ (					lowercase__  ,    lowercase__=0  ):






   '''simple docstring'''






   UpperCAmelCase_											=[]
   for old_item in old_list:
      UpperCAmelCase_											=old_item.replace("in_layers.0"  ,    "norm1"  )
      UpperCAmelCase_											=new_item.replace("in_layers.2"  ,    "conv1"  )

      UpperCAmelCase_											=new_item.replace("out_layers.0"  ,    "norm2"  )
      UpperCAmelCase_											=new_item.replace("out_layers.3"  ,    "conv2"  )

      UpperCAmelCase_											=new_item.replace("emb_layers.1"  ,    "time_emb_proj"  )
      UpperCAmelCase_											=new_item.replace("skip_connection"  ,    "conv_shortcut"  )

      UpperCAmelCase_											=shave_segments(snake_case__  ,    n_shave_prefix_segments=snake_case__  )

      mapping.append({"old": old_item, "new": new_item}  )

   return mapping




def     a__ (					lowercase__  ,    lowercase__=0  ):






   '''simple docstring'''






   UpperCAmelCase_											=[]
   for old_item in old_list:
      UpperCAmelCase_											=old_item

      UpperCAmelCase_											=new_item.replace("norm.weight"  ,    "group_norm.weight"  )
      UpperCAmelCase_											=new_item.replace("norm.bias"  ,    "group_norm.bias"  )

      UpperCAmelCase_											=new_item.replace("proj_out.weight"  ,    "proj_attn.weight"  )
      UpperCAmelCase_											=new_item.replace("proj_out.bias"  ,    "proj_attn.bias"  )

      UpperCAmelCase_											=shave_segments(snake_case__  ,    n_shave_prefix_segments=snake_case__  )

      mapping.append({"old": old_item, "new": new_item}  )

   return mapping




def     a__ (					lowercase__  ,    lowercase__  ,    lowercase__  ,    lowercase__=None  ,    lowercase__=None  ,    lowercase__=None  ):






   '''simple docstring'''






   assert isinstance(snake_case__  ,    snake_case__  ), "Paths should be a list of dicts containing 'old' and 'new' keys."

   # Splits the attention layers into three variables.
   if attention_paths_to_split is not None:
      for path, path_map in attention_paths_to_split.items():
         UpperCAmelCase_											=old_checkpoint[path]
         UpperCAmelCase_											=old_tensor.shape[0] // 3

         UpperCAmelCase_											=(-1, channels) if len(old_tensor.shape  ) == 3 else (-1)

         UpperCAmelCase_											=old_tensor.shape[0] // config["""num_head_channels"""] // 3

         UpperCAmelCase_											=old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]  )
         UpperCAmelCase_											=old_tensor.split(channels // num_heads  ,    dim=1  )

         UpperCAmelCase_											=query.reshape(snake_case__  )
         UpperCAmelCase_											=key.reshape(snake_case__  )
         UpperCAmelCase_											=value.reshape(snake_case__  )

   for path in paths:
      UpperCAmelCase_											=path["""new"""]

      # These have already been assigned
      if attention_paths_to_split is not None and new_path in attention_paths_to_split:
         continue

      # Global renaming happens here
      UpperCAmelCase_											=new_path.replace("middle_block.0"  ,    "mid_block.resnets.0"  )
      UpperCAmelCase_											=new_path.replace("middle_block.1"  ,    "mid_block.attentions.0"  )
      UpperCAmelCase_											=new_path.replace("middle_block.2"  ,    "mid_block.resnets.1"  )

      if additional_replacements is not None:
         for replacement in additional_replacements:
            UpperCAmelCase_											=new_path.replace(replacement["old"]  ,    replacement["new"]  )

            # proj_attn.weight has to be converted from conv 1D to linear
      if "proj_attn.weight" in new_path:
         UpperCAmelCase_											=old_checkpoint[path["""old"""]][:, :, 0]
      else:
         UpperCAmelCase_											=old_checkpoint[path["""old"""]]




def     a__ (					lowercase__  ,    lowercase__  ):






   '''simple docstring'''






   UpperCAmelCase_											={}

   UpperCAmelCase_											=checkpoint["""time_embed.0.weight"""]
   UpperCAmelCase_											=checkpoint["""time_embed.0.bias"""]
   UpperCAmelCase_											=checkpoint["""time_embed.2.weight"""]
   UpperCAmelCase_											=checkpoint["""time_embed.2.bias"""]

   UpperCAmelCase_											=checkpoint["""input_blocks.0.0.weight"""]
   UpperCAmelCase_											=checkpoint["""input_blocks.0.0.bias"""]

   UpperCAmelCase_											=checkpoint["""out.0.weight"""]
   UpperCAmelCase_											=checkpoint["""out.0.bias"""]
   UpperCAmelCase_											=checkpoint["""out.2.weight"""]
   UpperCAmelCase_											=checkpoint["""out.2.bias"""]

   # Retrieves the keys for the input blocks only
   UpperCAmelCase_											=len({".".join(layer.split("."  )[:2]  ) for layer in checkpoint if "input_blocks" in layer}  )
   UpperCAmelCase_											={
       layer_id: [key for key in checkpoint if F'input_blocks.{layer_id}' in key]
       for layer_id in range(snake_case__  )
   }

   # Retrieves the keys for the middle blocks only
   UpperCAmelCase_											=len({".".join(layer.split("."  )[:2]  ) for layer in checkpoint if "middle_block" in layer}  )
   UpperCAmelCase_											={
       layer_id: [key for key in checkpoint if F'middle_block.{layer_id}' in key]
       for layer_id in range(snake_case__  )
   }

   # Retrieves the keys for the output blocks only
   UpperCAmelCase_											=len({".".join(layer.split("."  )[:2]  ) for layer in checkpoint if "output_blocks" in layer}  )
   UpperCAmelCase_											={
       layer_id: [key for key in checkpoint if F'output_blocks.{layer_id}' in key]
       for layer_id in range(snake_case__  )
   }

   for i in range(1  ,    snake_case__  ):
      UpperCAmelCase_											=(i - 1) // (config["""num_res_blocks"""] + 1)
      UpperCAmelCase_											=(i - 1) % (config["""num_res_blocks"""] + 1)

      UpperCAmelCase_											=[key for key in input_blocks[i] if F'input_blocks.{i}.0' in key]
      UpperCAmelCase_											=[key for key in input_blocks[i] if F'input_blocks.{i}.1' in key]

      if F'input_blocks.{i}.0.op.weight' in checkpoint:
         UpperCAmelCase_											=checkpoint[
             F'input_blocks.{i}.0.op.weight'
         ]
         UpperCAmelCase_											=checkpoint[
             F'input_blocks.{i}.0.op.bias'
         ]
         continue

      UpperCAmelCase_											=renew_resnet_paths(snake_case__  )
      UpperCAmelCase_											={"""old""": F'input_blocks.{i}.0', """new""": F'down_blocks.{block_id}.resnets.{layer_in_block_id}'}
      UpperCAmelCase_											={"""old""": """resnets.2.op""", """new""": """downsamplers.0.op"""}
      assign_to_checkpoint(
          snake_case__  ,    snake_case__  ,    snake_case__  ,    additional_replacements=[meta_path, resnet_op]  ,    config=snake_case__  )

      if len(snake_case__  ):
         UpperCAmelCase_											=renew_attention_paths(snake_case__  )
         UpperCAmelCase_											={
             """old""": F'input_blocks.{i}.1',
             """new""": F'down_blocks.{block_id}.attentions.{layer_in_block_id}',
         }
         UpperCAmelCase_											={
             F'input_blocks.{i}.1.qkv.bias': {
                 """key""": F'down_blocks.{block_id}.attentions.{layer_in_block_id}.key.bias',
                 """query""": F'down_blocks.{block_id}.attentions.{layer_in_block_id}.query.bias',
                 """value""": F'down_blocks.{block_id}.attentions.{layer_in_block_id}.value.bias',
             },
             F'input_blocks.{i}.1.qkv.weight': {
                 """key""": F'down_blocks.{block_id}.attentions.{layer_in_block_id}.key.weight',
                 """query""": F'down_blocks.{block_id}.attentions.{layer_in_block_id}.query.weight',
                 """value""": F'down_blocks.{block_id}.attentions.{layer_in_block_id}.value.weight',
             },
         }
         assign_to_checkpoint(
             snake_case__  ,    snake_case__  ,    snake_case__  ,    additional_replacements=[meta_path]  ,    attention_paths_to_split=snake_case__  ,    config=snake_case__  ,    )

   UpperCAmelCase_											=middle_blocks[0]
   UpperCAmelCase_											=middle_blocks[1]
   UpperCAmelCase_											=middle_blocks[2]

   UpperCAmelCase_											=renew_resnet_paths(snake_case__  )
   assign_to_checkpoint(snake_case__  ,    snake_case__  ,    snake_case__  ,    config=snake_case__  )

   UpperCAmelCase_											=renew_resnet_paths(snake_case__  )
   assign_to_checkpoint(snake_case__  ,    snake_case__  ,    snake_case__  ,    config=snake_case__  )

   UpperCAmelCase_											=renew_attention_paths(snake_case__  )
   UpperCAmelCase_											={
       """middle_block.1.qkv.bias""": {
           """key""": """mid_block.attentions.0.key.bias""",
           """query""": """mid_block.attentions.0.query.bias""",
           """value""": """mid_block.attentions.0.value.bias""",
       },
       """middle_block.1.qkv.weight""": {
           """key""": """mid_block.attentions.0.key.weight""",
           """query""": """mid_block.attentions.0.query.weight""",
           """value""": """mid_block.attentions.0.value.weight""",
       },
   }
   assign_to_checkpoint(
       snake_case__  ,    snake_case__  ,    snake_case__  ,    attention_paths_to_split=snake_case__  ,    config=snake_case__  )

   for i in range(snake_case__  ):
      UpperCAmelCase_											=i // (config["""num_res_blocks"""] + 1)
      UpperCAmelCase_											=i % (config["""num_res_blocks"""] + 1)
      UpperCAmelCase_											=[shave_segments(snake_case__  ,    2  ) for name in output_blocks[i]]
      UpperCAmelCase_											={}

      for layer in output_block_layers:
         UpperCAmelCase_											=layer.split("."  )[0], shave_segments(snake_case__  ,    1  )
         if layer_id in output_block_list:
            output_block_list[layer_id].append(snake_case__  )
         else:
            UpperCAmelCase_											=[layer_name]

      if len(snake_case__  ) > 1:
         UpperCAmelCase_											=[key for key in output_blocks[i] if F'output_blocks.{i}.0' in key]
         UpperCAmelCase_											=[key for key in output_blocks[i] if F'output_blocks.{i}.1' in key]

         UpperCAmelCase_											=renew_resnet_paths(snake_case__  )
         UpperCAmelCase_											=renew_resnet_paths(snake_case__  )

         UpperCAmelCase_											={"""old""": F'output_blocks.{i}.0', """new""": F'up_blocks.{block_id}.resnets.{layer_in_block_id}'}
         assign_to_checkpoint(snake_case__  ,    snake_case__  ,    snake_case__  ,    additional_replacements=[meta_path]  ,    config=snake_case__  )

         if ["conv.weight", "conv.bias"] in output_block_list.values():
            UpperCAmelCase_											=list(output_block_list.values()  ).index(["conv.weight", "conv.bias"]  )
            UpperCAmelCase_											=checkpoint[
                F'output_blocks.{i}.{index}.conv.weight'
            ]
            UpperCAmelCase_											=checkpoint[
                F'output_blocks.{i}.{index}.conv.bias'
            ]

            # Clear attentions as they have been attributed above.
            if len(snake_case__  ) == 2:
               UpperCAmelCase_											=[]

         if len(snake_case__  ):
            UpperCAmelCase_											=renew_attention_paths(snake_case__  )
            UpperCAmelCase_											={
                """old""": F'output_blocks.{i}.1',
                """new""": F'up_blocks.{block_id}.attentions.{layer_in_block_id}',
            }
            UpperCAmelCase_											={
                F'output_blocks.{i}.1.qkv.bias': {
                    """key""": F'up_blocks.{block_id}.attentions.{layer_in_block_id}.key.bias',
                    """query""": F'up_blocks.{block_id}.attentions.{layer_in_block_id}.query.bias',
                    """value""": F'up_blocks.{block_id}.attentions.{layer_in_block_id}.value.bias',
                },
                F'output_blocks.{i}.1.qkv.weight': {
                    """key""": F'up_blocks.{block_id}.attentions.{layer_in_block_id}.key.weight',
                    """query""": F'up_blocks.{block_id}.attentions.{layer_in_block_id}.query.weight',
                    """value""": F'up_blocks.{block_id}.attentions.{layer_in_block_id}.value.weight',
                },
            }
            assign_to_checkpoint(
                snake_case__  ,    snake_case__  ,    snake_case__  ,    additional_replacements=[meta_path]  ,    attention_paths_to_split=to_split if any("qkv" in key for key in attentions  ) else None  ,    config=snake_case__  ,    )
      else:
         UpperCAmelCase_											=renew_resnet_paths(snake_case__  ,    n_shave_prefix_segments=1  )
         for path in resnet_0_paths:
            UpperCAmelCase_											=""".""".join(["output_blocks", str(snake_case__  ), path["old"]]  )
            UpperCAmelCase_											=""".""".join(["up_blocks", str(snake_case__  ), "resnets", str(snake_case__  ), path["new"]]  )

            UpperCAmelCase_											=checkpoint[old_path]

   return new_checkpoint


if __name__ == "__main__":
     __lowercase  :   Dict						      =argparse.ArgumentParser()

     parser.add_argument(
         """--checkpoint_path""", default=None, type=str, required=True, help="""Path to the checkpoint to convert."""
     )

     parser.add_argument(
         """--config_file""",
         default=None,
         type=str,
         required=True,
         help="""The config json file corresponding to the architecture.""",
     )

     parser.add_argument("""--dump_path""", default=None, type=str, required=True, help="""Path to the output model.""")

     __lowercase  :   str						      =parser.parse_args()

     __lowercase  :   Dict						      =torch.load(args.checkpoint_path)

     with open(args.config_file) as f:
          __lowercase  :   Any						      =json.loads(f.read())

     __lowercase  :   Union[str, Any]						      =convert_ldm_checkpoint(checkpoint, config)

     if "ldm" in config:
          del config["ldm"]

     __lowercase  :   Union[str, Any]						      =UNetaDModel(**config)
     model.load_state_dict(converted_checkpoint)

     try:
          __lowercase  :   Union[str, Any]						      =DDPMScheduler.from_config("""/""".join(args.checkpoint_path.split("""/""")[:-1]))
          __lowercase  :   List[str]						      =VQModel.from_pretrained("""/""".join(args.checkpoint_path.split("""/""")[:-1]))

          __lowercase  :   int						      =LDMPipeline(unet=model, scheduler=scheduler, vae=vqvae)
          pipe.save_pretrained(args.dump_path)
     except:  # noqa: E722
          model.save_pretrained(args.dump_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




