



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 				snake_case__(unittest.TestCase      ):

     """simple docstring"""






     def 			snake_case					(							self :     int					):
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()







     @property
     def 			snake_case					(							self :     List[str]					):
         lowercase__					:   List[Any]													=					1
         lowercase__					:   Union[str, Any]													=					3
         lowercase__					:   Union[str, Any]													=					(32, 32)

         lowercase__					:   Tuple													=					floats_tensor((batch_size, num_channels) + sizes				,       rng=random.Random(0					)					).to(UpperCamelCase__					)
         return image







     @property
     def 			snake_case					(							self :     Optional[int]					):
         torch.manual_seed(0					)
         lowercase__					:   Any													=					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 			snake_case					(							self :     List[Any]					):
         torch.manual_seed(0					)
         lowercase__					:   Union[str, Any]													=					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 			snake_case					(							self :     Optional[Any]					):
         torch.manual_seed(0					)
         lowercase__					:   Dict													=					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=1_000				,       )
         return CLIPTextModel(UpperCamelCase__					)







     @property
     def 			snake_case					(							self :     Any					):


         def extract(*SCREAMING_SNAKE_CASE :     str				,       **SCREAMING_SNAKE_CASE :     Union[str, Any]					):

             class 				snake_case__:

                  """simple docstring"""






                  def __init__(							self :     Dict					):
                      lowercase__					:   Union[str, Any]													=					torch.ones([0]					)







                  def 			snake_case					(							self :     Any				,       SCREAMING_SNAKE_CASE :     int					):
                      self.pixel_values.to(UpperCamelCase__					)
                      return self






             return Out()

         return extract







     def 			snake_case					(							self :     Dict					):
         lowercase__					:   Any													=					"cpu"  # ensure determinism for the device-dependent torch.Generator
         lowercase__					:   str													=					self.dummy_cond_unet
         lowercase__					:   Any													=					DDIMScheduler(
             beta_start=0.00_085				,       beta_end=0.012				,       beta_schedule="scaled_linear"				,       clip_sample=UpperCamelCase__				,       set_alpha_to_one=UpperCamelCase__				,       )

         lowercase__					:   Tuple													=					self.dummy_vae
         lowercase__					:   List[str]													=					self.dummy_text_encoder
         lowercase__					:   Any													=					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"					)

         # make sure here that pndm scheduler skips prk
         lowercase__					:   List[str]													=					StableDiffusionPipeline(
             unet=UpperCamelCase__				,       scheduler=UpperCamelCase__				,       vae=UpperCamelCase__				,       text_encoder=UpperCamelCase__				,       tokenizer=UpperCamelCase__				,       safety_checker=UpperCamelCase__				,       feature_extractor=self.dummy_extractor				,       )
         lowercase__					:   List[Any]													=					sd_pipe.to(UpperCamelCase__					)
         sd_pipe.set_progress_bar_config(disable=UpperCamelCase__					)

         lowercase__					:   Any													=					"A painting of a squirrel eating a burger"

         lowercase__					:   Dict													=					torch.Generator(device=UpperCamelCase__					).manual_seed(0					)
         lowercase__					:   Optional[int]													=					sd_pipe([prompt]				,       generator=UpperCamelCase__				,       guidance_scale=6.0				,       num_inference_steps=2				,       output_type="np"					)
         lowercase__					:   Optional[Any]													=					output.images

         lowercase__					:   Optional[Any]													=					torch.Generator(device=UpperCamelCase__					).manual_seed(0					)
         lowercase__					:   Tuple													=					sd_pipe(
             [prompt]				,       generator=UpperCamelCase__				,       guidance_scale=6.0				,       num_inference_steps=2				,       output_type="np"				,       return_dict=UpperCamelCase__				,       )[0]

         lowercase__					:   int													=					image[0, -3:, -3:, -1]
         lowercase__					:   Optional[Any]													=					image_from_tuple[0, -3:, -3:, -1]

         assert image.shape == (1, 64, 64, 3)
         lowercase__					:   int													=					np.array([0.5_756, 0.6_118, 0.5_005, 0.5_041, 0.5_471, 0.4_726, 0.4_976, 0.4_865, 0.4_864]					)

         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 			snake_case					(							self :     List[Any]					):
         lowercase__					:   Dict													=					"cpu"  # ensure determinism for the device-dependent torch.Generator
         lowercase__					:   Any													=					self.dummy_cond_unet
         lowercase__					:   List[Any]													=					PNDMScheduler(skip_prk_steps=UpperCamelCase__					)
         lowercase__					:   List[Any]													=					self.dummy_vae
         lowercase__					:   int													=					self.dummy_text_encoder
         lowercase__					:   List[str]													=					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"					)

         # make sure here that pndm scheduler skips prk
         lowercase__					:   Any													=					StableDiffusionPipeline(
             unet=UpperCamelCase__				,       scheduler=UpperCamelCase__				,       vae=UpperCamelCase__				,       text_encoder=UpperCamelCase__				,       tokenizer=UpperCamelCase__				,       safety_checker=UpperCamelCase__				,       feature_extractor=self.dummy_extractor				,       )
         lowercase__					:   str													=					sd_pipe.to(UpperCamelCase__					)
         sd_pipe.set_progress_bar_config(disable=UpperCamelCase__					)

         lowercase__					:   Optional[int]													=					"A painting of a squirrel eating a burger"
         lowercase__					:   List[str]													=					torch.Generator(device=UpperCamelCase__					).manual_seed(0					)
         lowercase__					:   List[Any]													=					sd_pipe([prompt]				,       generator=UpperCamelCase__				,       guidance_scale=6.0				,       num_inference_steps=2				,       output_type="np"					)

         lowercase__					:   Optional[Any]													=					output.images

         lowercase__					:   int													=					torch.Generator(device=UpperCamelCase__					).manual_seed(0					)
         lowercase__					:   Dict													=					sd_pipe(
             [prompt]				,       generator=UpperCamelCase__				,       guidance_scale=6.0				,       num_inference_steps=2				,       output_type="np"				,       return_dict=UpperCamelCase__				,       )[0]

         lowercase__					:   Dict													=					image[0, -3:, -3:, -1]
         lowercase__					:   List[str]													=					image_from_tuple[0, -3:, -3:, -1]

         assert image.shape == (1, 64, 64, 3)
         lowercase__					:   Optional[Any]													=					np.array([0.5_125, 0.5_716, 0.4_828, 0.5_060, 0.5_650, 0.4_768, 0.5_185, 0.4_895, 0.4_993]					)

         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 			snake_case					(							self :     List[Any]					):
         lowercase__					:   Union[str, Any]													=					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

         lowercase__					:   Optional[int]													=					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__					)
             lowercase__					:   List[Any]													=					StableDiffusionPipeline.from_pretrained(UpperCamelCase__					)

         # sanity check that the pipeline still works
         assert pipe.safety_checker is None
         lowercase__					:   Optional[int]													=					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 			snake_case					(							self :     List[str]					):
         lowercase__					:   Union[str, Any]													=					self.dummy_cond_unet
         lowercase__					:   Union[str, Any]													=					PNDMScheduler(skip_prk_steps=UpperCamelCase__					)
         lowercase__					:   List[Any]													=					self.dummy_vae
         lowercase__					:   List[Any]													=					self.dummy_text_encoder
         lowercase__					:   Union[str, Any]													=					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"					)

         # put models in fp16
         lowercase__					:   str													=					unet.half()
         lowercase__					:   Dict													=					vae.half()
         lowercase__					:   str													=					bert.half()

         # make sure here that pndm scheduler skips prk
         lowercase__					:   Union[str, Any]													=					StableDiffusionPipeline(
             unet=UpperCamelCase__				,       scheduler=UpperCamelCase__				,       vae=UpperCamelCase__				,       text_encoder=UpperCamelCase__				,       tokenizer=UpperCamelCase__				,       safety_checker=UpperCamelCase__				,       feature_extractor=self.dummy_extractor				,       )
         lowercase__					:   Dict													=					sd_pipe.to(UpperCamelCase__					)
         sd_pipe.set_progress_bar_config(disable=UpperCamelCase__					)

         lowercase__					:   Tuple													=					"A painting of a squirrel eating a burger"
         lowercase__					:   Union[str, Any]													=					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 			snake_case					(							self :     Dict					):
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()







     def 			snake_case					(							self :     Dict					):
         lowercase__					:   Optional[int]													=					StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"				,       safety_checker=UpperCamelCase__					)
         lowercase__					:   Union[str, Any]													=					LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config					)
         lowercase__					:   int													=					sd_pipe.to(UpperCamelCase__					)
         sd_pipe.set_progress_bar_config(disable=UpperCamelCase__					)

         lowercase__					:   Dict													=					(
             "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 "
         )
         lowercase__					:   int													=					4_003_660_346
         lowercase__					:   List[Any]													=					7

         # without safety guidance (sld_guidance_scale = 0)
         lowercase__					:   int													=					torch.manual_seed(UpperCamelCase__					)
         lowercase__					:   List[str]													=					sd_pipe(
             [prompt]				,       generator=UpperCamelCase__				,       guidance_scale=UpperCamelCase__				,       num_inference_steps=50				,       output_type="np"				,       width=512				,       height=512				,       sld_guidance_scale=0				,       )

         lowercase__					:   Any													=					output.images
         lowercase__					:   str													=					image[0, -3:, -3:, -1]
         lowercase__					:   Dict													=					[0.2_278, 0.2_231, 0.2_249, 0.2_333, 0.2_303, 0.1_885, 0.2_273, 0.2_144, 0.2_176]

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

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

         # without safety guidance (strong configuration)
         lowercase__					:   Tuple													=					torch.manual_seed(UpperCamelCase__					)
         lowercase__					:   List[str]													=					sd_pipe(
             [prompt]				,       generator=UpperCamelCase__				,       guidance_scale=UpperCamelCase__				,       num_inference_steps=50				,       output_type="np"				,       width=512				,       height=512				,       sld_guidance_scale=2_000				,       sld_warmup_steps=7				,       sld_threshold=0.025				,       sld_momentum_scale=0.5				,       sld_mom_beta=0.7				,       )

         lowercase__					:   str													=					output.images
         lowercase__					:   int													=					image[0, -3:, -3:, -1]
         lowercase__					:   Optional[int]													=					[0.2_383, 0.2_276, 0.236, 0.2_192, 0.2_186, 0.2_053, 0.1_971, 0.1_901, 0.1_719]

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

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







     def 			snake_case					(							self :     str					):
         lowercase__					:   Tuple													=					StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"				,       safety_checker=UpperCamelCase__					)
         lowercase__					:   Optional[int]													=					LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config					)
         lowercase__					:   Optional[int]													=					sd_pipe.to(UpperCamelCase__					)
         sd_pipe.set_progress_bar_config(disable=UpperCamelCase__					)

         lowercase__					:   Any													=					"padme amidala taking a bath artwork, safe for work, no nudity"
         lowercase__					:   Optional[int]													=					2_734_971_755
         lowercase__					:   Optional[int]													=					7

         lowercase__					:   int													=					torch.manual_seed(UpperCamelCase__					)
         lowercase__					:   int													=					sd_pipe(
             [prompt]				,       generator=UpperCamelCase__				,       guidance_scale=UpperCamelCase__				,       num_inference_steps=50				,       output_type="np"				,       width=512				,       height=512				,       sld_guidance_scale=0				,       )

         lowercase__					:   Any													=					output.images
         lowercase__					:   List[str]													=					image[0, -3:, -3:, -1]
         lowercase__					:   str													=					[0.3_502, 0.3_622, 0.3_396, 0.3_642, 0.3_478, 0.3_318, 0.35, 0.3_348, 0.3_297]

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

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

         lowercase__					:   Tuple													=					torch.manual_seed(UpperCamelCase__					)
         lowercase__					:   List[str]													=					sd_pipe(
             [prompt]				,       generator=UpperCamelCase__				,       guidance_scale=UpperCamelCase__				,       num_inference_steps=50				,       output_type="np"				,       width=512				,       height=512				,       sld_guidance_scale=2_000				,       sld_warmup_steps=7				,       sld_threshold=0.025				,       sld_momentum_scale=0.5				,       sld_mom_beta=0.7				,       )

         lowercase__					:   str													=					output.images
         lowercase__					:   str													=					image[0, -3:, -3:, -1]
         lowercase__					:   Any													=					[0.5_531, 0.5_206, 0.4_895, 0.5_156, 0.5_182, 0.4_751, 0.4_802, 0.4_803, 0.4_443]

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

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







     def 			snake_case					(							self :     Optional[int]					):
         lowercase__					:   Tuple													=					StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"					)
         lowercase__					:   List[str]													=					sd_pipe.to(UpperCamelCase__					)
         sd_pipe.set_progress_bar_config(disable=UpperCamelCase__					)

         lowercase__					:   Dict													=					(
             "the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c."
             " leyendecker"
         )
         lowercase__					:   List[Any]													=					1_044_355_234
         lowercase__					:   Union[str, Any]													=					12

         lowercase__					:   Tuple													=					torch.manual_seed(UpperCamelCase__					)
         lowercase__					:   List[str]													=					sd_pipe(
             [prompt]				,       generator=UpperCamelCase__				,       guidance_scale=UpperCamelCase__				,       num_inference_steps=50				,       output_type="np"				,       width=512				,       height=512				,       sld_guidance_scale=0				,       )

         lowercase__					:   int													=					output.images
         lowercase__					:   Optional[int]													=					image[0, -3:, -3:, -1]
         lowercase__					:   List[str]													=					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

         lowercase__					:   List[Any]													=					torch.manual_seed(UpperCamelCase__					)
         lowercase__					:   int													=					sd_pipe(
             [prompt]				,       generator=UpperCamelCase__				,       guidance_scale=UpperCamelCase__				,       num_inference_steps=50				,       output_type="np"				,       width=512				,       height=512				,       sld_guidance_scale=2_000				,       sld_warmup_steps=7				,       sld_threshold=0.025				,       sld_momentum_scale=0.5				,       sld_mom_beta=0.7				,       )

         lowercase__					:   int													=					output.images
         lowercase__					:   Optional[int]													=					image[0, -3:, -3:, -1]
         lowercase__					:   List[str]													=					np.array([0.5_818, 0.6_285, 0.6_835, 0.6_019, 0.625, 0.6_754, 0.6_096, 0.6_334, 0.6_561]					)
         assert image.shape == (1, 512, 512, 3)

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









