



import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()




class 							lowerCAmelCase__  (	unittest.TestCase						):







					def  _lowercase			(       self      :     Any):
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()







					@property
					def  _lowercase			(       self      :     Dict):
											torch.manual_seed(0)
											A__	:							Union[str, Any]						    = UNetaDModel(
											    sample_size=(32, 64)			,		in_channels=1			,		out_channels=1			,		layers_per_block=2			,		block_out_channels=(128, 128)			,		down_block_types=("AttnDownBlock2D", "DownBlock2D")			,		up_block_types=("UpBlock2D", "AttnUpBlock2D")			,		)
											return model







					@property
					def  _lowercase			(       self      :     Tuple):
											torch.manual_seed(0)
											A__	:							List[str]						    = UNetaDConditionModel(
											    sample_size=(64, 32)			,		in_channels=1			,		out_channels=1			,		layers_per_block=2			,		block_out_channels=(128, 128)			,		down_block_types=("CrossAttnDownBlock2D", "DownBlock2D")			,		up_block_types=("UpBlock2D", "CrossAttnUpBlock2D")			,		cross_attention_dim=10			,		)
											return model







					@property
					def  _lowercase			(       self      :     Optional[int]):
											torch.manual_seed(0)
											A__	:							Union[str, Any]						    = AutoencoderKL(
											    sample_size=(128, 64)			,		in_channels=1			,		out_channels=1			,		latent_channels=1			,		layers_per_block=2			,		block_out_channels=(128, 128)			,		down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D")			,		up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D")			,		)
											A__	:							str						    = UNetaDModel(
											    sample_size=(64, 32)			,		in_channels=1			,		out_channels=1			,		layers_per_block=2			,		block_out_channels=(128, 128)			,		down_block_types=("AttnDownBlock2D", "DownBlock2D")			,		up_block_types=("UpBlock2D", "AttnUpBlock2D")			,		)
											return vqvae, unet







					@slow
					def  _lowercase			(       self      :     Union[str, Any]):
											A__	:							Dict						    = "cpu"  # ensure determinism for the device-dependent torch.Generator
											A__	:							Union[str, Any]						    = Mel(
											    x_res=self.dummy_unet.config.sample_size[1]			,		y_res=self.dummy_unet.config.sample_size[0]			,		)

											A__	:							Dict						    = DDPMScheduler()
											A__	:							Optional[Any]						    = AudioDiffusionPipeline(vqvae=_A			,		unet=self.dummy_unet			,		mel=_A			,		scheduler=_A)
											A__	:							Union[str, Any]						    = pipe.to(_A)
											pipe.set_progress_bar_config(disable=_A)

											A__	:							int						    = torch.Generator(device=_A).manual_seed(42)
											A__	:							str						    = pipe(generator=_A			,		steps=4)
											A__	:							Optional[int]						    = output.audios[0]
											A__	:							Union[str, Any]						    = output.images[0]

											A__	:							Dict						    = torch.Generator(device=_A).manual_seed(42)
											A__	:							Dict						    = pipe(generator=_A			,		steps=4			,		return_dict=_A)
											A__	:							Optional[Any]						    = output[0][0]

											assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
											assert (
											    image.height == self.dummy_unet.config.sample_size[0]
											    and image.width == self.dummy_unet.config.sample_size[1]
											)
											A__	:							Any						    = np.frombuffer(image.tobytes()			,		dtype="uint8")[:10]
											A__	:							List[str]						    = np.frombuffer(image_from_tuple.tobytes()			,		dtype="uint8")[:10]
											A__	:							Optional[Any]						    = np.array([69, 255, 255, 255, 0, 0, 77, 181, 12, 127])

											assert np.abs(image_slice.flatten() - expected_slice).max() == 0
											assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() == 0

											A__	:							str						    = Mel(
											    x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1]			,		y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0]			,		)

											A__	:							Union[str, Any]						    = DDIMScheduler()
											A__	:							Any						    = self.dummy_vqvae_and_unet
											A__	:							Optional[int]						    = AudioDiffusionPipeline(
											    vqvae=self.dummy_vqvae_and_unet[0]			,		unet=dummy_vqvae_and_unet[1]			,		mel=_A			,		scheduler=_A)
											A__	:							int						    = pipe.to(_A)
											pipe.set_progress_bar_config(disable=_A)

											np.random.seed(0)
											A__	:							Optional[Any]						    = np.random.uniform(-1			,		1			,		((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,))
											A__	:							Optional[int]						    = torch.Generator(device=_A).manual_seed(42)
											A__	:							int						    = pipe(raw_audio=_A			,		generator=_A			,		start_step=5			,		steps=10)
											A__	:							List[str]						    = output.images[0]

											assert (
											    image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
											    and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
											)
											A__	:							Optional[Any]						    = np.frombuffer(image.tobytes()			,		dtype="uint8")[:10]
											A__	:							Tuple						    = np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121])

											assert np.abs(image_slice.flatten() - expected_slice).max() == 0

											A__	:							str						    = self.dummy_unet_condition
											A__	:							List[Any]						    = AudioDiffusionPipeline(
											    vqvae=self.dummy_vqvae_and_unet[0]			,		unet=_A			,		mel=_A			,		scheduler=_A)
											A__	:							Optional[Any]						    = pipe.to(_A)
											pipe.set_progress_bar_config(disable=_A)

											np.random.seed(0)
											A__	:							Tuple						    = torch.rand((1, 1, 10))
											A__	:							Optional[Any]						    = pipe(generator=_A			,		encoding=_A)
											A__	:							Dict						    = output.images[0]
											A__	:							Tuple						    = np.frombuffer(image.tobytes()			,		dtype="uint8")[:10]
											A__	:							Any						    = np.array([107, 103, 120, 127, 142, 122, 113, 122, 97, 111])

											assert np.abs(image_slice.flatten() - expected_slice).max() == 0




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







					def  _lowercase			(       self      :     int):
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()







					def  _lowercase			(       self      :     Dict):
											A__	:							Optional[int]						    = torch_device

											A__	:							Optional[Any]						    = DiffusionPipeline.from_pretrained("teticio/audio-diffusion-ddim-256")
											A__	:							Optional[int]						    = pipe.to(_A)
											pipe.set_progress_bar_config(disable=_A)

											A__	:							List[Any]						    = torch.Generator(device=_A).manual_seed(42)
											A__	:							Tuple						    = pipe(generator=_A)
											A__	:							Tuple						    = output.audios[0]
											A__	:							str						    = output.images[0]

											assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
											assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
											A__	:							List[Any]						    = np.frombuffer(image.tobytes()			,		dtype="uint8")[:10]
											A__	:							str						    = np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26])

											assert np.abs(image_slice.flatten() - expected_slice).max() == 0




