




from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block
@dataclass
class 		UpperCAmelCase   (				UpperCAmelCase__			):







						'''simple docstring'''





						lowerCAmelCase_     		=   42







class 		UpperCAmelCase   (				nn.Module			):







						'''simple docstring'''





						def __init__(						self							:     Dict       ,	__lowercase							:     str=3       ,	__lowercase							:     Optional[Any]=3       ,	__lowercase							:     List[Any]=("DownEncoderBlock2D",)       ,	__lowercase							:     List[str]=(64,)       ,	__lowercase							:     Any=2       ,	__lowercase							:     Dict=32       ,	__lowercase							:     Union[str, Any]="silu"       ,	__lowercase							:     Any=True       ,	):






							"""simple docstring"""






							super().__init__()
							snake_case_		    =     layers_per_block

							snake_case_		    =     torch.nn.Convad(
							    __lowercase       ,	block_out_channels[0]       ,	kernel_size=3       ,	stride=1       ,	padding=1       ,	)

							snake_case_		    =     None
							snake_case_		    =     nn.ModuleList([]						)

							# down
							snake_case_		    =     block_out_channels[0]
							for i, down_block_type in enumerate(__lowercase						):
								snake_case_		    =     output_channel
								snake_case_		    =     block_out_channels[i]
								snake_case_		    =     i == len(__lowercase						) - 1

								snake_case_		    =     get_down_block(
								    __lowercase       ,	num_layers=self.layers_per_block       ,	in_channels=__lowercase       ,	out_channels=__lowercase       ,	add_downsample=not is_final_block       ,	resnet_eps=1E-6       ,	downsample_padding=0       ,	resnet_act_fn=__lowercase       ,	resnet_groups=__lowercase       ,	attention_head_dim=__lowercase       ,	temb_channels=__lowercase       ,	)
								self.down_blocks.append(__lowercase						)

							# mid
							snake_case_		    =     UNetMidBlockaD(
							    in_channels=block_out_channels[-1]       ,	resnet_eps=1E-6       ,	resnet_act_fn=__lowercase       ,	output_scale_factor=1       ,	resnet_time_scale_shift="default"       ,	attention_head_dim=block_out_channels[-1]       ,	resnet_groups=__lowercase       ,	temb_channels=__lowercase       ,	)

							# out
							snake_case_		    =     nn.GroupNorm(num_channels=block_out_channels[-1]       ,	num_groups=__lowercase       ,	eps=1E-6						)
							snake_case_		    =     nn.SiLU()

							snake_case_		    =     2 * out_channels if double_z else out_channels
							snake_case_		    =     nn.Convad(block_out_channels[-1]       ,	__lowercase       ,	3       ,	padding=1						)

							snake_case_		    =     False






						def        snake_case__						(						self							:     List[str]       ,	__lowercase							:     Dict						):






							"""simple docstring"""






							snake_case_		    =     x
							snake_case_		    =     self.conv_in(__lowercase						)

							if self.training and self.gradient_checkpointing:

								def create_custom_forward(__lowercase							:     Any						):
									def custom_forward(*__lowercase							:     Union[str, Any]						):
										return module(*__lowercase						)

									return custom_forward

								# down
								if is_torch_version(">="       ,	"1.11.0"						):
									for down_block in self.down_blocks:
										snake_case_		    =     torch.utils.checkpoint.checkpoint(
										    create_custom_forward(__lowercase						)       ,	__lowercase       ,	use_reentrant=__lowercase						)
									# middle
									snake_case_		    =     torch.utils.checkpoint.checkpoint(
									    create_custom_forward(self.mid_block						)       ,	__lowercase       ,	use_reentrant=__lowercase						)
								else:
									for down_block in self.down_blocks:
										snake_case_		    =     torch.utils.checkpoint.checkpoint(create_custom_forward(__lowercase						)       ,	__lowercase						)
									# middle
									snake_case_		    =     torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block						)       ,	__lowercase						)

							else:
								# down
								for down_block in self.down_blocks:
									snake_case_		    =     down_block(__lowercase						)

								# middle
								snake_case_		    =     self.mid_block(__lowercase						)

							# post-process
							snake_case_		    =     self.conv_norm_out(__lowercase						)
							snake_case_		    =     self.conv_act(__lowercase						)
							snake_case_		    =     self.conv_out(__lowercase						)

							return sample







