





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 lowerCAmelCase__(__snake_case       ,__snake_case							)				->      Union[str, Any]:







	'''simple docstring'''







	lowerCamelCase__								=		checkpoint

	lowerCamelCase__								=		{}

	lowerCamelCase__								=		vae_state_dict['''encoder.conv_in.weight''']
	lowerCamelCase__								=		vae_state_dict['''encoder.conv_in.bias''']
	lowerCamelCase__								=		vae_state_dict['''encoder.conv_out.weight''']
	lowerCamelCase__								=		vae_state_dict['''encoder.conv_out.bias''']
	lowerCamelCase__								=		vae_state_dict['''encoder.norm_out.weight''']
	lowerCamelCase__								=		vae_state_dict['''encoder.norm_out.bias''']

	lowerCamelCase__								=		vae_state_dict['''decoder.conv_in.weight''']
	lowerCamelCase__								=		vae_state_dict['''decoder.conv_in.bias''']
	lowerCamelCase__								=		vae_state_dict['''decoder.conv_out.weight''']
	lowerCamelCase__								=		vae_state_dict['''decoder.conv_out.bias''']
	lowerCamelCase__								=		vae_state_dict['''decoder.norm_out.weight''']
	lowerCamelCase__								=		vae_state_dict['''decoder.norm_out.bias''']

	lowerCamelCase__								=		vae_state_dict['''quant_conv.weight''']
	lowerCamelCase__								=		vae_state_dict['''quant_conv.bias''']
	lowerCamelCase__								=		vae_state_dict['''post_quant_conv.weight''']
	lowerCamelCase__								=		vae_state_dict['''post_quant_conv.bias''']

	# Retrieves the keys for the encoder down blocks only
	lowerCamelCase__								=		len({'''.'''.join(layer.split('''.'''							)[:3]							) for layer in vae_state_dict if '''encoder.down''' in layer}							)
	lowerCamelCase__								=		{
	    layer_id: [key for key in vae_state_dict if F'down.{layer_id}' in key] for layer_id in range(__snake_case							)
	}

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

	for i in range(__snake_case							):
		lowerCamelCase__								=		[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:
			lowerCamelCase__								=		vae_state_dict.pop(
			    F'encoder.down.{i}.downsample.conv.weight'							)
			lowerCamelCase__								=		vae_state_dict.pop(
			    F'encoder.down.{i}.downsample.conv.bias'							)

		lowerCamelCase__								=		renew_vae_resnet_paths(__snake_case							)
		lowerCamelCase__								=		{'''old''': F'down.{i}.block', '''new''': F'down_blocks.{i}.resnets'}
		assign_to_checkpoint(__snake_case       ,__snake_case       ,__snake_case       ,additional_replacements=[meta_path]       ,config=__snake_case							)

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

		lowerCamelCase__								=		renew_vae_resnet_paths(__snake_case							)
		lowerCamelCase__								=		{'''old''': F'mid.block_{i}', '''new''': F'mid_block.resnets.{i - 1}'}
		assign_to_checkpoint(__snake_case       ,__snake_case       ,__snake_case       ,additional_replacements=[meta_path]       ,config=__snake_case							)

	lowerCamelCase__								=		[key for key in vae_state_dict if '''encoder.mid.attn''' in key]
	lowerCamelCase__								=		renew_vae_attention_paths(__snake_case							)
	lowerCamelCase__								=		{'''old''': '''mid.attn_1''', '''new''': '''mid_block.attentions.0'''}
	assign_to_checkpoint(__snake_case       ,__snake_case       ,__snake_case       ,additional_replacements=[meta_path]       ,config=__snake_case							)
	conv_attn_to_linear(__snake_case							)

	for i in range(__snake_case							):
		lowerCamelCase__								=		num_up_blocks - 1 - i
		lowerCamelCase__								=		[
		    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:
			lowerCamelCase__								=		vae_state_dict[
			    F'decoder.up.{block_id}.upsample.conv.weight'
			]
			lowerCamelCase__								=		vae_state_dict[
			    F'decoder.up.{block_id}.upsample.conv.bias'
			]

		lowerCamelCase__								=		renew_vae_resnet_paths(__snake_case							)
		lowerCamelCase__								=		{'''old''': F'up.{block_id}.block', '''new''': F'up_blocks.{i}.resnets'}
		assign_to_checkpoint(__snake_case       ,__snake_case       ,__snake_case       ,additional_replacements=[meta_path]       ,config=__snake_case							)

	lowerCamelCase__								=		[key for key in vae_state_dict if '''decoder.mid.block''' in key]
	lowerCamelCase__								=		2
	for i in range(1       ,num_mid_res_blocks + 1							):
		lowerCamelCase__								=		[key for key in mid_resnets if F'decoder.mid.block_{i}' in key]

		lowerCamelCase__								=		renew_vae_resnet_paths(__snake_case							)
		lowerCamelCase__								=		{'''old''': F'mid.block_{i}', '''new''': F'mid_block.resnets.{i - 1}'}
		assign_to_checkpoint(__snake_case       ,__snake_case       ,__snake_case       ,additional_replacements=[meta_path]       ,config=__snake_case							)

	lowerCamelCase__								=		[key for key in vae_state_dict if '''decoder.mid.attn''' in key]
	lowerCamelCase__								=		renew_vae_attention_paths(__snake_case							)
	lowerCamelCase__								=		{'''old''': '''mid.attn_1''', '''new''': '''mid_block.attentions.0'''}
	assign_to_checkpoint(__snake_case       ,__snake_case       ,__snake_case       ,additional_replacements=[meta_path]       ,config=__snake_case							)
	conv_attn_to_linear(__snake_case							)
	return new_checkpoint





def lowerCAmelCase__(__snake_case       ,__snake_case       ,)				->      Union[str, Any]:







	'''simple docstring'''







	lowerCamelCase__								=		requests.get(
	    ''' https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml'''							)
	lowerCamelCase__								=		io.BytesIO(r.content							)

	lowerCamelCase__								=		OmegaConf.load(__snake_case							)
	lowerCamelCase__								=		512
	lowerCamelCase__								=		'''cuda''' if torch.cuda.is_available() else '''cpu'''
	if checkpoint_path.endswith('''safetensors'''							):
		from safetensors import safe_open

		lowerCamelCase__								=		{}
		with safe_open(__snake_case       ,framework='''pt'''       ,device='''cpu'''							) as f:
			for key in f.keys():
				lowerCamelCase__								=		f.get_tensor(__snake_case							)
	else:
		lowerCamelCase__								=		torch.load(__snake_case       ,map_location=__snake_case							)['''state_dict''']

	# Convert the VAE model.
	lowerCamelCase__								=		create_vae_diffusers_config(__snake_case       ,image_size=__snake_case							)
	lowerCamelCase__								=		custom_convert_ldm_vae_checkpoint(__snake_case       ,__snake_case							)

	lowerCamelCase__								=		AutoencoderKL(**__snake_case							)
	vae.load_state_dict(__snake_case							)
	vae.save_pretrained(__snake_case							)


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

			_a 						=      parser.parse_args()

			vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)










import pytest

from datasets.parallel import ParallelBackendConfig, parallel_backend
from datasets.utils.py_utils import map_nested

from .utils import require_dill_gt_0_3_2, require_joblibspark, require_not_windows





def lowerCAmelCase__(__snake_case							)				->      int:  # picklable for multiprocessing







	'''simple docstring'''







	return i + 1





@require_dill_gt_0_3_2
@require_joblibspark
@require_not_windows
def lowerCAmelCase__()				->      Any:







	'''simple docstring'''







	with parallel_backend('''spark'''							):
		assert ParallelBackendConfig.backend_name == "spark"

	lowerCamelCase__								=		[1, 2, 3]
	with pytest.raises(__snake_case							):
		with parallel_backend('''unsupported backend'''							):
			map_nested(__snake_case       ,__snake_case       ,num_proc=2							)

	with pytest.raises(__snake_case							):
		with parallel_backend('''unsupported backend'''							):
			map_nested(__snake_case       ,__snake_case       ,num_proc=-1							)





@require_dill_gt_0_3_2
@require_joblibspark
@require_not_windows
@pytest.mark.parametrize('''num_proc'''       ,[2, -1]							)
def lowerCAmelCase__(__snake_case							)				->      Tuple:







	'''simple docstring'''







	lowerCamelCase__								=		[1, 2]
	lowerCamelCase__								=		{'''a''': 1, '''b''': 2}
	lowerCamelCase__								=		{'''a''': [1, 2], '''b''': [3, 4]}
	lowerCamelCase__								=		{'''a''': {'''1''': 1}, '''b''': 2}
	lowerCamelCase__								=		{'''a''': 1, '''b''': 2, '''c''': 3, '''d''': 4}
	lowerCamelCase__								=		[2, 3]
	lowerCamelCase__								=		{'''a''': 2, '''b''': 3}
	lowerCamelCase__								=		{'''a''': [2, 3], '''b''': [4, 5]}
	lowerCamelCase__								=		{'''a''': {'''1''': 2}, '''b''': 3}
	lowerCamelCase__								=		{'''a''': 2, '''b''': 3, '''c''': 4, '''d''': 5}

	with parallel_backend('''spark'''							):
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa



