




from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
					from PIL import Image

					from ..image_utils import load_image

if is_tf_available():
					from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_VISION_2_SEQ_MAPPING

if is_torch_available():
					import torch

					from ..models.auto.modeling_auto import MODEL_FOR_VISION_2_SEQ_MAPPING

__A        =    logging.get_logger(__name__)







@add_end_docstrings(__magic_name__ )
class     __lowerCAmelCase (       __magic_name__ ):

			"""simple docstring"""







			def __init__(      self			,						*lowerCamelCase__			,						**lowerCamelCase__							)    ->      str:



							'''simple docstring'''
							super().__init__(*lowerCamelCase__			,						**lowerCamelCase__							)
							requires_backends(self			,						'vision'							)
							self.check_model_type(
							    TF_MODEL_FOR_VISION_2_SEQ_MAPPING if self.framework == 'tf' else MODEL_FOR_VISION_2_SEQ_MAPPING							)


			def       lowercase_    (      self			,						lowerCamelCase__=None			,						lowerCamelCase__=None			,						lowerCamelCase__=None							)    ->      Optional[int]:



							'''simple docstring'''
							__lowerCamelCase		       =    {}
							__lowerCamelCase		       =    {}

							if prompt is not None:
											__lowerCamelCase		       =    prompt

							if generate_kwargs is not None:
											__lowerCamelCase		       =    generate_kwargs
							if max_new_tokens is not None:
											if "generate_kwargs" not in forward_kwargs:
															__lowerCamelCase		       =    {}
											if "max_new_tokens" in forward_kwargs["generate_kwargs"]:
															raise ValueError(
															    '\'max_new_tokens\' is defined twice, once in \'generate_kwargs\' and once as a direct parameter,'
															    ' please use only one'							)
											__lowerCamelCase		       =    max_new_tokens
							return preprocess_params, forward_kwargs, {}


			def __call__(      self			,						lowerCamelCase__			,						**lowerCamelCase__							)    ->      int:



							'''simple docstring'''
							return super().__call__(lowerCamelCase__			,						**lowerCamelCase__							)


			def       lowercase_    (      self			,						lowerCamelCase__			,						lowerCamelCase__=None							)    ->      Dict:



							'''simple docstring'''
							__lowerCamelCase		       =    load_image(lowerCamelCase__							)

							if prompt is not None:
											if not isinstance(lowerCamelCase__			,						lowerCamelCase__							):
															raise ValueError(
															    f"""Received an invalid text input, got - {type(lowerCamelCase__							)} - but expected a single string. """
															    'Note also that one single text can be provided for conditional image to text generation.'							)

											__lowerCamelCase		       =    self.model.config.model_type

											if model_type == "git":
															__lowerCamelCase		       =    self.image_processor(images=lowerCamelCase__			,						return_tensors=self.framework							)
															__lowerCamelCase		       =    self.tokenizer(text=lowerCamelCase__			,						add_special_tokens=lowerCamelCase__							).input_ids
															__lowerCamelCase		       =    [self.tokenizer.cls_token_id] + input_ids
															__lowerCamelCase		       =    torch.tensor(lowerCamelCase__							).unsqueeze(0							)
															model_inputs.update({'input_ids': input_ids}							)

											elif model_type == "pix2struct":
															__lowerCamelCase		       =    self.image_processor(images=lowerCamelCase__			,						header_text=lowerCamelCase__			,						return_tensors=self.framework							)

											elif model_type != "vision-encoder-decoder":
															# vision-encoder-decoder does not support conditional generation
															__lowerCamelCase		       =    self.image_processor(images=lowerCamelCase__			,						return_tensors=self.framework							)
															__lowerCamelCase		       =    self.tokenizer(lowerCamelCase__			,						return_tensors=self.framework							)
															model_inputs.update(lowerCamelCase__							)

											else:
															raise ValueError(f"""Model type {model_type} does not support conditional text generation"""							)

							else:
											__lowerCamelCase		       =    self.image_processor(images=lowerCamelCase__			,						return_tensors=self.framework							)

							if self.model.config.model_type == "git" and prompt is None:
											__lowerCamelCase		       =    None

							return model_inputs


			def       lowercase_    (      self			,						lowerCamelCase__			,						lowerCamelCase__=None							)    ->      List[str]:



							'''simple docstring'''
							# Git model sets `model_inputs["input_ids"] = None` in `preprocess` (when `prompt=None`). In batch model, the
							# pipeline will group them into a list of `None`, which fail `_forward`. Avoid this by checking it first.
							if (
							    "input_ids" in model_inputs
							    and isinstance(model_inputs['input_ids']			,						lowerCamelCase__							)
							    and all(x is None for x in model_inputs['input_ids']							)
							):
											__lowerCamelCase		       =    None

							if generate_kwargs is None:
											__lowerCamelCase		       =    {}
							# FIXME: We need to pop here due to a difference in how `generation.py` and `generation.tf_utils.py`
							#  parse inputs. In the Tensorflow version, `generate` raises an error if we don't use `input_ids` whereas
							#  the PyTorch version matches it with `self.model.main_input_name` or `self.model.encoder.main_input_name`
							#  in the `_prepare_model_inputs` method.
							__lowerCamelCase		       =    model_inputs.pop(self.model.main_input_name							)
							__lowerCamelCase		       =    self.model.generate(lowerCamelCase__			,						**lowerCamelCase__			,						**lowerCamelCase__							)
							return model_outputs






			def       lowercase_    (      self			,						lowerCamelCase__							)    ->      str:



							'''simple docstring'''
							__lowerCamelCase		       =    []
							for output_ids in model_outputs:
											__lowerCamelCase		       =    {
											    'generated_text': self.tokenizer.decode(
											        lowerCamelCase__			,						skip_special_tokens=lowerCamelCase__			,						)
											}
											records.append(lowerCamelCase__							)
							return records