class 		UpperCAmelCase   (				nn.Module			):







						'''simple docstring'''





						def __init__(						self							:     Optional[int]       ,	__lowercase							:     Dict=3       ,	__lowercase							:     List[str]=3       ,	__lowercase							:     Optional[int]=("UpDecoderBlock2D",)       ,	__lowercase							:     List[str]=(64,)       ,	__lowercase							:     Union[str, Any]=2       ,	__lowercase							:     Union[str, Any]=32       ,	__lowercase							:     List[Any]="silu"       ,	__lowercase							:     Tuple="group"       ,	):






							"""simple docstring"""






							super().__init__()
							snake_case_		    =     layers_per_block

							snake_case_		    =     nn.Convad(
							    __lowercase       ,	block_out_channels[-1]       ,	kernel_size=3       ,	stride=1       ,	padding=1       ,	)

							snake_case_		    =     None
							snake_case_		    =     nn.ModuleList([]						)

							snake_case_		    =     in_channels if norm_type == "spatial" else None

							# mid
							snake_case_		    =     UNetMidBlockaD(
							    in_channels=block_out_channels[-1]       ,	resnet_eps=1E-6       ,	resnet_act_fn=__lowercase       ,	output_scale_factor=1       ,	resnet_time_scale_shift="default" if norm_type == "group" else norm_type       ,	attention_head_dim=block_out_channels[-1]       ,	resnet_groups=__lowercase       ,	temb_channels=__lowercase       ,	)

							# up
							snake_case_		    =     list(reversed(__lowercase						)						)
							snake_case_		    =     reversed_block_out_channels[0]
							for i, up_block_type in enumerate(__lowercase						):
								snake_case_		    =     output_channel
								snake_case_		    =     reversed_block_out_channels[i]

								snake_case_		    =     i == len(__lowercase						) - 1

								snake_case_		    =     get_up_block(
								    __lowercase       ,	num_layers=self.layers_per_block + 1       ,	in_channels=__lowercase       ,	out_channels=__lowercase       ,	prev_output_channel=__lowercase       ,	add_upsample=not is_final_block       ,	resnet_eps=1E-6       ,	resnet_act_fn=__lowercase       ,	resnet_groups=__lowercase       ,	attention_head_dim=__lowercase       ,	temb_channels=__lowercase       ,	resnet_time_scale_shift=__lowercase       ,	)
								self.up_blocks.append(__lowercase						)
								snake_case_		    =     output_channel

							# out
							if norm_type == "spatial":
								snake_case_		    =     SpatialNorm(block_out_channels[0]       ,	__lowercase						)
							else:
								snake_case_		    =     nn.GroupNorm(num_channels=block_out_channels[0]       ,	num_groups=__lowercase       ,	eps=1E-6						)
							snake_case_		    =     nn.SiLU()
							snake_case_		    =     nn.Convad(block_out_channels[0]       ,	__lowercase       ,	3       ,	padding=1						)

							snake_case_		    =     False






						def        snake_case__						(						self							:     List[Any]       ,	__lowercase							:     Any       ,	__lowercase							:     List[Any]=None						):






							"""simple docstring"""






							snake_case_		    =     z
							snake_case_		    =     self.conv_in(__lowercase						)

							snake_case_		    =     next(iter(self.up_blocks.parameters()						)						).dtype
							if self.training and self.gradient_checkpointing:

								def create_custom_forward(__lowercase							:     str						):
									def custom_forward(*__lowercase							:     List[Any]						):
										return module(*__lowercase						)

									return custom_forward

								if is_torch_version(">="       ,	"1.11.0"						):
									# middle
									snake_case_		    =     torch.utils.checkpoint.checkpoint(
									    create_custom_forward(self.mid_block						)       ,	__lowercase       ,	__lowercase       ,	use_reentrant=__lowercase						)
									snake_case_		    =     sample.to(__lowercase						)

									# up
									for up_block in self.up_blocks:
										snake_case_		    =     torch.utils.checkpoint.checkpoint(
										    create_custom_forward(__lowercase						)       ,	__lowercase       ,	__lowercase       ,	use_reentrant=__lowercase						)
								else:
									# middle
									snake_case_		    =     torch.utils.checkpoint.checkpoint(
									    create_custom_forward(self.mid_block						)       ,	__lowercase       ,	__lowercase						)
									snake_case_		    =     sample.to(__lowercase						)

									# up
									for up_block in self.up_blocks:
										snake_case_		    =     torch.utils.checkpoint.checkpoint(create_custom_forward(__lowercase						)       ,	__lowercase       ,	__lowercase						)
							else:
								# middle
								snake_case_		    =     self.mid_block(__lowercase       ,	__lowercase						)
								snake_case_		    =     sample.to(__lowercase						)

								# up
								for up_block in self.up_blocks:
									snake_case_		    =     up_block(__lowercase       ,	__lowercase						)

        # post-process
							if latent_embeds is None:
								snake_case_		    =     self.conv_norm_out(__lowercase						)
							else:
								snake_case_		    =     self.conv_norm_out(__lowercase       ,	__lowercase						)
							snake_case_		    =     self.conv_act(__lowercase						)
							snake_case_		    =     self.conv_out(__lowercase						)

							return sample







