





import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




@skip_mps
class 	_UpperCAmelCase						( _lowercase			,							_lowercase			,							_lowercase			,							unittest.TestCase       ):
					__A        = StableDiffusionModelEditingPipeline
					__A        = TEXT_TO_IMAGE_PARAMS
					__A        = TEXT_TO_IMAGE_BATCH_PARAMS
					__A        = TEXT_TO_IMAGE_IMAGE_PARAMS
					__A        = TEXT_TO_IMAGE_IMAGE_PARAMS




					def 						snake_case_		(    self							)    ->	int:
							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     ,   )
							UpperCamelCase__								=						DDIMScheduler()
							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     ,   )
							torch.manual_seed(0							)
							UpperCamelCase__								=						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     ,   )
							UpperCamelCase__								=						CLIPTextModel(__lowerCAmelCase							)
							UpperCamelCase__								=						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""							)

							UpperCamelCase__								=						{
							    """unet""": unet,
							    """scheduler""": scheduler,
							    """vae""": vae,
							    """text_encoder""": text_encoder,
							    """tokenizer""": tokenizer,
							    """safety_checker""": None,
							    """feature_extractor""": None,
							}
							return components




					def 						snake_case_		(    self     ,   __lowerCAmelCase     ,   __lowerCAmelCase=0							)    ->	Tuple:
							UpperCamelCase__								=						torch.manual_seed(__lowerCAmelCase							)
							UpperCamelCase__								=						{
							    """prompt""": """A field of roses""",
							    """generator""": generator,
							    # Setting height and width to None to prevent OOMs on CPU.
							    """height""": None,
							    """width""": None,
							    """num_inference_steps""": 2,
							    """guidance_scale""": 6.0,
							    """output_type""": """numpy""",
							}
							return inputs




					def 						snake_case_		(    self							)    ->	Any:
							UpperCamelCase__								=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
							UpperCamelCase__								=						self.get_dummy_components()
							UpperCamelCase__								=						StableDiffusionModelEditingPipeline(**__lowerCAmelCase							)
							UpperCamelCase__								=						sd_pipe.to(__lowerCAmelCase							)
							sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase							)

							UpperCamelCase__								=						self.get_dummy_inputs(__lowerCAmelCase							)
							UpperCamelCase__								=						sd_pipe(**__lowerCAmelCase							).images
							UpperCamelCase__								=						image[0, -3:, -3:, -1]
							assert image.shape == (1, 64, 64, 3)

							UpperCamelCase__								=						np.array([0.4_7_5_5, 0.5_1_3_2, 0.4_9_7_6, 0.3_9_0_4, 0.3_5_5_4, 0.4_7_6_5, 0.5_1_3_9, 0.5_1_5_8, 0.4_8_8_9]							)

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




					def 						snake_case_		(    self							)    ->	Any:
							UpperCamelCase__								=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
							UpperCamelCase__								=						self.get_dummy_components()
							UpperCamelCase__								=						StableDiffusionModelEditingPipeline(**__lowerCAmelCase							)
							UpperCamelCase__								=						sd_pipe.to(__lowerCAmelCase							)
							sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase							)

							UpperCamelCase__								=						self.get_dummy_inputs(__lowerCAmelCase							)
							UpperCamelCase__								=						"""french fries"""
							UpperCamelCase__								=						sd_pipe(**__lowerCAmelCase     ,   negative_prompt=__lowerCAmelCase							)
							UpperCamelCase__								=						output.images
							UpperCamelCase__								=						image[0, -3:, -3:, -1]

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

							UpperCamelCase__								=						np.array([0.4_9_9_2, 0.5_1_0_1, 0.5_0_0_4, 0.3_9_4_9, 0.3_6_0_4, 0.4_7_3_5, 0.5_2_1_6, 0.5_2_0_4, 0.4_9_1_3]							)

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




					def 						snake_case_		(    self							)    ->	Tuple:
							UpperCamelCase__								=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
							UpperCamelCase__								=						self.get_dummy_components()
							UpperCamelCase__								=						EulerAncestralDiscreteScheduler(
							    beta_start=0.0_0_0_8_5     ,   beta_end=0.0_1_2     ,   beta_schedule="""scaled_linear"""							)
							UpperCamelCase__								=						StableDiffusionModelEditingPipeline(**__lowerCAmelCase							)
							UpperCamelCase__								=						sd_pipe.to(__lowerCAmelCase							)
							sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase							)

							UpperCamelCase__								=						self.get_dummy_inputs(__lowerCAmelCase							)
							UpperCamelCase__								=						sd_pipe(**__lowerCAmelCase							).images
							UpperCamelCase__								=						image[0, -3:, -3:, -1]

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

							UpperCamelCase__								=						np.array([0.4_7_4_7, 0.5_3_7_2, 0.4_7_7_9, 0.4_9_8_2, 0.5_5_4_3, 0.4_8_1_6, 0.5_2_3_8, 0.4_9_0_4, 0.5_0_2_7]							)

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




					def 						snake_case_		(    self							)    ->	Tuple:
							UpperCamelCase__								=						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
							UpperCamelCase__								=						self.get_dummy_components()
							UpperCamelCase__								=						PNDMScheduler()
							UpperCamelCase__								=						StableDiffusionModelEditingPipeline(**__lowerCAmelCase							)
							UpperCamelCase__								=						sd_pipe.to(__lowerCAmelCase							)
							sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase							)

							UpperCamelCase__								=						self.get_dummy_inputs(__lowerCAmelCase							)
							# the pipeline does not expect pndm so test if it raises error.
							with self.assertRaises(__lowerCAmelCase							):
									UpperCamelCase__								=						sd_pipe(**__lowerCAmelCase							).images




					def 						snake_case_		(    self							)    ->	List[Any]:
							super().test_inference_batch_single_identical(expected_max_diff=5E-3							)




					def 						snake_case_		(    self							)    ->	List[Any]:
							super().test_attention_slicing_forward_pass(expected_max_diff=5E-3							)




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




					def 						snake_case_		(    self							)    ->	Union[str, Any]:
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()




					def 						snake_case_		(    self     ,   __lowerCAmelCase=0							)    ->	str:
							UpperCamelCase__								=						torch.manual_seed(__lowerCAmelCase							)
							UpperCamelCase__								=						{
							    """prompt""": """A field of roses""",
							    """generator""": generator,
							    """num_inference_steps""": 3,
							    """guidance_scale""": 7.5,
							    """output_type""": """numpy""",
							}
							return inputs




					def 						snake_case_		(    self							)    ->	Tuple:
							UpperCamelCase__								=						"""CompVis/stable-diffusion-v1-4"""
							UpperCamelCase__								=						StableDiffusionModelEditingPipeline.from_pretrained(__lowerCAmelCase     ,   safety_checker=__lowerCAmelCase							)
							pipe.to(__lowerCAmelCase							)
							pipe.set_progress_bar_config(disable=__lowerCAmelCase							)
							pipe.enable_attention_slicing()

							UpperCamelCase__								=						self.get_inputs()
							UpperCamelCase__								=						pipe(**__lowerCAmelCase							).images
							UpperCamelCase__								=						image[0, -3:, -3:, -1].flatten()

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

							UpperCamelCase__								=						np.array(
							    [0.6_7_4_9_4_9_6, 0.6_3_8_6_4_5_3, 0.5_1_4_4_3_2_6_7, 0.6_6_0_9_4_9_0_5, 0.6_1_9_2_1_2_1_5, 0.5_4_9_1_3_3_2, 0.5_7_4_4_4_1_7, 0.5_8_0_7_5_1_0_6, 0.5_1_7_4_6_5_8]							)

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

							# make sure image changes after editing
							pipe.edit_model("""A pack of roses"""     ,   """A pack of blue roses"""							)

							UpperCamelCase__								=						pipe(**__lowerCAmelCase							).images
							UpperCamelCase__								=						image[0, -3:, -3:, -1].flatten()

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

							assert np.abs(expected_slice - image_slice							).max() > 1E-1




					def 						snake_case_		(    self							)    ->	str:
							torch.cuda.empty_cache()
							torch.cuda.reset_max_memory_allocated()
							torch.cuda.reset_peak_memory_stats()

							UpperCamelCase__								=						"""CompVis/stable-diffusion-v1-4"""
							UpperCamelCase__								=						DDIMScheduler.from_pretrained(__lowerCAmelCase     ,   subfolder="""scheduler"""							)
							UpperCamelCase__								=						StableDiffusionModelEditingPipeline.from_pretrained(
							    __lowerCAmelCase     ,   scheduler=__lowerCAmelCase     ,   safety_checker=__lowerCAmelCase							)
							UpperCamelCase__								=						pipe.to(__lowerCAmelCase							)
							pipe.set_progress_bar_config(disable=__lowerCAmelCase							)
							pipe.enable_attention_slicing(1							)
							pipe.enable_sequential_cpu_offload()

							UpperCamelCase__								=						self.get_inputs()
							UpperCamelCase__								=						pipe(**__lowerCAmelCase							)

							UpperCamelCase__								=						torch.cuda.max_memory_allocated()
							# make sure that less than 4.4 GB is allocated
							assert mem_bytes < 4.4 * 10**9