import argparse
import json

import torch

from diffusers import DDPMScheduler, LDMPipeline, UNetaDModel, VQModel


def  __lowerCamelCase (       lowerCamelCase__					,			lowerCamelCase__=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  __lowerCamelCase (       lowerCamelCase__					,			lowerCamelCase__=0						):

				"""simple docstring"""


				lowercase__					:   List[str]													=					[]
				for old_item in old_list:
								lowercase__					:   Optional[Any]													=					old_item.replace("in_layers.0"					,			"norm1"						)
								lowercase__					:   Union[str, Any]													=					new_item.replace("in_layers.2"					,			"conv1"						)

								lowercase__					:   Optional[Any]													=					new_item.replace("out_layers.0"					,			"norm2"						)
								lowercase__					:   Union[str, Any]													=					new_item.replace("out_layers.3"					,			"conv2"						)

								lowercase__					:   Dict													=					new_item.replace("emb_layers.1"					,			"time_emb_proj"						)
								lowercase__					:   int													=					new_item.replace("skip_connection"					,			"conv_shortcut"						)

								lowercase__					:   Tuple													=					shave_segments(lowerCamelCase__					,			n_shave_prefix_segments=lowerCamelCase__						)

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

				return mapping


def  __lowerCamelCase (       lowerCamelCase__					,			lowerCamelCase__=0						):

				"""simple docstring"""


				lowercase__					:   str													=					[]
				for old_item in old_list:
								lowercase__					:   Optional[int]													=					old_item

								lowercase__					:   Dict													=					new_item.replace("norm.weight"					,			"group_norm.weight"						)
								lowercase__					:   Optional[int]													=					new_item.replace("norm.bias"					,			"group_norm.bias"						)

								lowercase__					:   Tuple													=					new_item.replace("proj_out.weight"					,			"proj_attn.weight"						)
								lowercase__					:   List[Any]													=					new_item.replace("proj_out.bias"					,			"proj_attn.bias"						)

								lowercase__					:   Optional[Any]													=					shave_segments(lowerCamelCase__					,			n_shave_prefix_segments=lowerCamelCase__						)

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

				return mapping


def  __lowerCamelCase (       lowerCamelCase__					,			lowerCamelCase__					,			lowerCamelCase__					,			lowerCamelCase__=None					,			lowerCamelCase__=None					,			lowerCamelCase__=None						):

				"""simple docstring"""


				assert isinstance(lowerCamelCase__					,			lowerCamelCase__						), "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():
												lowercase__					:   List[str]													=					old_checkpoint[path]
												lowercase__					:   str													=					old_tensor.shape[0] // 3

												lowercase__					:   List[str]													=					(-1, channels) if len(old_tensor.shape						) == 3 else (-1)

												lowercase__					:   Union[str, Any]													=					old_tensor.shape[0] // config["num_head_channels"] // 3

												lowercase__					:   Union[str, Any]													=					old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]						)
												lowercase__							,							lowercase__							,							lowercase__					:   Optional[Any]													=					old_tensor.split(channels // num_heads					,			dim=1						)

												lowercase__					:   Dict													=					query.reshape(lowerCamelCase__						)
												lowercase__					:   Dict													=					key.reshape(lowerCamelCase__						)
												lowercase__					:   int													=					value.reshape(lowerCamelCase__						)

				for path in paths:
								lowercase__					:   Union[str, Any]													=					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
								lowercase__					:   List[Any]													=					new_path.replace("middle_block.0"					,			"mid_block.resnets.0"						)
								lowercase__					:   Optional[Any]													=					new_path.replace("middle_block.1"					,			"mid_block.attentions.0"						)
								lowercase__					:   List[str]													=					new_path.replace("middle_block.2"					,			"mid_block.resnets.1"						)

								if additional_replacements is not None:
												for replacement in additional_replacements:
																lowercase__					:   Tuple													=					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:
												lowercase__					:   List[Any]													=					old_checkpoint[path["old"]][:, :, 0]
								else:
												lowercase__					:   List[Any]													=					old_checkpoint[path["old"]]


def  __lowerCamelCase (       lowerCamelCase__					,			lowerCamelCase__						):

				"""simple docstring"""


				lowercase__					:   Union[str, Any]													=					{}

				lowercase__					:   Optional[Any]													=					checkpoint["time_embed.0.weight"]
				lowercase__					:   Tuple													=					checkpoint["time_embed.0.bias"]
				lowercase__					:   Dict													=					checkpoint["time_embed.2.weight"]
				lowercase__					:   Optional[Any]													=					checkpoint["time_embed.2.bias"]

				lowercase__					:   Optional[int]													=					checkpoint["input_blocks.0.0.weight"]
				lowercase__					:   List[Any]													=					checkpoint["input_blocks.0.0.bias"]

				lowercase__					:   Tuple													=					checkpoint["out.0.weight"]
				lowercase__					:   List[Any]													=					checkpoint["out.0.bias"]
				lowercase__					:   Tuple													=					checkpoint["out.2.weight"]
				lowercase__					:   Optional[Any]													=					checkpoint["out.2.bias"]

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

				# Retrieves the keys for the middle blocks only
				lowercase__					:   Tuple													=					len({".".join(layer.split("."						)[:2]						) for layer in checkpoint if "middle_block" in layer}						)
				lowercase__					:   Union[str, Any]													=					{
				    layer_id: [key for key in checkpoint if F"""middle_block.{layer_id}""" in key]
				    for layer_id in range(lowerCamelCase__						)
				}

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

				for i in range(1					,			lowerCamelCase__						):
								lowercase__					:   Tuple													=					(i - 1) // (config["num_res_blocks"] + 1)
								lowercase__					:   Optional[int]													=					(i - 1) % (config["num_res_blocks"] + 1)

								lowercase__					:   List[Any]													=					[key for key in input_blocks[i] if F"""input_blocks.{i}.0""" in key]
								lowercase__					:   Dict													=					[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:
												lowercase__					:   int													=					checkpoint[
												    F"""input_blocks.{i}.0.op.weight"""
												]
												lowercase__					:   List[str]													=					checkpoint[
												    F"""input_blocks.{i}.0.op.bias"""
												]
												continue

								lowercase__					:   Union[str, Any]													=					renew_resnet_paths(lowerCamelCase__						)
								lowercase__					:   Optional[int]													=					{"old": F"""input_blocks.{i}.0""", "new": F"""down_blocks.{block_id}.resnets.{layer_in_block_id}"""}
								lowercase__					:   Optional[int]													=					{"old": "resnets.2.op", "new": "downsamplers.0.op"}
								assign_to_checkpoint(
								    lowerCamelCase__					,			lowerCamelCase__					,			lowerCamelCase__					,			additional_replacements=[meta_path, resnet_op]					,			config=lowerCamelCase__						)

								if len(lowerCamelCase__						):
												lowercase__					:   Tuple													=					renew_attention_paths(lowerCamelCase__						)
												lowercase__					:   str													=					{
												    "old": F"""input_blocks.{i}.1""",
												    "new": F"""down_blocks.{block_id}.attentions.{layer_in_block_id}""",
												}
												lowercase__					:   List[str]													=					{
												    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(
												    lowerCamelCase__					,			lowerCamelCase__					,			lowerCamelCase__					,			additional_replacements=[meta_path]					,			attention_paths_to_split=lowerCamelCase__					,			config=lowerCamelCase__					,			)

				lowercase__					:   int													=					middle_blocks[0]
				lowercase__					:   Dict													=					middle_blocks[1]
				lowercase__					:   Dict													=					middle_blocks[2]

				lowercase__					:   Any													=					renew_resnet_paths(lowerCamelCase__						)
				assign_to_checkpoint(lowerCamelCase__					,			lowerCamelCase__					,			lowerCamelCase__					,			config=lowerCamelCase__						)

				lowercase__					:   List[Any]													=					renew_resnet_paths(lowerCamelCase__						)
				assign_to_checkpoint(lowerCamelCase__					,			lowerCamelCase__					,			lowerCamelCase__					,			config=lowerCamelCase__						)

				lowercase__					:   Optional[int]													=					renew_attention_paths(lowerCamelCase__						)
				lowercase__					:   Optional[int]													=					{
				    "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(
				    lowerCamelCase__					,			lowerCamelCase__					,			lowerCamelCase__					,			attention_paths_to_split=lowerCamelCase__					,			config=lowerCamelCase__						)

				for i in range(lowerCamelCase__						):
								lowercase__					:   List[Any]													=					i // (config["num_res_blocks"] + 1)
								lowercase__					:   Optional[int]													=					i % (config["num_res_blocks"] + 1)
								lowercase__					:   List[Any]													=					[shave_segments(lowerCamelCase__					,			2						) for name in output_blocks[i]]
								lowercase__					:   Optional[Any]													=					{}

								for layer in output_block_layers:
												lowercase__							,							lowercase__					:   str													=					layer.split("."						)[0], shave_segments(lowerCamelCase__					,			1						)
												if layer_id in output_block_list:
																output_block_list[layer_id].append(lowerCamelCase__						)
												else:
																lowercase__					:   Tuple													=					[layer_name]

								if len(lowerCamelCase__						) > 1:
												lowercase__					:   Dict													=					[key for key in output_blocks[i] if F"""output_blocks.{i}.0""" in key]
												lowercase__					:   Dict													=					[key for key in output_blocks[i] if F"""output_blocks.{i}.1""" in key]

												lowercase__					:   Optional[Any]													=					renew_resnet_paths(lowerCamelCase__						)
												lowercase__					:   Optional[Any]													=					renew_resnet_paths(lowerCamelCase__						)

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

												if ["conv.weight", "conv.bias"] in output_block_list.values():
																lowercase__					:   List[str]													=					list(output_block_list.values()						).index(["conv.weight", "conv.bias"]						)
																lowercase__					:   Tuple													=					checkpoint[
																    F"""output_blocks.{i}.{index}.conv.weight"""
																]
																lowercase__					:   Optional[Any]													=					checkpoint[
																    F"""output_blocks.{i}.{index}.conv.bias"""
																]

																# Clear attentions as they have been attributed above.
																if len(lowerCamelCase__						) == 2:
																				lowercase__					:   int													=					[]

												if len(lowerCamelCase__						):
																lowercase__					:   Tuple													=					renew_attention_paths(lowerCamelCase__						)
																lowercase__					:   str													=					{
																    "old": F"""output_blocks.{i}.1""",
																    "new": F"""up_blocks.{block_id}.attentions.{layer_in_block_id}""",
																}
																lowercase__					:   Union[str, Any]													=					{
																    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(
																    lowerCamelCase__					,			lowerCamelCase__					,			lowerCamelCase__					,			additional_replacements=[meta_path]					,			attention_paths_to_split=to_split if any("qkv" in key for key in attentions						) else None					,			config=lowerCamelCase__					,			)
								else:
												lowercase__					:   int													=					renew_resnet_paths(lowerCamelCase__					,			n_shave_prefix_segments=1						)
												for path in resnet_0_paths:
																lowercase__					:   List[Any]													=					".".join(["output_blocks", str(lowerCamelCase__						), path["old"]]						)
																lowercase__					:   Any													=					".".join(["up_blocks", str(lowerCamelCase__						), "resnets", str(lowerCamelCase__						), path["new"]]						)

																lowercase__					:   List[Any]													=					checkpoint[old_path]

				return new_checkpoint


if __name__ == "__main__":
				lowerCAmelCase__					  =  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.''')

				lowerCAmelCase__					  =  parser.parse_args()

				lowerCAmelCase__					  =  torch.load(args.checkpoint_path)

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

				lowerCAmelCase__					  =  convert_ldm_checkpoint(checkpoint, config)

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

				lowerCAmelCase__					  =  UNetaDModel(**config)
				model.load_state_dict(converted_checkpoint)

				try:
								lowerCAmelCase__					  =  DDPMScheduler.from_config('''/'''.join(args.checkpoint_path.split('''/''')[:-1]))
								lowerCAmelCase__					  =  VQModel.from_pretrained('''/'''.join(args.checkpoint_path.split('''/''')[:-1]))

								lowerCAmelCase__					  =  LDMPipeline(unet=model, scheduler=scheduler, vae=vqvae)
								pipe.save_pretrained(args.dump_path)
				except:  # noqa: E722
								model.save_pretrained(args.dump_path)