class 		UpperCAmelCase   (				nn.Module			):







						'''simple docstring'''





						def __init__(						self							:     Tuple       ,	__lowercase							:     Tuple       ,	__lowercase							:     int       ,	__lowercase							:     Tuple       ,	__lowercase							:     List[Any]=None       ,	__lowercase							:     Tuple="random"       ,	__lowercase							:     Tuple=False       ,	__lowercase							:     List[str]=True						):






							"""simple docstring"""






							super().__init__()
							snake_case_		    =     n_e
							snake_case_		    =     vq_embed_dim
							snake_case_		    =     beta
							snake_case_		    =     legacy

							snake_case_		    =     nn.Embedding(self.n_e       ,	self.vq_embed_dim						)
							self.embedding.weight.data.uniform_(-1.0 / self.n_e       ,	1.0 / self.n_e						)

							snake_case_		    =     remap
							if self.remap is not None:
								self.register_buffer("used"       ,	torch.tensor(np.load(self.remap						)						)						)
								snake_case_		    =     self.used.shape[0]
								snake_case_		    =     unknown_index  # "random" or "extra" or integer
								if self.unknown_index == "extra":
									snake_case_		    =     self.re_embed
									snake_case_		    =     self.re_embed + 1
								print(
								    f"Remapping {self.n_e} indices to {self.re_embed} indices. "
								    f"Using {self.unknown_index} for unknown indices."						)
							else:
								snake_case_		    =     n_e

							snake_case_		    =     sane_index_shape






						def        snake_case__						(						self							:     Tuple       ,	__lowercase							:     Any						):






							"""simple docstring"""






							snake_case_		    =     inds.shape
							assert len(__lowercase						) > 1
							snake_case_		    =     inds.reshape(ishape[0]       ,	-1						)
							snake_case_		    =     self.used.to(__lowercase						)
							snake_case_		    =     (inds[:, :, None] == used[None, None, ...]).long()
							snake_case_		    =     match.argmax(-1						)
							snake_case_		    =     match.sum(2						) < 1
							if self.unknown_index == "random":
								snake_case_		    =     torch.randint(0       ,	self.re_embed       ,	size=new[unknown].shape						).to(device=new.device						)
							else:
								snake_case_		    =     self.unknown_index
							return new.reshape(__lowercase						)






						def        snake_case__						(						self							:     Optional[Any]       ,	__lowercase							:     Any						):






							"""simple docstring"""






							snake_case_		    =     inds.shape
							assert len(__lowercase						) > 1
							snake_case_		    =     inds.reshape(ishape[0]       ,	-1						)
							snake_case_		    =     self.used.to(__lowercase						)
							if self.re_embed > self.used.shape[0]:  # extra token
								snake_case_		    =     0  # simply set to zero
							snake_case_		    =     torch.gather(used[None, :][inds.shape[0] * [0], :]       ,	1       ,	__lowercase						)
							return back.reshape(__lowercase						)






						def        snake_case__						(						self							:     List[str]       ,	__lowercase							:     Dict						):






							"""simple docstring"""






							snake_case_		    =     z.permute(0       ,	2       ,	3       ,	1						).contiguous()
							snake_case_		    =     z.view(-1       ,	self.vq_embed_dim						)

							# distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
							snake_case_		    =     torch.argmin(torch.cdist(__lowercase       ,	self.embedding.weight						)       ,	dim=1						)

							snake_case_		    =     self.embedding(__lowercase						).view(z.shape						)
							snake_case_		    =     None
							snake_case_		    =     None

							# compute loss for embedding
							if not self.legacy:
								snake_case_		    =     self.beta * torch.mean((z_q.detach() - z) ** 2						) + torch.mean((z_q - z.detach()) ** 2						)
							else:
								snake_case_		    =     torch.mean((z_q.detach() - z) ** 2						) + self.beta * torch.mean((z_q - z.detach()) ** 2						)

							# preserve gradients
							snake_case_		    =     z + (z_q - z).detach()

							# reshape back to match original input shape
							snake_case_		    =     z_q.permute(0       ,	3       ,	1       ,	2						).contiguous()

							if self.remap is not None:
								snake_case_		    =     min_encoding_indices.reshape(z.shape[0]       ,	-1						)  # add batch axis
								snake_case_		    =     self.remap_to_used(__lowercase						)
								snake_case_		    =     min_encoding_indices.reshape(-1       ,	1						)  # flatten

							if self.sane_index_shape:
								snake_case_		    =     min_encoding_indices.reshape(z_q.shape[0]       ,	z_q.shape[2]       ,	z_q.shape[3]						)

							return z_q, loss, (perplexity, min_encodings, min_encoding_indices)






						def        snake_case__						(						self							:     Union[str, Any]       ,	__lowercase							:     Optional[int]       ,	__lowercase							:     Optional[int]						):






							"""simple docstring"""






							if self.remap is not None:
								snake_case_		    =     indices.reshape(shape[0]       ,	-1						)  # add batch axis
								snake_case_		    =     self.unmap_to_all(__lowercase						)
								snake_case_		    =     indices.reshape(-1						)  # flatten again

							# get quantized latent vectors
							snake_case_		    =     self.embedding(__lowercase						)

							if shape is not None:
								snake_case_		    =     z_q.view(__lowercase						)
								# reshape back to match original input shape
								snake_case_		    =     z_q.permute(0       ,	3       ,	1       ,	2						).contiguous()

							return z_q