import argparse

import torch

from diffusers import UNetaDConditionModel

def      a    (   _lowercase :  Dict			,    _lowercase :  Dict			,    _lowercase :  Any			,    _lowercase :  Tuple=None			,    _lowercase :  List[Any]=None			,    _lowercase :  Optional[Any]=None							)  ->Any:
		assert isinstance(_lowercase			,    _lowercase							), "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__	,      UpperCamelCase__	,      UpperCamelCase__								=						old_tensor.split(channels // num_heads			,    dim=1							)

						UpperCamelCase__								=						query.reshape(_lowercase							)
						UpperCamelCase__								=						key.reshape(_lowercase							)
						UpperCamelCase__								=						value.reshape(_lowercase							)

		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

				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
				UpperCamelCase__								=						old_checkpoint[path["""old"""]]
				UpperCamelCase__								=						["""proj_attn.weight"""]
				UpperCamelCase__								=						["""proj_out.weight""", """proj_in.weight"""]
				if any(k in new_path for k in names							):
						UpperCamelCase__								=						weight[:, :, 0]
				elif any(k in new_path for k in names_a							) and len(weight.shape							) > 2 and ".attentions." not in new_path:
						UpperCamelCase__								=						weight[:, :, 0]
				else:
						UpperCamelCase__								=						weight

