




import importlib
import inspect
import json
import os
import re
import shutil
import sys
from pathlib import Path
from typing import Dict, Optional, Union
from urllib import request

from huggingface_hub import HfFolder, cached_download, hf_hub_download, model_info
from packaging import version

from .. import __version__
from . import DIFFUSERS_DYNAMIC_MODULE_NAME, HF_MODULES_CACHE, logging


_lowerCamelCase     :		Dict								=       (
    '''https://raw.githubusercontent.com/huggingface/diffusers/{revision}/examples/community/{pipeline}.py'''
)


_lowerCamelCase     :		Union[str, Any]								=       logging.get_logger(__name__)  # pylint: disable=invalid-name







def SCREAMING_SNAKE_CASE    (	)       ->	int:



       """simple docstring"""

       A__							=   '''https://pypi.org/pypi/diffusers/json'''
       A__							=   json.loads(request.urlopen(A_						).read()						)['''releases'''].keys()
       return sorted(A_			,    key=lambda lowercase_						: version.Version(A_						)						)







def SCREAMING_SNAKE_CASE    (	)       ->	Optional[int]:



       """simple docstring"""

       if HF_MODULES_CACHE in sys.path:
              return

       sys.path.append(A_						)
       os.makedirs(A_			,    exist_ok=A_						)
       A__							=   Path(A_						) / '''__init__.py'''
       if not init_path.exists():
              init_path.touch()







def SCREAMING_SNAKE_CASE    (	lowercase_						)       ->	List[Any]:



       """simple docstring"""

       init_hf_modules()
       A__							=   Path(A_						) / name
       # If the parent module does not exist yet, recursively create it.
       if not dynamic_module_path.parent.exists():
              create_dynamic_module(dynamic_module_path.parent						)
       os.makedirs(A_			,    exist_ok=A_						)
       A__							=   dynamic_module_path / '''__init__.py'''
       if not init_path.exists():
              init_path.touch()







def SCREAMING_SNAKE_CASE    (	lowercase_						)       ->	Optional[Any]:



       """simple docstring"""

       with open(A_			,    '''r'''			,    encoding='''utf-8'''						) as f:
              A__							=   f.read()

       # Imports of the form `import .xxx`
       A__							=   re.findall('''^\s*import\s+\.(\S+)\s*$'''			,    A_			,    flags=re.MULTILINE						)
       # Imports of the form `from .xxx import yyy`
       relative_imports += re.findall('''^\s*from\s+\.(\S+)\s+import'''			,    A_			,    flags=re.MULTILINE						)
       # Unique-ify
       return list(set(A_						)						)







def SCREAMING_SNAKE_CASE    (	lowercase_						)       ->	List[Any]:



       """simple docstring"""

       A__							=   False
       A__							=   [module_file]
       A__							=   []

       # Let's recurse through all relative imports
       while not no_change:
              A__							=   []
              for f in files_to_check:
                     new_imports.extend(get_relative_imports(A_						)						)

              A__							=   Path(A_						).parent
              A__							=   [str(module_path / m						) for m in new_imports]
              A__							=   [f for f in new_import_files if f not in all_relative_imports]
              A__							=   [f"""{f}.py""" for f in new_import_files]

              A__							=   len(A_						) == 0
              all_relative_imports.extend(A_						)

       return all_relative_imports







def SCREAMING_SNAKE_CASE    (	lowercase_						)       ->	Tuple:



       """simple docstring"""

       with open(A_			,    '''r'''			,    encoding='''utf-8'''						) as f:
              A__							=   f.read()

       # Imports of the form `import xxx`
       A__							=   re.findall('''^\s*import\s+(\S+)\s*$'''			,    A_			,    flags=re.MULTILINE						)
       # Imports of the form `from xxx import yyy`
       imports += re.findall('''^\s*from\s+(\S+)\s+import'''			,    A_			,    flags=re.MULTILINE						)
       # Only keep the top-level module
       A__							=   [imp.split('''.'''						)[0] for imp in imports if not imp.startswith('''.'''						)]

       # Unique-ify and test we got them all
       A__							=   list(set(A_						)						)
       A__							=   []
       for imp in imports:
              try:
                     importlib.import_module(A_						)
              except ImportError:
                     missing_packages.append(A_						)

       if len(A_						) > 0:
              raise ImportError(
                  '''This modeling file requires the following packages that were not found in your environment: '''
                  f"""{", ".join(A_						)}. Run `pip install {" ".join(A_						)}`"""						)

       return get_relative_imports(A_						)







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_						)       ->	Union[str, Any]:



       """simple docstring"""

       A__							=   module_path.replace(os.path.sep			,    '''.'''						)
       A__							=   importlib.import_module(A_						)

       if class_name is None:
              return find_pipeline_class(A_						)
       return getattr(A_			,    A_						)