"""simple docstring"""





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 	__a     (		__lowerCamelCase,		__lowerCamelCase  ):
   UpperCAmelCase_					:					str			     =    checkpoint

   UpperCAmelCase_					:					int			     =    {}

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

   UpperCAmelCase_					:					Any			     =    vae_state_dict["decoder.conv_in.weight"]
   UpperCAmelCase_					:					int			     =    vae_state_dict["decoder.conv_in.bias"]
   UpperCAmelCase_					:					Any			     =    vae_state_dict["decoder.conv_out.weight"]
   UpperCAmelCase_					:					Tuple			     =    vae_state_dict["decoder.conv_out.bias"]
   UpperCAmelCase_					:					List[Any]			     =    vae_state_dict["decoder.norm_out.weight"]
   UpperCAmelCase_					:					Tuple			     =    vae_state_dict["decoder.norm_out.bias"]

   UpperCAmelCase_					:					str			     =    vae_state_dict["quant_conv.weight"]
   UpperCAmelCase_					:					Optional[Any]			     =    vae_state_dict["quant_conv.bias"]
   UpperCAmelCase_					:					List[str]			     =    vae_state_dict["post_quant_conv.weight"]
   UpperCAmelCase_					:					List[Any]			     =    vae_state_dict["post_quant_conv.bias"]

   # Retrieves the keys for the encoder down blocks only
   UpperCAmelCase_					:					Optional[Any]			     =    len({".".join(layer.split("."  )[:3]  ) for layer in vae_state_dict if "encoder.down" in layer}  )
   UpperCAmelCase_					:					Optional[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
   UpperCAmelCase_					:					Dict			     =    len({".".join(layer.split("."  )[:3]  ) for layer in vae_state_dict if "decoder.up" in layer}  )
   UpperCAmelCase_					:					Optional[int]			     =    {
       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  ):
      UpperCAmelCase_					:					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:
         UpperCAmelCase_					:					Dict			     =    vae_state_dict.pop(
             f"""encoder.down.{i}.downsample.conv.weight"""  )
         UpperCAmelCase_					:					Dict			     =    vae_state_dict.pop(
             f"""encoder.down.{i}.downsample.conv.bias"""  )

      UpperCAmelCase_					:					List[str]			     =    renew_vae_resnet_paths(__lowerCamelCase  )
      UpperCAmelCase_					:					Optional[Any]			     =    {"old": f"""down.{i}.block""", "new": f"""down_blocks.{i}.resnets"""}
      assign_to_checkpoint(__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		additional_replacements=[meta_path],		config=__lowerCamelCase  )

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

      UpperCAmelCase_					:					List[Any]			     =    renew_vae_resnet_paths(__lowerCamelCase  )
      UpperCAmelCase_					:					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  )

   UpperCAmelCase_					:					Optional[int]			     =    [key for key in vae_state_dict if "encoder.mid.attn" in key]
   UpperCAmelCase_					:					Union[str, Any]			     =    renew_vae_attention_paths(__lowerCamelCase  )
   UpperCAmelCase_					:					int			     =    {"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  ):
      UpperCAmelCase_					:					Optional[Any]			     =    num_up_blocks - 1 - i
      UpperCAmelCase_					:					Any			     =    [
          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:
         UpperCAmelCase_					:					str			     =    vae_state_dict[
             f"""decoder.up.{block_id}.upsample.conv.weight"""
         ]
         UpperCAmelCase_					:					Optional[Any]			     =    vae_state_dict[
             f"""decoder.up.{block_id}.upsample.conv.bias"""
         ]

      UpperCAmelCase_					:					Dict			     =    renew_vae_resnet_paths(__lowerCamelCase  )
      UpperCAmelCase_					:					List[str]			     =    {"old": f"""up.{block_id}.block""", "new": f"""up_blocks.{i}.resnets"""}
      assign_to_checkpoint(__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		additional_replacements=[meta_path],		config=__lowerCamelCase  )

   UpperCAmelCase_					:					Optional[Any]			     =    [key for key in vae_state_dict if "decoder.mid.block" in key]
   UpperCAmelCase_					:					List[Any]			     =    2
   for i in range(1,		num_mid_res_blocks + 1  ):
      UpperCAmelCase_					:					str			     =    [key for key in mid_resnets if f"""decoder.mid.block_{i}""" in key]

      UpperCAmelCase_					:					Tuple			     =    renew_vae_resnet_paths(__lowerCamelCase  )
      UpperCAmelCase_					:					Tuple			     =    {"old": f"""mid.block_{i}""", "new": f"""mid_block.resnets.{i - 1}"""}
      assign_to_checkpoint(__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		additional_replacements=[meta_path],		config=__lowerCamelCase  )

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

   UpperCAmelCase_					:					Any			     =    OmegaConf.load(__lowerCamelCase  )
   UpperCAmelCase_					:					Union[str, Any]			     =    512
   UpperCAmelCase_					:					Optional[Any]			     =    "cuda" if torch.cuda.is_available() else "cpu"
   if checkpoint_path.endswith("safetensors"  ):
      from safetensors import safe_open

      UpperCAmelCase_					:					int			     =    {}
      with safe_open(__lowerCamelCase,		framework="pt",		device="cpu"  ) as f:
         for key in f.keys():
            UpperCAmelCase_					:					Tuple			     =    f.get_tensor(__lowerCamelCase  )
   else:
      UpperCAmelCase_					:					Any			     =    torch.load(__lowerCamelCase,		map_location=__lowerCamelCase  )["state_dict"]

   # Convert the VAE model.
   UpperCAmelCase_					:					Dict			     =    create_vae_diffusers_config(__lowerCamelCase,		image_size=__lowerCamelCase  )
   UpperCAmelCase_					:					Union[str, Any]			     =    custom_convert_ldm_vae_checkpoint(__lowerCamelCase,		__lowerCamelCase  )

   UpperCAmelCase_					:					int			     =    AutoencoderKL(**__lowerCamelCase  )
   vae.load_state_dict(__lowerCamelCase  )
   vae.save_pretrained(__lowerCamelCase  )


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)