def      a    (   _lowercase :  Tuple			,    _lowercase :  List[Any]=0							)  ->int:
		UpperCamelCase__								=						[]
		for old_item in old_list:
				UpperCamelCase__								=						old_item

				#         new_item = new_item.replace('norm.weight', 'group_norm.weight')
				#         new_item = new_item.replace('norm.bias', 'group_norm.bias')

				#         new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')
				#         new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')

				#         new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)

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

		return mapping

def      a    (   _lowercase :  Union[str, Any]			,    _lowercase :  int=1							)  ->str:
		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 :  List[str]			,    _lowercase :  Any=0							)  ->Any:
		UpperCamelCase__								=						[]
		for old_item in old_list:
				mapping.append({"""old""": old_item, """new""": old_item}							)

		return mapping

def      a    (   _lowercase :  List[str]			,    _lowercase :  int=0							)  ->Tuple:
		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(_lowercase			,    n_shave_prefix_segments=_lowercase							)

				if "temopral_conv" not in old_item:
						mapping.append({"""old""": old_item, """new""": new_item}							)

		return mapping

def      a    (   _lowercase :  List[Any]			,    _lowercase :  List[Any]			,    _lowercase :  int=None			,    _lowercase :  Optional[Any]=False							)  ->List[str]:
		UpperCamelCase__								=						{}
		UpperCamelCase__								=						list(checkpoint.keys()							)

		UpperCamelCase__								=						"""model.diffusion_model."""

		# at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA
		if sum(k.startswith("""model_ema"""							) for k in keys							) > 1_0_0 and extract_ema:
				print(F"""Checkpoint {path} has both EMA and non-EMA weights."""							)
				print(
				    """In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA"""
				    """ weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag."""							)
				for key in keys:
						if key.startswith("""model.diffusion_model"""							):
								UpperCamelCase__								=						"""model_ema.""" + """""".join(key.split("""."""							)[1:]							)
								UpperCamelCase__								=						checkpoint.pop(_lowercase							)
		else:
				if sum(k.startswith("""model_ema"""							) for k in keys							) > 1_0_0:
						print(
						    """In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA"""
						    """ weights (usually better for inference), please make sure to add the `--extract_ema` flag."""							)

				for key in keys:
						UpperCamelCase__								=						checkpoint.pop(_lowercase							)

		UpperCamelCase__								=						{}

		UpperCamelCase__								=						unet_state_dict["""time_embed.0.weight"""]
		UpperCamelCase__								=						unet_state_dict["""time_embed.0.bias"""]
		UpperCamelCase__								=						unet_state_dict["""time_embed.2.weight"""]
		UpperCamelCase__								=						unet_state_dict["""time_embed.2.bias"""]

		if config["class_embed_type"] is None:
				# No parameters to port
				...
		elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection":
				UpperCamelCase__								=						unet_state_dict["""label_emb.0.0.weight"""]
				UpperCamelCase__								=						unet_state_dict["""label_emb.0.0.bias"""]
				UpperCamelCase__								=						unet_state_dict["""label_emb.0.2.weight"""]
				UpperCamelCase__								=						unet_state_dict["""label_emb.0.2.bias"""]
		else:
				raise NotImplementedError(F"""Not implemented `class_embed_type`: {config['class_embed_type']}"""							)

		UpperCamelCase__								=						unet_state_dict["""input_blocks.0.0.weight"""]
		UpperCamelCase__								=						unet_state_dict["""input_blocks.0.0.bias"""]

		UpperCamelCase__								=						[v for v in unet_state_dict if v.startswith("""input_blocks.0.1"""							)]
		UpperCamelCase__								=						renew_attention_paths(_lowercase							)
		UpperCamelCase__								=						{"""old""": """input_blocks.0.1""", """new""": """transformer_in"""}
		assign_to_checkpoint(_lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

		UpperCamelCase__								=						unet_state_dict["""out.0.weight"""]
		UpperCamelCase__								=						unet_state_dict["""out.0.bias"""]
		UpperCamelCase__								=						unet_state_dict["""out.2.weight"""]
		UpperCamelCase__								=						unet_state_dict["""out.2.bias"""]

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

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

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

		for i in range(1			,    _lowercase							):
				UpperCamelCase__								=						(i - 1) // (config["""layers_per_block"""] + 1)
				UpperCamelCase__								=						(i - 1) % (config["""layers_per_block"""] + 1)

				UpperCamelCase__								=						[
				    key for key in input_blocks[i] if F"""input_blocks.{i}.0""" in key and F"""input_blocks.{i}.0.op""" not in key
				]
				UpperCamelCase__								=						[key for key in input_blocks[i] if F"""input_blocks.{i}.1""" in key]
				UpperCamelCase__								=						[key for key in input_blocks[i] if F"""input_blocks.{i}.2""" in key]

				if F"""input_blocks.{i}.op.weight""" in unet_state_dict:
						UpperCamelCase__								=						unet_state_dict.pop(
						    F"""input_blocks.{i}.op.weight"""							)
						UpperCamelCase__								=						unet_state_dict.pop(
						    F"""input_blocks.{i}.op.bias"""							)

				UpperCamelCase__								=						renew_resnet_paths(_lowercase							)
				UpperCamelCase__								=						{"""old""": F"""input_blocks.{i}.0""", """new""": F"""down_blocks.{block_id}.resnets.{layer_in_block_id}"""}
				assign_to_checkpoint(
				    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

				UpperCamelCase__								=						[key for key in resnets if """temopral_conv""" in key]
				UpperCamelCase__								=						renew_temp_conv_paths(_lowercase							)
				UpperCamelCase__								=						{
				    """old""": F"""input_blocks.{i}.0.temopral_conv""",
				    """new""": F"""down_blocks.{block_id}.temp_convs.{layer_in_block_id}""",
				}
				assign_to_checkpoint(
				    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

				if len(_lowercase							):
						UpperCamelCase__								=						renew_attention_paths(_lowercase							)
						UpperCamelCase__								=						{"""old""": F"""input_blocks.{i}.1""", """new""": F"""down_blocks.{block_id}.attentions.{layer_in_block_id}"""}
						assign_to_checkpoint(
						    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

				if len(_lowercase							):
						UpperCamelCase__								=						renew_attention_paths(_lowercase							)
						UpperCamelCase__								=						{
						    """old""": F"""input_blocks.{i}.2""",
						    """new""": F"""down_blocks.{block_id}.temp_attentions.{layer_in_block_id}""",
						}
						assign_to_checkpoint(
						    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

		UpperCamelCase__								=						middle_blocks[0]
		UpperCamelCase__								=						[key for key in resnet_a if """temopral_conv""" in key]
		UpperCamelCase__								=						middle_blocks[1]
		UpperCamelCase__								=						middle_blocks[2]
		UpperCamelCase__								=						middle_blocks[3]
		UpperCamelCase__								=						[key for key in resnet_a if """temopral_conv""" in key]

		UpperCamelCase__								=						renew_resnet_paths(_lowercase							)
		UpperCamelCase__								=						{"""old""": """middle_block.0""", """new""": """mid_block.resnets.0"""}
		assign_to_checkpoint(
		    _lowercase			,    _lowercase			,    _lowercase			,    config=_lowercase			,    additional_replacements=[meta_path]							)

		UpperCamelCase__								=						renew_temp_conv_paths(_lowercase							)
		UpperCamelCase__								=						{"""old""": """middle_block.0.temopral_conv""", """new""": """mid_block.temp_convs.0"""}
		assign_to_checkpoint(
		    _lowercase			,    _lowercase			,    _lowercase			,    config=_lowercase			,    additional_replacements=[meta_path]							)

		UpperCamelCase__								=						renew_resnet_paths(_lowercase							)
		UpperCamelCase__								=						{"""old""": """middle_block.3""", """new""": """mid_block.resnets.1"""}
		assign_to_checkpoint(
		    _lowercase			,    _lowercase			,    _lowercase			,    config=_lowercase			,    additional_replacements=[meta_path]							)

		UpperCamelCase__								=						renew_temp_conv_paths(_lowercase							)
		UpperCamelCase__								=						{"""old""": """middle_block.3.temopral_conv""", """new""": """mid_block.temp_convs.1"""}
		assign_to_checkpoint(
		    _lowercase			,    _lowercase			,    _lowercase			,    config=_lowercase			,    additional_replacements=[meta_path]							)

		UpperCamelCase__								=						renew_attention_paths(_lowercase							)
		UpperCamelCase__								=						{"""old""": """middle_block.1""", """new""": """mid_block.attentions.0"""}
		assign_to_checkpoint(
		    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

		UpperCamelCase__								=						renew_attention_paths(_lowercase							)
		UpperCamelCase__								=						{"""old""": """middle_block.2""", """new""": """mid_block.temp_attentions.0"""}
		assign_to_checkpoint(
		    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

		for i in range(_lowercase							):
				UpperCamelCase__								=						i // (config["""layers_per_block"""] + 1)
				UpperCamelCase__								=						i % (config["""layers_per_block"""] + 1)
				UpperCamelCase__								=						[shave_segments(_lowercase			,    2							) for name in output_blocks[i]]
				UpperCamelCase__								=						{}

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

				if len(_lowercase							) > 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__								=						[key for key in output_blocks[i] if F"""output_blocks.{i}.2""" in key]

						UpperCamelCase__								=						renew_resnet_paths(_lowercase							)
						UpperCamelCase__								=						renew_resnet_paths(_lowercase							)

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

						UpperCamelCase__								=						[key for key in resnets if """temopral_conv""" in key]
						UpperCamelCase__								=						renew_temp_conv_paths(_lowercase							)
						UpperCamelCase__								=						{
						    """old""": F"""output_blocks.{i}.0.temopral_conv""",
						    """new""": F"""up_blocks.{block_id}.temp_convs.{layer_in_block_id}""",
						}
						assign_to_checkpoint(
						    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

						UpperCamelCase__								=						{k: sorted(_lowercase							) for k, v in output_block_list.items()}
						if ["conv.bias", "conv.weight"] in output_block_list.values():
								UpperCamelCase__								=						list(output_block_list.values()							).index(["""conv.bias""", """conv.weight"""]							)
								UpperCamelCase__								=						unet_state_dict[
								    F"""output_blocks.{i}.{index}.conv.weight"""
								]
								UpperCamelCase__								=						unet_state_dict[
								    F"""output_blocks.{i}.{index}.conv.bias"""
								]

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

						if len(_lowercase							):
								UpperCamelCase__								=						renew_attention_paths(_lowercase							)
								UpperCamelCase__								=						{
								    """old""": F"""output_blocks.{i}.1""",
								    """new""": F"""up_blocks.{block_id}.attentions.{layer_in_block_id}""",
								}
								assign_to_checkpoint(
								    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)

						if len(_lowercase							):
								UpperCamelCase__								=						renew_attention_paths(_lowercase							)
								UpperCamelCase__								=						{
								    """old""": F"""output_blocks.{i}.2""",
								    """new""": F"""up_blocks.{block_id}.temp_attentions.{layer_in_block_id}""",
								}
								assign_to_checkpoint(
								    _lowercase			,    _lowercase			,    _lowercase			,    additional_replacements=[meta_path]			,    config=_lowercase							)
				else:
						UpperCamelCase__								=						renew_resnet_paths(_lowercase			,    n_shave_prefix_segments=1							)
						for path in resnet_0_paths:
								UpperCamelCase__								=						""".""".join(["""output_blocks""", str(_lowercase							), path["""old"""]]							)
								UpperCamelCase__								=						""".""".join(["""up_blocks""", str(_lowercase							), """resnets""", str(_lowercase							), path["""new"""]]							)
								UpperCamelCase__								=						unet_state_dict[old_path]

						UpperCamelCase__								=						[l for l in output_block_layers if """temopral_conv""" in l]
						for path in temopral_conv_paths:
								UpperCamelCase__								=						path.split("""temopral_conv."""							)[-1]
								UpperCamelCase__								=						""".""".join(["""output_blocks""", str(_lowercase							), str(_lowercase							), """temopral_conv""", pruned_path]							)
								UpperCamelCase__								=						""".""".join(["""up_blocks""", str(_lowercase							), """temp_convs""", str(_lowercase							), pruned_path]							)
								UpperCamelCase__								=						unet_state_dict[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("--dump_path", default=None, type=str, required=True, help="Path to the output model.")
				__lowerCAmelCase 	=      parser.parse_args()

				__lowerCAmelCase 	=      torch.load(args.checkpoint_path, map_location="cpu")
				__lowerCAmelCase 	=      UNetaDConditionModel()

				__lowerCAmelCase 	=      convert_ldm_unet_checkpoint(unet_checkpoint, unet.config)

				__lowerCAmelCase 	=      set(unet.state_dict().keys()) - set(converted_ckpt.keys())
				__lowerCAmelCase 	=      set(converted_ckpt.keys()) - set(unet.state_dict().keys())

				assert len(diff_a) == len(diff_a) == 0, "Converted weights don't match"

				# load state_dict
				unet.load_state_dict(converted_ckpt)

				unet.save_pretrained(args.dump_path)

				# -- finish converting the unet --