class 		UpperCAmelCase   (				UpperCAmelCase__			):







						'''simple docstring'''





						def __init__(						self							:     Union[str, Any]       ,	__lowercase							:     List[Any]       ,	__lowercase							:     Any=False						):






							"""simple docstring"""






							snake_case_		    =     parameters
							snake_case_    ,      snake_case_		    =     torch.chunk(__lowercase       ,	2       ,	dim=1						)
							snake_case_		    =     torch.clamp(self.logvar       ,	-30.0       ,	20.0						)
							snake_case_		    =     deterministic
							snake_case_		    =     torch.exp(0.5 * self.logvar						)
							snake_case_		    =     torch.exp(self.logvar						)
							if self.deterministic:
								snake_case_		    =     snake_case_		    =     torch.zeros_like(
								    self.mean       ,	device=self.parameters.device       ,	dtype=self.parameters.dtype						)






						def        snake_case__						(						self							:     Optional[Any]       ,	__lowercase							:     Optional[torch.Generator] = None						):






							"""simple docstring"""






							snake_case_		    =     randn_tensor(
							    self.mean.shape       ,	generator=__lowercase       ,	device=self.parameters.device       ,	dtype=self.parameters.dtype						)
							snake_case_		    =     self.mean + self.std * sample
							return x






						def        snake_case__						(						self							:     List[Any]       ,	__lowercase							:     Dict=None						):






							"""simple docstring"""






							if self.deterministic:
								return torch.Tensor([0.0]						)
							else:
								if other is None:
									return 0.5 * torch.sum(torch.pow(self.mean       ,	2						) + self.var - 1.0 - self.logvar       ,	dim=[1, 2, 3]						)
								else:
									return 0.5 * torch.sum(
									    torch.pow(self.mean - other.mean       ,	2						) / other.var
									    + self.var / other.var
									    - 1.0
									    - self.logvar
									    + other.logvar       ,	dim=[1, 2, 3]       ,	)






						def        snake_case__						(						self							:     Optional[Any]       ,	__lowercase							:     Union[str, Any]       ,	__lowercase							:     Dict=[1, 2, 3]						):






							"""simple docstring"""






							if self.deterministic:
								return torch.Tensor([0.0]						)
							snake_case_		    =     np.log(2.0 * np.pi						)
							return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean       ,	2						) / self.var       ,	dim=__lowercase						)






						def        snake_case__						(						self							:     str						):






							"""simple docstring"""






							return self.mean