import argparse
import io

import requests
import torch
from omegaconf import OmegaConf

from diffusers import AutoencoderKL
from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
    assign_to_checkpoint,
    conv_attn_to_linear,
    create_vae_diffusers_config,
    renew_vae_attention_paths,
    renew_vae_resnet_paths,
)
def 	snake_case__       (					__lowercase							,     __lowercase       )		->		int:







						"""simple docstring"""

						A__	:							List[str]						    = checkpoint

						A__	:							Dict						    = {}

						A__	:							Union[str, Any]						    = vae_state_dict["encoder.conv_in.weight"]
						A__	:							int						    = vae_state_dict["encoder.conv_in.bias"]
						A__	:							List[Any]						    = vae_state_dict["encoder.conv_out.weight"]
						A__	:							Optional[int]						    = vae_state_dict["encoder.conv_out.bias"]
						A__	:							int						    = vae_state_dict["encoder.norm_out.weight"]
						A__	:							Optional[int]						    = vae_state_dict["encoder.norm_out.bias"]

						A__	:							Optional[int]						    = vae_state_dict["decoder.conv_in.weight"]
						A__	:							Optional[Any]						    = vae_state_dict["decoder.conv_in.bias"]
						A__	:							str						    = vae_state_dict["decoder.conv_out.weight"]
						A__	:							Optional[int]						    = vae_state_dict["decoder.conv_out.bias"]
						A__	:							Union[str, Any]						    = vae_state_dict["decoder.norm_out.weight"]
						A__	:							Optional[Any]						    = vae_state_dict["decoder.norm_out.bias"]

						A__	:							Optional[Any]						    = vae_state_dict["quant_conv.weight"]
						A__	:							Optional[Any]						    = vae_state_dict["quant_conv.bias"]
						A__	:							Dict						    = vae_state_dict["post_quant_conv.weight"]
						A__	:							Union[str, Any]						    = vae_state_dict["post_quant_conv.bias"]

						# Retrieves the keys for the encoder down blocks only
						A__	:							str						    = len({".".join(layer.split("."       )[:3]       ) for layer in vae_state_dict if "encoder.down" in layer}       )
						A__	:							Optional[int]						    = {
						    layer_id: [key for key in vae_state_dict if F'down.{layer_id}' in key] for layer_id in range(__lowercase       )
						}

						# Retrieves the keys for the decoder up blocks only
						A__	:							Any						    = len({".".join(layer.split("."       )[:3]       ) for layer in vae_state_dict if "decoder.up" in layer}       )
						A__	:							Any						    = {
						    layer_id: [key for key in vae_state_dict if F'up.{layer_id}' in key] for layer_id in range(__lowercase       )
						}

						for i in range(__lowercase       ):
												A__	:							Any						    = [key for key in down_blocks[i] if F'down.{i}' in key and F'down.{i}.downsample' not in key]

												if F'encoder.down.{i}.downsample.conv.weight' in vae_state_dict:
																		A__	:							Any						    = vae_state_dict.pop(
																		    F'encoder.down.{i}.downsample.conv.weight'       )
																		A__	:							Optional[int]						    = vae_state_dict.pop(
																		    F'encoder.down.{i}.downsample.conv.bias'       )

												A__	:							Any						    = renew_vae_resnet_paths(__lowercase       )
												A__	:							int						    = {"old": F'down.{i}.block', "new": F'down_blocks.{i}.resnets'}
												assign_to_checkpoint(__lowercase							,     __lowercase							,     __lowercase							,     additional_replacements=[meta_path]							,     config=__lowercase       )

						A__	:							int						    = [key for key in vae_state_dict if "encoder.mid.block" in key]
						A__	:							Any						    = 2
						for i in range(1							,     num_mid_res_blocks + 1       ):
												A__	:							Optional[int]						    = [key for key in mid_resnets if F'encoder.mid.block_{i}' in key]

												A__	:							Optional[Any]						    = renew_vae_resnet_paths(__lowercase       )
												A__	:							int						    = {"old": F'mid.block_{i}', "new": F'mid_block.resnets.{i - 1}'}
												assign_to_checkpoint(__lowercase							,     __lowercase							,     __lowercase							,     additional_replacements=[meta_path]							,     config=__lowercase       )

						A__	:							str						    = [key for key in vae_state_dict if "encoder.mid.attn" in key]
						A__	:							List[Any]						    = renew_vae_attention_paths(__lowercase       )
						A__	:							Optional[int]						    = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
						assign_to_checkpoint(__lowercase							,     __lowercase							,     __lowercase							,     additional_replacements=[meta_path]							,     config=__lowercase       )
						conv_attn_to_linear(__lowercase       )

						for i in range(__lowercase       ):
												A__	:							Dict						    = num_up_blocks - 1 - i
												A__	:							int						    = [
												    key for key in up_blocks[block_id] if F'up.{block_id}' in key and F'up.{block_id}.upsample' not in key
												]

												if F'decoder.up.{block_id}.upsample.conv.weight' in vae_state_dict:
																		A__	:							Tuple						    = vae_state_dict[
																		    F'decoder.up.{block_id}.upsample.conv.weight'
																		]
																		A__	:							List[Any]						    = vae_state_dict[
																		    F'decoder.up.{block_id}.upsample.conv.bias'
																		]

												A__	:							str						    = renew_vae_resnet_paths(__lowercase       )
												A__	:							List[str]						    = {"old": F'up.{block_id}.block', "new": F'up_blocks.{i}.resnets'}
												assign_to_checkpoint(__lowercase							,     __lowercase							,     __lowercase							,     additional_replacements=[meta_path]							,     config=__lowercase       )

						A__	:							int						    = [key for key in vae_state_dict if "decoder.mid.block" in key]
						A__	:							Optional[int]						    = 2
						for i in range(1							,     num_mid_res_blocks + 1       ):
												A__	:							Union[str, Any]						    = [key for key in mid_resnets if F'decoder.mid.block_{i}' in key]

												A__	:							Dict						    = renew_vae_resnet_paths(__lowercase       )
												A__	:							int						    = {"old": F'mid.block_{i}', "new": F'mid_block.resnets.{i - 1}'}
												assign_to_checkpoint(__lowercase							,     __lowercase							,     __lowercase							,     additional_replacements=[meta_path]							,     config=__lowercase       )

						A__	:							Optional[Any]						    = [key for key in vae_state_dict if "decoder.mid.attn" in key]
						A__	:							Union[str, Any]						    = renew_vae_attention_paths(__lowercase       )
						A__	:							List[str]						    = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
						assign_to_checkpoint(__lowercase							,     __lowercase							,     __lowercase							,     additional_replacements=[meta_path]							,     config=__lowercase       )
						conv_attn_to_linear(__lowercase       )
						return new_checkpoint