def SCREAMING_SNAKE_CASE    (	lowercase_						)       ->	int:



       """simple docstring"""

       from ..pipelines import DiffusionPipeline

       A__							=   dict(inspect.getmembers(A_			,    inspect.isclass						)						)

       A__							=   None
       for cls_name, cls in cls_members.items():
              if (
                  cls_name != DiffusionPipeline.__name__
                  and issubclass(cls			,    A_						)
                  and cls.__module__.split('''.'''						)[0] != "diffusers"
              ):
                     if pipeline_class is not None:
                            raise ValueError(
                                f"""Multiple classes that inherit from {DiffusionPipeline.__name__} have been found:"""
                                f""" {pipeline_class.__name__}, and {cls_name}. Please make sure to define only one in"""
                                f""" {loaded_module}."""						)
                     A__							=   cls

       return pipeline_class







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_			,    lowercase_ = None			,    lowercase_ = False			,    lowercase_ = False			,    lowercase_ = None			,    lowercase_ = None			,    lowercase_ = None			,    lowercase_ = False			,    )       ->	List[Any]:



       """simple docstring"""

       A__							=   str(A_						)

       A__							=   os.path.join(A_			,    A_						)

       if os.path.isfile(A_						):
              A__							=   module_file_or_url
              A__							=   '''local'''
       elif pretrained_model_name_or_path.count('''/'''						) == 0:
              A__							=   get_diffusers_versions()
              # cut ".dev0"
              A__							=   '''v''' + '''.'''.join(__version__.split('''.'''						)[:3]						)

              # retrieve github version that matches
              if revision is None:
                     A__							=   latest_version if latest_version[1:] in available_versions else '''main'''
                     logger.info(f"""Defaulting to latest_version: {revision}."""						)
              elif revision in available_versions:
                     A__							=   f"""v{revision}"""
              elif revision == "main":
                     A__							=   revision
              else:
                     raise ValueError(
                         f"""`custom_revision`: {revision} does not exist. Please make sure to choose one of"""
                         f""" {", ".join(available_versions + ["main"]						)}."""						)

              # community pipeline on GitHub
              A__							=   COMMUNITY_PIPELINES_URL.format(revision=A_			,    pipeline=A_						)
              try:
                     A__							=   cached_download(
                         A_			,    cache_dir=A_			,    force_download=A_			,    proxies=A_			,    resume_download=A_			,    local_files_only=A_			,    use_auth_token=A_			,    )
                     A__							=   '''git'''
                     A__							=   pretrained_model_name_or_path + '''.py'''
              except EnvironmentError:
                     logger.error(f"""Could not locate the {module_file} inside {pretrained_model_name_or_path}."""						)
                     raise
       else:
              try:
                     # Load from URL or cache if already cached
                     A__							=   hf_hub_download(
                         A_			,    A_			,    cache_dir=A_			,    force_download=A_			,    proxies=A_			,    resume_download=A_			,    local_files_only=A_			,    use_auth_token=A_			,    )
                     A__							=   os.path.join('''local'''			,    '''--'''.join(pretrained_model_name_or_path.split('''/'''						)						)						)
              except EnvironmentError:
                     logger.error(f"""Could not locate the {module_file} inside {pretrained_model_name_or_path}."""						)
                     raise

    # Check we have all the requirements in our environment
       A__							=   check_imports(A_						)

       # Now we move the module inside our cached dynamic modules.
       A__							=   DIFFUSERS_DYNAMIC_MODULE_NAME + os.path.sep + submodule
       create_dynamic_module(A_						)
       A__							=   Path(A_						) / full_submodule
       if submodule == "local" or submodule == "git":
              # We always copy local files (we could hash the file to see if there was a change, and give them the name of
              # that hash, to only copy when there is a modification but it seems overkill for now).
              # The only reason we do the copy is to avoid putting too many folders in sys.path.
              shutil.copy(A_			,    submodule_path / module_file						)
              for module_needed in modules_needed:
                     A__							=   f"""{module_needed}.py"""
                     shutil.copy(os.path.join(A_			,    A_						)			,    submodule_path / module_needed						)
       else:
              # Get the commit hash
              # TODO: we will get this info in the etag soon, so retrieve it from there and not here.
              if isinstance(A_			,    A_						):
                     A__							=   use_auth_token
              elif use_auth_token is True:
                     A__							=   HfFolder.get_token()
              else:
                     A__							=   None

              A__							=   model_info(A_			,    revision=A_			,    token=A_						).sha

              # The module file will end up being placed in a subfolder with the git hash of the repo. This way we get the
              # benefit of versioning.
              A__							=   submodule_path / commit_hash
              A__							=   full_submodule + os.path.sep + commit_hash
              create_dynamic_module(A_						)

              if not (submodule_path / module_file).exists():
                     shutil.copy(A_			,    submodule_path / module_file						)
              # Make sure we also have every file with relative
              for module_needed in modules_needed:
                     if not (submodule_path / module_needed).exists():
                            get_cached_module_file(
                                A_			,    f"""{module_needed}.py"""			,    cache_dir=A_			,    force_download=A_			,    resume_download=A_			,    proxies=A_			,    use_auth_token=A_			,    revision=A_			,    local_files_only=A_			,    )
       return os.path.join(A_			,    A_						)







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_			,    lowercase_ = None			,    lowercase_ = None			,    lowercase_ = False			,    lowercase_ = False			,    lowercase_ = None			,    lowercase_ = None			,    lowercase_ = None			,    lowercase_ = False			,    **lowercase_			,    )       ->	Union[str, Any]:



       """simple docstring"""

       A__							=   get_cached_module_file(
           A_			,    A_			,    cache_dir=A_			,    force_download=A_			,    resume_download=A_			,    proxies=A_			,    use_auth_token=A_			,    revision=A_			,    local_files_only=A_			,    )
       return get_class_in_module(A_			,    final_module.replace('''.py'''			,    ''''''						)						)








import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor







def SCREAMING_SNAKE_CASE    (	lowercase_						)       ->	str:



       """simple docstring"""

       if "cls_token" in name:
              A__							=   name.replace('''cls_token'''			,    '''vit.embeddings.cls_token'''						)
       if "mask_token" in name:
              A__							=   name.replace('''mask_token'''			,    '''decoder.mask_token'''						)
       if "decoder_pos_embed" in name:
              A__							=   name.replace('''decoder_pos_embed'''			,    '''decoder.decoder_pos_embed'''						)
       if "pos_embed" in name and "decoder" not in name:
              A__							=   name.replace('''pos_embed'''			,    '''vit.embeddings.position_embeddings'''						)
       if "patch_embed.proj" in name:
              A__							=   name.replace('''patch_embed.proj'''			,    '''vit.embeddings.patch_embeddings.projection'''						)
       if "patch_embed.norm" in name:
              A__							=   name.replace('''patch_embed.norm'''			,    '''vit.embeddings.norm'''						)
       if "decoder_blocks" in name:
              A__							=   name.replace('''decoder_blocks'''			,    '''decoder.decoder_layers'''						)
       if "blocks" in name:
              A__							=   name.replace('''blocks'''			,    '''vit.encoder.layer'''						)
       if "attn.proj" in name:
              A__							=   name.replace('''attn.proj'''			,    '''attention.output.dense'''						)
       if "attn" in name:
              A__							=   name.replace('''attn'''			,    '''attention.self'''						)
       if "norm1" in name:
              A__							=   name.replace('''norm1'''			,    '''layernorm_before'''						)
       if "norm2" in name:
              A__							=   name.replace('''norm2'''			,    '''layernorm_after'''						)
       if "mlp.fc1" in name:
              A__							=   name.replace('''mlp.fc1'''			,    '''intermediate.dense'''						)
       if "mlp.fc2" in name:
              A__							=   name.replace('''mlp.fc2'''			,    '''output.dense'''						)
       if "decoder_embed" in name:
              A__							=   name.replace('''decoder_embed'''			,    '''decoder.decoder_embed'''						)
       if "decoder_norm" in name:
              A__							=   name.replace('''decoder_norm'''			,    '''decoder.decoder_norm'''						)
       if "decoder_pred" in name:
              A__							=   name.replace('''decoder_pred'''			,    '''decoder.decoder_pred'''						)
       if "norm.weight" in name and "decoder" not in name:
              A__							=   name.replace('''norm.weight'''			,    '''vit.layernorm.weight'''						)
       if "norm.bias" in name and "decoder" not in name:
              A__							=   name.replace('''norm.bias'''			,    '''vit.layernorm.bias'''						)

       return name







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_						)       ->	str:



       """simple docstring"""

       for key in orig_state_dict.copy().keys():
              A__							=   orig_state_dict.pop(lowercase_						)

              if "qkv" in key:
                     A__							=   key.split('''.'''						)
                     A__							=   int(key_split[1]						)
                     if "decoder_blocks" in key:
                            A__							=   config.decoder_hidden_size
                            A__							=   '''decoder.decoder_layers.'''
                            if "weight" in key:
                                   A__							=   val[:dim, :]
                                   A__							=   val[dim : dim * 2, :]
                                   A__							=   val[-dim:, :]
                            elif "bias" in key:
                                   A__							=   val[:dim]
                                   A__							=   val[dim : dim * 2]
                                   A__							=   val[-dim:]
                     else:
                            A__							=   config.hidden_size
                            A__							=   '''vit.encoder.layer.'''
                            if "weight" in key:
                                   A__							=   val[:dim, :]
                                   A__							=   val[dim : dim * 2, :]
                                   A__							=   val[-dim:, :]
                            elif "bias" in key:
                                   A__							=   val[:dim]
                                   A__							=   val[dim : dim * 2]
                                   A__							=   val[-dim:]

              else:
                     A__							=   val

       return orig_state_dict







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_						)       ->	Tuple:



       """simple docstring"""

       A__							=   ViTMAEConfig()
       if "large" in checkpoint_url:
              A__							=   1_024
              A__							=   4_096
              A__							=   24
              A__							=   16
       elif "huge" in checkpoint_url:
              A__							=   14
              A__							=   1_280
              A__							=   5_120
              A__							=   32
              A__							=   16

       A__							=   ViTMAEForPreTraining(lowercase_						)

       A__							=   torch.hub.load_state_dict_from_url(lowercase_			,    map_location='''cpu'''						)['''model''']

       A__							=   ViTMAEImageProcessor(size=config.image_size						)

       A__							=   convert_state_dict(lowercase_			,    lowercase_						)

       model.load_state_dict(lowercase_						)
       model.eval()

       A__							=   '''https://user-images.githubusercontent.com/11435359/147738734-196fd92f-9260-48d5-ba7e-bf103d29364d.jpg'''

       A__							=   Image.open(requests.get(lowercase_			,    stream=lowercase_						).raw						)
       A__							=   ViTMAEImageProcessor(size=config.image_size						)
       A__							=   image_processor(images=lowercase_			,    return_tensors='''pt'''						)

       # forward pass
       torch.manual_seed(2						)
       A__							=   model(**lowercase_						)
       A__							=   outputs.logits

       if "large" in checkpoint_url:
              A__							=   torch.tensor(
                  [[-0.73_09, -0.71_28, -1.01_69], [-1.01_61, -0.90_58, -1.18_78], [-1.04_78, -0.94_11, -1.19_11]]						)
       elif "huge" in checkpoint_url:
              A__							=   torch.tensor(
                  [[-1.15_99, -0.91_99, -1.22_21], [-1.19_52, -0.92_69, -1.23_07], [-1.21_43, -0.93_37, -1.22_62]]						)
       else:
              A__							=   torch.tensor(
                  [[-0.91_92, -0.84_81, -1.12_59], [-1.13_49, -1.00_34, -1.25_99], [-1.17_57, -1.04_29, -1.27_26]]						)

       # verify logits
       assert torch.allclose(logits[0, :3, :3]			,    lowercase_			,    atol=1E-4						)

       print(f"""Saving model to {pytorch_dump_folder_path}"""						)
       model.save_pretrained(lowercase_						)

       print(f"""Saving image processor to {pytorch_dump_folder_path}"""						)
       image_processor.save_pretrained(lowercase_						)


if __name__ == "__main__":
  _lowerCamelCase     :		Any								=       argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--checkpoint_url""",
      default="""https://dl.fbaipublicfiles.com/mae/visualize/mae_visualize_vit_base.pth""",
      type=str,
      help="""URL of the checkpoint you'd like to convert.""",
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
  )

  _lowerCamelCase     :		Optional[int]								=       parser.parse_args()
  convert_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)