import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


# =================#
# UNet Conversion #
# =================#

lowercase__				:     Optional[int] 		=						[
    # (stable-diffusion, HF Diffusers)
    ("time_embed.0.weight", "time_embedding.linear_1.weight"),
    ("time_embed.0.bias", "time_embedding.linear_1.bias"),
    ("time_embed.2.weight", "time_embedding.linear_2.weight"),
    ("time_embed.2.bias", "time_embedding.linear_2.bias"),
    ("input_blocks.0.0.weight", "conv_in.weight"),
    ("input_blocks.0.0.bias", "conv_in.bias"),
    ("out.0.weight", "conv_norm_out.weight"),
    ("out.0.bias", "conv_norm_out.bias"),
    ("out.2.weight", "conv_out.weight"),
    ("out.2.bias", "conv_out.bias"),
]

lowercase__				:     List[Any] 		=						[
    # (stable-diffusion, HF Diffusers)
    ("in_layers.0", "norm1"),
    ("in_layers.2", "conv1"),
    ("out_layers.0", "norm2"),
    ("out_layers.3", "conv2"),
    ("emb_layers.1", "time_emb_proj"),
    ("skip_connection", "conv_shortcut"),
]

lowercase__				:     Optional[Any] 		=						[]
# hardcoded number of downblocks and resnets/attentions...
# would need smarter logic for other networks.
for i in range(4):
		# loop over downblocks/upblocks

		for j in range(2):
				# loop over resnets/attentions for downblocks
				lowercase__				:     str 		=						f'''down_blocks.{i}.resnets.{j}.'''
				lowercase__				:     Union[str, Any] 		=						f'''input_blocks.{3*i + j + 1}.0.'''
				unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))

				if i < 3:
						# no attention layers in down_blocks.3
						lowercase__				:     Tuple 		=						f'''down_blocks.{i}.attentions.{j}.'''
						lowercase__				:     Dict 		=						f'''input_blocks.{3*i + j + 1}.1.'''
						unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))

		for j in range(3):
				# loop over resnets/attentions for upblocks
				lowercase__				:     List[Any] 		=						f'''up_blocks.{i}.resnets.{j}.'''
				lowercase__				:     int 		=						f'''output_blocks.{3*i + j}.0.'''
				unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))

				if i > 0:
						# no attention layers in up_blocks.0
						lowercase__				:     List[str] 		=						f'''up_blocks.{i}.attentions.{j}.'''
						lowercase__				:     Tuple 		=						f'''output_blocks.{3*i + j}.1.'''
						unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))

		if i < 3:
				# no downsample in down_blocks.3
				lowercase__				:     List[str] 		=						f'''down_blocks.{i}.downsamplers.0.conv.'''
				lowercase__				:     Any 		=						f'''input_blocks.{3*(i+1)}.0.op.'''
				unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))

				# no upsample in up_blocks.3
				lowercase__				:     Optional[int] 		=						f'''up_blocks.{i}.upsamplers.0.'''
				lowercase__				:     int 		=						f'''output_blocks.{3*i + 2}.{1 if i == 0 else 2}.'''
				unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))

lowercase__				:     Union[str, Any] 		=						"mid_block.attentions.0."
lowercase__				:     List[Any] 		=						"middle_block.1."
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))

for j in range(2):
		lowercase__				:     Tuple 		=						f'''mid_block.resnets.{j}.'''
		lowercase__				:     List[str] 		=						f'''middle_block.{2*j}.'''
		unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	snake_case_		    =     {k: k for k in unet_state_dict.keys()}
	for sd_name, hf_name in unet_conversion_map:
		snake_case_		    =     sd_name
	for k, v in mapping.items():
		if "resnets" in k:
			for sd_part, hf_part in unet_conversion_map_resnet:
				snake_case_		    =     v.replace(_A     ,	_A							)
			snake_case_		    =     v
	for k, v in mapping.items():
		for sd_part, hf_part in unet_conversion_map_layer:
			snake_case_		    =     v.replace(_A     ,	_A							)
		snake_case_		    =     v
	snake_case_		    =     {v: unet_state_dict[k] for k, v in mapping.items()}
	return new_state_dict


# ================#
# VAE Conversion #
# ================#

lowercase__				:     Dict 		=						[
    # (stable-diffusion, HF Diffusers)
    ("nin_shortcut", "conv_shortcut"),
    ("norm_out", "conv_norm_out"),
    ("mid.attn_1.", "mid_block.attentions.0."),
]

for i in range(4):
		# down_blocks have two resnets
		for j in range(2):
				lowercase__				:     Any 		=						f'''encoder.down_blocks.{i}.resnets.{j}.'''
				lowercase__				:     List[Any] 		=						f'''encoder.down.{i}.block.{j}.'''
				vae_conversion_map.append((sd_down_prefix, hf_down_prefix))

		if i < 3:
				lowercase__				:     Optional[int] 		=						f'''down_blocks.{i}.downsamplers.0.'''
				lowercase__				:     Tuple 		=						f'''down.{i}.downsample.'''
				vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))

				lowercase__				:     List[str] 		=						f'''up_blocks.{i}.upsamplers.0.'''
				lowercase__				:     Optional[int] 		=						f'''up.{3-i}.upsample.'''
				vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))

		# up_blocks have three resnets
		# also, up blocks in hf are numbered in reverse from sd
		for j in range(3):
				lowercase__				:     int 		=						f'''decoder.up_blocks.{i}.resnets.{j}.'''
				lowercase__				:     Union[str, Any] 		=						f'''decoder.up.{3-i}.block.{j}.'''
				vae_conversion_map.append((sd_up_prefix, hf_up_prefix))

# this part accounts for mid blocks in both the encoder and the decoder
for i in range(2):
		lowercase__				:     Dict 		=						f'''mid_block.resnets.{i}.'''
		lowercase__				:     int 		=						f'''mid.block_{i+1}.'''
		vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))


lowercase__				:     Optional[Any] 		=						[
    # (stable-diffusion, HF Diffusers)
    ("norm.", "group_norm."),
    ("q.", "query."),
    ("k.", "key."),
    ("v.", "value."),
    ("proj_out.", "proj_attn."),
]


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	return w.reshape(*w.shape     ,	1     ,	1							)


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	snake_case_		    =     {k: k for k in vae_state_dict.keys()}
	for k, v in mapping.items():
		for sd_part, hf_part in vae_conversion_map:
			snake_case_		    =     v.replace(_A     ,	_A							)
		snake_case_		    =     v
	for k, v in mapping.items():
		if "attentions" in k:
			for sd_part, hf_part in vae_conversion_map_attn:
				snake_case_		    =     v.replace(_A     ,	_A							)
			snake_case_		    =     v
	snake_case_		    =     {v: vae_state_dict[k] for k, v in mapping.items()}
	snake_case_		    =     ["q", "k", "v", "proj_out"]
	for k, v in new_state_dict.items():
		for weight_name in weights_to_convert:
			if f"mid.attn_1.{weight_name}.weight" in k:
				print(f"Reshaping {k} for SD format"							)
				snake_case_		    =     reshape_weight_for_sd(_A							)
	return new_state_dict


# =========================#
# Text Encoder Conversion #
# =========================#


lowercase__				:     int 		=						[
    # (stable-diffusion, HF Diffusers)
    ("resblocks.", "text_model.encoder.layers."),
    ("ln_1", "layer_norm1"),
    ("ln_2", "layer_norm2"),
    (".c_fc.", ".fc1."),
    (".c_proj.", ".fc2."),
    (".attn", ".self_attn"),
    ("ln_final.", "transformer.text_model.final_layer_norm."),
    ("token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"),
    ("positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"),
]
lowercase__				:     Dict 		=						{re.escape(x[1]): x[0] for x in textenc_conversion_lst}
lowercase__				:     Tuple 		=						re.compile("|".join(protected.keys()))

# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
lowercase__				:     Dict 		=						{"q": 0, "k": 1, "v": 2}


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	snake_case_		    =     {}
	snake_case_		    =     {}
	snake_case_		    =     {}
	for k, v in text_enc_dict.items():
		if (
		    k.endswith(".self_attn.q_proj.weight"							)
		    or k.endswith(".self_attn.k_proj.weight"							)
		    or k.endswith(".self_attn.v_proj.weight"							)
		):
			snake_case_		    =     k[: -len(".q_proj.weight"							)]
			snake_case_		    =     k[-len("q_proj.weight"							)]
			if k_pre not in capture_qkv_weight:
				snake_case_		    =     [None, None, None]
			snake_case_		    =     v
			continue

		if (
		    k.endswith(".self_attn.q_proj.bias"							)
		    or k.endswith(".self_attn.k_proj.bias"							)
		    or k.endswith(".self_attn.v_proj.bias"							)
		):
			snake_case_		    =     k[: -len(".q_proj.bias"							)]
			snake_case_		    =     k[-len("q_proj.bias"							)]
			if k_pre not in capture_qkv_bias:
				snake_case_		    =     [None, None, None]
			snake_case_		    =     v
			continue

		snake_case_		    =     textenc_pattern.sub(lambda _A							: protected[re.escape(m.group(0							)							)]     ,	_A							)
		snake_case_		    =     v

	for k_pre, tensors in capture_qkv_weight.items():
		if None in tensors:
			raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing"							)
		snake_case_		    =     textenc_pattern.sub(lambda _A							: protected[re.escape(m.group(0							)							)]     ,	_A							)
		snake_case_		    =     torch.cat(_A							)

	for k_pre, tensors in capture_qkv_bias.items():
		if None in tensors:
			raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing"							)
		snake_case_		    =     textenc_pattern.sub(lambda _A							: protected[re.escape(m.group(0							)							)]     ,	_A							)
		snake_case_		    =     torch.cat(_A							)

	return new_state_dict


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	return text_enc_dict


