
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        UpperCamelCase    (							lowerCAmelCase__ ,			lowerCAmelCase__		):
					'''simple docstring'''
					lowercase							=						checkpoint

					lowercase							=						{}

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

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

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

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

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

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

										lowercase							=						renew_vae_resnet_paths(lowerCAmelCase__		)
										lowercase							=						{'''old''': f'down.{i}.block', '''new''': f'down_blocks.{i}.resnets'}
										assign_to_checkpoint(lowerCAmelCase__ ,			lowerCAmelCase__ ,			lowerCAmelCase__ ,			additional_replacements=[meta_path] ,			config=lowerCAmelCase__		)

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

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

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

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

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

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

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

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



def        UpperCamelCase    (							lowerCAmelCase__ ,			lowerCAmelCase__ ,			):
					'''simple docstring'''
					# Only support V1
					lowercase							=						requests.get(
					    ''' https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml'''		)
					lowercase							=						io.BytesIO(r.content		)

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

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

					# Convert the VAE model.
					lowercase							=						create_vae_diffusers_config(lowerCAmelCase__ ,			image_size=lowerCAmelCase__		)
					lowercase							=						custom_convert_ldm_vae_checkpoint(lowerCAmelCase__ ,			lowerCAmelCase__		)

					lowercase							=						AutoencoderKL(**lowerCAmelCase__		)
					vae.load_state_dict(lowerCAmelCase__		)
					vae.save_pretrained(lowerCAmelCase__		)


if __name__ == "__main__":
							lowercase__		:List[str]   =    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.")

							lowercase__		:int   =    parser.parse_args()

							vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)






import gc
import unittest

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

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference







class 				lowercase		(					SCREAMING_SNAKE_CASE__       ,  unittest.TestCase		):
							lowercase_       :       Tuple       					=ShapEPipeline
							lowercase_       :       List[Any]       					=['''prompt''']
							lowercase_       :       int       					=['''prompt''']
							lowercase_       :       Union[str, Any]       					=[
							    '''num_images_per_prompt''',
							    '''num_inference_steps''',
							    '''generator''',
							    '''latents''',
							    '''guidance_scale''',
							    '''frame_size''',
							    '''output_type''',
							    '''return_dict''',
							]
							lowercase_       :       Optional[int]       					=False




							@property
							def  A__      (				self):
												return 3_2




							@property
							def  A__      (				self):
												return 3_2




							@property
							def  A__      (				self):
												return self.time_input_dim * 4




							@property
							def  A__      (				self):
												return 8




							@property
							def  A__      (				self):
												lowercase							=						CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')
												return tokenizer




							@property
							def  A__      (				self):
												torch.manual_seed(0)
												lowercase							=						CLIPTextConfig(
												    bos_token_id=0   ,eos_token_id=2   ,hidden_size=self.text_embedder_hidden_size   ,projection_dim=self.text_embedder_hidden_size   ,intermediate_size=3_7   ,layer_norm_eps=1E-05   ,num_attention_heads=4   ,num_hidden_layers=5   ,pad_token_id=1   ,vocab_size=1_0_0_0   ,)
												return CLIPTextModelWithProjection(A__)




							@property
							def  A__      (				self):
												torch.manual_seed(0)

												lowercase							=						{
												    '''num_attention_heads''': 2,
												    '''attention_head_dim''': 1_6,
												    '''embedding_dim''': self.time_input_dim,
												    '''num_embeddings''': 3_2,
												    '''embedding_proj_dim''': self.text_embedder_hidden_size,
												    '''time_embed_dim''': self.time_embed_dim,
												    '''num_layers''': 1,
												    '''clip_embed_dim''': self.time_input_dim * 2,
												    '''additional_embeddings''': 0,
												    '''time_embed_act_fn''': '''gelu''',
												    '''norm_in_type''': '''layer''',
												    '''encoder_hid_proj_type''': None,
												    '''added_emb_type''': None,
												}

												lowercase							=						PriorTransformer(**A__)
												return model




							@property
							def  A__      (				self):
												torch.manual_seed(0)

												lowercase							=						{
												    '''param_shapes''': (
												        (self.renderer_dim, 9_3),
												        (self.renderer_dim, 8),
												        (self.renderer_dim, 8),
												        (self.renderer_dim, 8),
												    ),
												    '''d_latent''': self.time_input_dim,
												    '''d_hidden''': self.renderer_dim,
												    '''n_output''': 1_2,
												    '''background''': (
												        0.1,
												        0.1,
												        0.1,
												    ),
												}
												lowercase							=						ShapERenderer(**A__)
												return model




							def  A__      (				self):
												lowercase							=						self.dummy_prior
												lowercase							=						self.dummy_text_encoder
												lowercase							=						self.dummy_tokenizer
												lowercase							=						self.dummy_renderer

												lowercase							=						HeunDiscreteScheduler(
												    beta_schedule='''exp'''   ,num_train_timesteps=1_0_2_4   ,prediction_type='''sample'''   ,use_karras_sigmas=A__   ,clip_sample=A__   ,clip_sample_range=1.0   ,)
												lowercase							=						{
												    '''prior''': prior,
												    '''text_encoder''': text_encoder,
												    '''tokenizer''': tokenizer,
												    '''renderer''': renderer,
												    '''scheduler''': scheduler,
												}

												return components




							def  A__      (				self   ,A__   ,A__=0):
												if str(A__).startswith('''mps'''):
																	lowercase							=						torch.manual_seed(A__)
												else:
																	lowercase							=						torch.Generator(device=A__).manual_seed(A__)
												lowercase							=						{
												    '''prompt''': '''horse''',
												    '''generator''': generator,
												    '''num_inference_steps''': 1,
												    '''frame_size''': 3_2,
												    '''output_type''': '''np''',
												}
												return inputs




							def  A__      (				self):
												lowercase							=						'''cpu'''

												lowercase							=						self.get_dummy_components()

												lowercase							=						self.pipeline_class(**A__)
												lowercase							=						pipe.to(A__)

												pipe.set_progress_bar_config(disable=A__)

												lowercase							=						pipe(**self.get_dummy_inputs(A__))
												lowercase							=						output.images[0]
												lowercase							=						image[0, -3:, -3:, -1]

												assert image.shape == (2_0, 3_2, 3_2, 3)

												lowercase							=						np.array(
												    [
												        0.00039216,
												        0.00039216,
												        0.00039216,
												        0.00039216,
												        0.00039216,
												        0.00039216,
												        0.00039216,
												        0.00039216,
												        0.00039216,
												    ])

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




							def  A__      (				self):
												# NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
												self._test_inference_batch_consistent(batch_sizes=[1, 2])




							def  A__      (				self):
												lowercase							=						torch_device == '''cpu'''
												lowercase							=						True

												self._test_inference_batch_single_identical(
												    batch_size=2   ,test_max_difference=A__   ,relax_max_difference=A__   ,)





							def  A__      (				self):
												lowercase							=						self.get_dummy_components()
												lowercase							=						self.pipeline_class(**A__)
												lowercase							=						pipe.to(A__)
												pipe.set_progress_bar_config(disable=A__)

												lowercase							=						1
												lowercase							=						2

												lowercase							=						self.get_dummy_inputs(A__)

												for key in inputs.keys():
																	if key in self.batch_params:
																						lowercase							=						batch_size * [inputs[key]]

												lowercase							=						pipe(**A__   ,num_images_per_prompt=A__)[0]

												assert images.shape[0] == batch_size * num_images_per_prompt







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




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





							def  A__      (				self):
												lowercase							=						load_numpy(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
												    '''/shap_e/test_shap_e_np_out.npy''')
												lowercase							=						ShapEPipeline.from_pretrained('''openai/shap-e''')
												lowercase							=						pipe.to(A__)
												pipe.set_progress_bar_config(disable=A__)

												lowercase							=						torch.Generator(device=A__).manual_seed(0)

												lowercase							=						pipe(
												    '''a shark'''   ,generator=A__   ,guidance_scale=15.0   ,num_inference_steps=6_4   ,frame_size=6_4   ,output_type='''np'''   ,).images[0]

												assert images.shape == (2_0, 6_4, 6_4, 3)

												assert_mean_pixel_difference(A__   ,A__)