def 	snake_case__       (					__lowercase							,     __lowercase							,     )		->		Tuple:







						"""simple docstring"""

						A__	:							List[str]						    = requests.get(
						    " https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml"       )
						A__	:							Dict						    = io.BytesIO(r.content       )

						A__	:							str						    = OmegaConf.load(__lowercase       )
						A__	:							Any						    = 5_1_2
						A__	:							Optional[Any]						    = "cuda" if torch.cuda.is_available() else "cpu"
						if checkpoint_path.endswith("safetensors"       ):
												from safetensors import safe_open

												A__	:							str						    = {}
												with safe_open(__lowercase							,     framework="pt"							,     device="cpu"       ) as f:
																		for key in f.keys():
																								A__	:							List[str]						    = f.get_tensor(__lowercase       )
						else:
												A__	:							str						    = torch.load(__lowercase							,     map_location=__lowercase       )["state_dict"]

						# Convert the VAE model.
						A__	:							List[str]						    = create_vae_diffusers_config(__lowercase							,     image_size=__lowercase       )
						A__	:							List[Any]						    = custom_convert_ldm_vae_checkpoint(__lowercase							,     __lowercase       )

						A__	:							Dict						    = AutoencoderKL(**__lowercase       )
						vae.load_state_dict(__lowercase       )
						vae.save_pretrained(__lowercase       )


if __name__ == "__main__":
		snake_case    :						List[Any]  =						argparse.ArgumentParser()

		parser.add_argument('--vae_pt_path', default=None, type=str, required=True, help='Path to the VAE.pt to convert.')
		parser.add_argument('--dump_path', default=None, type=str, required=True, help='Path to the VAE.pt to convert.')

		snake_case    :						int  =						parser.parse_args()

		vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)