
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	) ->						List[Any]:
    __lowercase		:			str	    =       checkpoint

    __lowercase		:			Any	    =       {}

    __lowercase		:			Any	    =       vae_state_dict["""encoder.conv_in.weight"""]
    __lowercase		:			int	    =       vae_state_dict["""encoder.conv_in.bias"""]
    __lowercase		:			Optional[Any]	    =       vae_state_dict["""encoder.conv_out.weight"""]
    __lowercase		:			Dict	    =       vae_state_dict["""encoder.conv_out.bias"""]
    __lowercase		:			Optional[int]	    =       vae_state_dict["""encoder.norm_out.weight"""]
    __lowercase		:			str	    =       vae_state_dict["""encoder.norm_out.bias"""]

    __lowercase		:			str	    =       vae_state_dict["""decoder.conv_in.weight"""]
    __lowercase		:			Optional[int]	    =       vae_state_dict["""decoder.conv_in.bias"""]
    __lowercase		:			Optional[Any]	    =       vae_state_dict["""decoder.conv_out.weight"""]
    __lowercase		:			List[str]	    =       vae_state_dict["""decoder.conv_out.bias"""]
    __lowercase		:			Dict	    =       vae_state_dict["""decoder.norm_out.weight"""]
    __lowercase		:			Optional[int]	    =       vae_state_dict["""decoder.norm_out.bias"""]

    __lowercase		:			Tuple	    =       vae_state_dict["""quant_conv.weight"""]
    __lowercase		:			Union[str, Any]	    =       vae_state_dict["""quant_conv.bias"""]
    __lowercase		:			Union[str, Any]	    =       vae_state_dict["""post_quant_conv.weight"""]
    __lowercase		:			List[str]	    =       vae_state_dict["""post_quant_conv.bias"""]

    # Retrieves the keys for the encoder down blocks only
    __lowercase		:			Optional[Any]	    =       len({'''.'''.join(layer.split('''.'''	)[:3]	) for layer in vae_state_dict if '''encoder.down''' in layer}	)
    __lowercase		:			Any	    =       {
        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		:			Tuple	    =       len({'''.'''.join(layer.split('''.'''	)[:3]	) for layer in vae_state_dict if '''decoder.up''' in layer}	)
    __lowercase		:			Optional[Any]	    =       {
        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		:			Union[str, 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:
            __lowercase		:			Union[str, Any]	    =       vae_state_dict.pop(
                F'encoder.down.{i}.downsample.conv.weight'	)
            __lowercase		:			List[str]	    =       vae_state_dict.pop(
                F'encoder.down.{i}.downsample.conv.bias'	)

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

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

        __lowercase		:			Tuple	    =       renew_vae_resnet_paths(_lowerCamelCase	)
        __lowercase		:			List[Any]	    =       {"""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		:			List[str]	    =       [key for key in vae_state_dict if """encoder.mid.attn""" in key]
    __lowercase		:			Dict	    =       renew_vae_attention_paths(_lowerCamelCase	)
    __lowercase		:			List[Any]	    =       {"""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		:			str	    =       num_up_blocks - 1 - i
        __lowercase		:			str	    =       [
            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		:			Any	    =       vae_state_dict[
                F'decoder.up.{block_id}.upsample.conv.weight'
            ]
            __lowercase		:			Any	    =       vae_state_dict[
                F'decoder.up.{block_id}.upsample.conv.bias'
            ]

        __lowercase		:			List[Any]	    =       renew_vae_resnet_paths(_lowerCamelCase	)
        __lowercase		:			Tuple	    =       {"""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		:			str	    =       [key for key in vae_state_dict if """decoder.mid.block""" in key]
    __lowercase		:			Optional[int]	    =       2
    for i in range(1	,					num_mid_res_blocks + 1	):
        __lowercase		:			Tuple	    =       [key for key in mid_resnets if F'decoder.mid.block_{i}' in key]

        __lowercase		:			Dict	    =       renew_vae_resnet_paths(_lowerCamelCase	)
        __lowercase		:			Optional[int]	    =       {"""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		:			Optional[Any]	    =       [key for key in vae_state_dict if """decoder.mid.attn""" in key]
    __lowercase		:			Dict	    =       renew_vae_attention_paths(_lowerCamelCase	)
    __lowercase		:			Any	    =       {"""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	,					) ->						int:
    # Only support V1
    __lowercase		:			Any	    =       requests.get(
        ''' https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml'''	)
    __lowercase		:			Dict	    =       io.BytesIO(r.content	)

    __lowercase		:			Dict	    =       OmegaConf.load(_lowerCamelCase	)
    __lowercase		:			Optional[Any]	    =       512
    __lowercase		:			List[Any]	    =       """cuda""" if torch.cuda.is_available() else """cpu"""
    if checkpoint_path.endswith('''safetensors'''	):
        from safetensors import safe_open

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

    # Convert the VAE model.
    __lowercase		:			Optional[int]	    =       create_vae_diffusers_config(_lowerCamelCase	,					image_size=_lowerCamelCase	)
    __lowercase		:			Dict	    =       custom_convert_ldm_vae_checkpoint(_lowerCamelCase	,					_lowerCamelCase	)

    __lowercase		:			Optional[Any]	    =       AutoencoderKL(**_lowerCamelCase	)
    vae.load_state_dict(_lowerCamelCase	)
    vae.save_pretrained(_lowerCamelCase	)


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

 __lowerCAmelCase					:			List[Any]			      =		parser.parse_args()

 vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)


from .constants import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    TORCH_LAUNCH_PARAMS,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
)
from .dataclasses import (
    BnbQuantizationConfig,
    ComputeEnvironment,
    CustomDtype,
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    DynamoBackend,
    FPaRecipeKwargs,
    FullyShardedDataParallelPlugin,
    GradientAccumulationPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    KwargsHandler,
    LoggerType,
    MegatronLMPlugin,
    PrecisionType,
    ProjectConfiguration,
    RNGType,
    SageMakerDistributedType,
    TensorInformation,
    TorchDynamoPlugin,
)
from .environment import get_int_from_env, parse_choice_from_env, parse_flag_from_env
from .imports import (
    get_ccl_version,
    is_abit_bnb_available,
    is_abit_bnb_available,
    is_aim_available,
    is_bfaa_available,
    is_bnb_available,
    is_botoa_available,
    is_ccl_available,
    is_comet_ml_available,
    is_datasets_available,
    is_deepspeed_available,
    is_fpa_available,
    is_ipex_available,
    is_megatron_lm_available,
    is_mlflow_available,
    is_mps_available,
    is_npu_available,
    is_rich_available,
    is_safetensors_available,
    is_sagemaker_available,
    is_tensorboard_available,
    is_tpu_available,
    is_transformers_available,
    is_wandb_available,
    is_xpu_available,
)
from .modeling import (
    check_device_map,
    check_tied_parameters_in_config,
    check_tied_parameters_on_same_device,
    compute_module_sizes,
    convert_file_size_to_int,
    dtype_byte_size,
    find_tied_parameters,
    get_balanced_memory,
    get_max_layer_size,
    get_max_memory,
    get_mixed_precision_context_manager,
    id_tensor_storage,
    infer_auto_device_map,
    load_checkpoint_in_model,
    load_offloaded_weights,
    load_state_dict,
    named_module_tensors,
    retie_parameters,
    set_module_tensor_to_device,
    shard_checkpoint,
)
from .offload import (
    OffloadedWeightsLoader,
    PrefixedDataset,
    extract_submodules_state_dict,
    load_offloaded_weight,
    offload_state_dict,
    offload_weight,
    save_offload_index,
)
from .operations import (
    broadcast,
    broadcast_object_list,
    concatenate,
    convert_outputs_to_fpaa,
    convert_to_fpaa,
    find_batch_size,
    find_device,
    gather,
    gather_object,
    get_data_structure,
    honor_type,
    initialize_tensors,
    is_namedtuple,
    is_tensor_information,
    is_torch_tensor,
    listify,
    pad_across_processes,
    recursively_apply,
    reduce,
    send_to_device,
    slice_tensors,
)
from .versions import compare_versions, is_torch_version


if is_deepspeed_available():
 from .deepspeed import (
     DeepSpeedEngineWrapper,
     DeepSpeedOptimizerWrapper,
     DeepSpeedSchedulerWrapper,
     DummyOptim,
     DummyScheduler,
     HfDeepSpeedConfig,
 )

from .bnb import has_abit_bnb_layers, load_and_quantize_model
from .fsdp_utils import load_fsdp_model, load_fsdp_optimizer, save_fsdp_model, save_fsdp_optimizer
from .launch import (
    PrepareForLaunch,
    _filter_args,
    prepare_deepspeed_cmd_env,
    prepare_multi_gpu_env,
    prepare_sagemager_args_inputs,
    prepare_simple_launcher_cmd_env,
    prepare_tpu,
)
from .megatron_lm import (
    AbstractTrainStep,
    BertTrainStep,
    GPTTrainStep,
    MegatronEngine,
    MegatronLMDummyDataLoader,
    MegatronLMDummyScheduler,
    MegatronLMOptimizerWrapper,
    MegatronLMSchedulerWrapper,
    TaTrainStep,
    avg_losses_across_data_parallel_group,
    gather_across_data_parallel_groups,
)
from .megatron_lm import initialize as megatron_lm_initialize
from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader
from .megatron_lm import prepare_model as megatron_lm_prepare_model
from .megatron_lm import prepare_optimizer as megatron_lm_prepare_optimizer
from .megatron_lm import prepare_scheduler as megatron_lm_prepare_scheduler
from .memory import find_executable_batch_size, release_memory
from .other import (
    extract_model_from_parallel,
    get_pretty_name,
    is_port_in_use,
    merge_dicts,
    patch_environment,
    save,
    wait_for_everyone,
    write_basic_config,
)
from .random import set_seed, synchronize_rng_state, synchronize_rng_states
from .torch_xla import install_xla
from .tqdm import tqdm
from .transformer_engine import convert_model, has_transformer_engine_layers