if __name__ == "__main__":
		lowercase__				:     int 		=						argparse.ArgumentParser()

		parser.add_argument("--model_path", default=None, type=str, required=True, help="Path to the model to convert.")
		parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the output model.")
		parser.add_argument("--half", action="store_true", help="Save weights in half precision.")
		parser.add_argument(
		    "--use_safetensors", action="store_true", help="Save weights use safetensors, default is ckpt."
		)

		lowercase__				:     Dict 		=						parser.parse_args()

		assert args.model_path is not None, "Must provide a model path!"

		assert args.checkpoint_path is not None, "Must provide a checkpoint path!"

		# Path for safetensors
		lowercase__				:     Tuple 		=						osp.join(args.model_path, "unet", "diffusion_pytorch_model.safetensors")
		lowercase__				:     int 		=						osp.join(args.model_path, "vae", "diffusion_pytorch_model.safetensors")
		lowercase__				:     Any 		=						osp.join(args.model_path, "text_encoder", "model.safetensors")

		# Load models from safetensors if it exists, if it doesn't pytorch
		if osp.exists(unet_path):
				lowercase__				:     str 		=						load_file(unet_path, device="cpu")
		else:
				lowercase__				:     Optional[Any] 		=						osp.join(args.model_path, "unet", "diffusion_pytorch_model.bin")
				lowercase__				:     Any 		=						torch.load(unet_path, map_location="cpu")

		if osp.exists(vae_path):
				lowercase__				:     Tuple 		=						load_file(vae_path, device="cpu")
		else:
				lowercase__				:     Any 		=						osp.join(args.model_path, "vae", "diffusion_pytorch_model.bin")
				lowercase__				:     Dict 		=						torch.load(vae_path, map_location="cpu")

		if osp.exists(text_enc_path):
				lowercase__				:     Union[str, Any] 		=						load_file(text_enc_path, device="cpu")
		else:
				lowercase__				:     Union[str, Any] 		=						osp.join(args.model_path, "text_encoder", "pytorch_model.bin")
				lowercase__				:     Optional[int] 		=						torch.load(text_enc_path, map_location="cpu")

		# Convert the UNet model
		lowercase__				:     Dict 		=						convert_unet_state_dict(unet_state_dict)
		lowercase__				:     Any 		=						{"model.diffusion_model." + k: v for k, v in unet_state_dict.items()}

		# Convert the VAE model
		lowercase__				:     Dict 		=						convert_vae_state_dict(vae_state_dict)
		lowercase__				:     Union[str, Any] 		=						{"first_stage_model." + k: v for k, v in vae_state_dict.items()}

		# Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper
		lowercase__				:     Optional[Any] 		=						"text_model.encoder.layers.22.layer_norm2.bias" in text_enc_dict

		if is_vaa_model:
				# Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm
				lowercase__				:     Any 		=						{"transformer." + k: v for k, v in text_enc_dict.items()}
				lowercase__				:     List[Any] 		=						convert_text_enc_state_dict_vaa(text_enc_dict)
				lowercase__				:     List[str] 		=						{"cond_stage_model.model." + k: v for k, v in text_enc_dict.items()}
		else:
				lowercase__				:     Tuple 		=						convert_text_enc_state_dict(text_enc_dict)
				lowercase__				:     Any 		=						{"cond_stage_model.transformer." + k: v for k, v in text_enc_dict.items()}

		# Put together new checkpoint
		lowercase__				:     Tuple 		=						{**unet_state_dict, **vae_state_dict, **text_enc_dict}
		if args.half:
				lowercase__				:     Any 		=						{k: v.half() for k, v in state_dict.items()}

		if args.use_safetensors:
				save_file(state_dict, args.checkpoint_path)
		else:
				lowercase__				:     Union[str, Any] 		=						{"state_dict": state_dict}
				torch.save(state_dict, args.checkpoint_path)
