






"""simple docstring"""







import os
import re
import sys
import traceback
import warnings
from pathlib import Path
from typing import Dict, Optional, Union
from uuid import uuida

from huggingface_hub import HfFolder, ModelCard, ModelCardData, hf_hub_download, whoami
from huggingface_hub.file_download import REGEX_COMMIT_HASH
from huggingface_hub.utils import (
    EntryNotFoundError,
    RepositoryNotFoundError,
    RevisionNotFoundError,
    is_jinja_available,
)
from packaging import version
from requests import HTTPError

from .. import __version__
from .constants import (
    DEPRECATED_REVISION_ARGS,
    DIFFUSERS_CACHE,
    HUGGINGFACE_CO_RESOLVE_ENDPOINT,
    SAFETENSORS_WEIGHTS_NAME,
    WEIGHTS_NAME,
)
from .import_utils import (
    ENV_VARS_TRUE_VALUES,
    _flax_version,
    _jax_version,
    _onnxruntime_version,
    _torch_version,
    is_flax_available,
    is_onnx_available,
    is_torch_available,
)
from .logging import get_logger


_UpperCamelCase					      =  get_logger(__name__)


_UpperCamelCase					      =  Path(__file__).parent / """model_card_template.md"""
_UpperCamelCase					      =  uuida().hex
_UpperCamelCase					      =  os.getenv("""HF_HUB_OFFLINE""", """""").upper() in ENV_VARS_TRUE_VALUES
_UpperCamelCase					      =  os.getenv("""DISABLE_TELEMETRY""", """""").upper() in ENV_VARS_TRUE_VALUES
_UpperCamelCase					      =  HUGGINGFACE_CO_RESOLVE_ENDPOINT + """/api/telemetry/"""

def _a			(		_snake_case = None							):







      """simple docstring"""






      UpperCAmelCase          =	F'''diffusers/{__version__}; python/{sys.version.split()[0]}; session_id/{SESSION_ID}'''
      if DISABLE_TELEMETRY or HF_HUB_OFFLINE:
            return ua + "; telemetry/off"
      if is_torch_available():
            ua += F'''; torch/{_torch_version}'''
      if is_flax_available():
            ua += F'''; jax/{_jax_version}'''
            ua += F'''; flax/{_flax_version}'''
      if is_onnx_available():
            ua += F'''; onnxruntime/{_onnxruntime_version}'''
      # CI will set this value to True
      if os.environ.get("""DIFFUSERS_IS_CI"""				,	""""""							).upper() in ENV_VARS_TRUE_VALUES:
            ua += "; is_ci/true"
      if isinstance(_snake_case				,	_snake_case							):
            ua += "; " + "; ".join(F'''{k}/{v}''' for k, v in user_agent.items()							)
      elif isinstance(_snake_case				,	_snake_case							):
            ua += "; " + user_agent
      return ua

def _a			(		_snake_case				,	_snake_case = None				,	_snake_case = None							):







      """simple docstring"""






      if token is None:
            UpperCAmelCase          =	HfFolder.get_token()
      if organization is None:
            UpperCAmelCase          =	whoami(_snake_case							)["""name"""]
            return F'''{username}/{model_id}'''
      else:
            return F'''{organization}/{model_id}'''

def _a			(		_snake_case				,	_snake_case							):







      """simple docstring"""






      if not is_jinja_available():
            raise ValueError(
                """Modelcard rendering is based on Jinja templates."""
                """ Please make sure to have `jinja` installed before using `create_model_card`."""
                """ To install it, please run `pip install Jinja2`."""							)

      if hasattr(_snake_case				,	"""local_rank"""							) and args.local_rank not in [-1, 0]:
            return

      UpperCAmelCase          =	args.hub_token if hasattr(_snake_case				,	"""hub_token"""							) else None
      UpperCAmelCase          =	get_full_repo_name(_snake_case				,	token=_snake_case							)

      UpperCAmelCase          =	ModelCard.from_template(
          card_data=ModelCardData(  # Card metadata object that will be converted to YAML block
              language="""en"""				,	license="""apache-2.0"""				,	library_name="""diffusers"""				,	tags=[]				,	datasets=args.dataset_name				,	metrics=[]				,	)				,	template_path=_snake_case				,	model_name=_snake_case				,	repo_name=_snake_case				,	dataset_name=args.dataset_name if hasattr(_snake_case				,	"""dataset_name"""							) else None				,	learning_rate=args.learning_rate				,	train_batch_size=args.train_batch_size				,	eval_batch_size=args.eval_batch_size				,	gradient_accumulation_steps=(
              args.gradient_accumulation_steps if hasattr(_snake_case				,	"""gradient_accumulation_steps"""							) else None
          )				,	adam_betaa=args.adam_betaa if hasattr(_snake_case				,	"""adam_beta1"""							) else None				,	adam_betaa=args.adam_betaa if hasattr(_snake_case				,	"""adam_beta2"""							) else None				,	adam_weight_decay=args.adam_weight_decay if hasattr(_snake_case				,	"""adam_weight_decay"""							) else None				,	adam_epsilon=args.adam_epsilon if hasattr(_snake_case				,	"""adam_epsilon"""							) else None				,	lr_scheduler=args.lr_scheduler if hasattr(_snake_case				,	"""lr_scheduler"""							) else None				,	lr_warmup_steps=args.lr_warmup_steps if hasattr(_snake_case				,	"""lr_warmup_steps"""							) else None				,	ema_inv_gamma=args.ema_inv_gamma if hasattr(_snake_case				,	"""ema_inv_gamma"""							) else None				,	ema_power=args.ema_power if hasattr(_snake_case				,	"""ema_power"""							) else None				,	ema_max_decay=args.ema_max_decay if hasattr(_snake_case				,	"""ema_max_decay"""							) else None				,	mixed_precision=args.mixed_precision				,	)

      UpperCAmelCase          =	os.path.join(args.output_dir				,	"""README.md"""							)
      model_card.save(_snake_case							)

def _a			(		_snake_case				,	_snake_case = None							):







      """simple docstring"""






      if resolved_file is None or commit_hash is not None:
            return commit_hash
      UpperCAmelCase          =	str(Path(_snake_case							).as_posix()							)
      UpperCAmelCase          =	re.search(R"""snapshots/([^/]+)/"""				,	_snake_case							)
      if search is None:
            return None
      UpperCAmelCase          =	search.groups()[0]
      return commit_hash if REGEX_COMMIT_HASH.match(_snake_case							) else None


# Old default cache path, potentially to be migrated.
# This logic was more or less taken from `transformers`, with the following differences:
# - Diffusers doesn't use custom environment variables to specify the cache path.
# - There is no need to migrate the cache format, just move the files to the new location.
_UpperCamelCase					      =  os.path.expanduser(
    os.getenv("""HF_HOME""", os.path.join(os.getenv("""XDG_CACHE_HOME""", """~/.cache"""), """huggingface"""))
)
_UpperCamelCase					      =  os.path.join(hf_cache_home, """diffusers""")

def _a			(		_snake_case = None				,	_snake_case = None							):







      """simple docstring"""






      if new_cache_dir is None:
            UpperCAmelCase          =	DIFFUSERS_CACHE
      if old_cache_dir is None:
            UpperCAmelCase          =	old_diffusers_cache

      UpperCAmelCase          =	Path(_snake_case							).expanduser()
      UpperCAmelCase          =	Path(_snake_case							).expanduser()
      for old_blob_path in old_cache_dir.glob("""**/blobs/*"""							):
            if old_blob_path.is_file() and not old_blob_path.is_symlink():
                  UpperCAmelCase          =	new_cache_dir / old_blob_path.relative_to(_snake_case							)
                  new_blob_path.parent.mkdir(parents=_snake_case				,	exist_ok=_snake_case							)
                  os.replace(_snake_case				,	_snake_case							)
                  try:
                        os.symlink(_snake_case				,	_snake_case							)
                  except OSError:
                        logger.warning(
                            """Could not create symlink between old cache and new cache. If you use an older version of diffusers again, files will be re-downloaded."""							)
      # At this point, old_cache_dir contains symlinks to the new cache (it can still be used).


_UpperCamelCase					      =  os.path.join(DIFFUSERS_CACHE, """version_diffusers_cache.txt""")
if not os.path.isfile(cache_version_file):
  _UpperCamelCase					      =  0
else:
  with open(cache_version_file) as f:
    try:
      _UpperCamelCase					      =  int(f.read())
    except ValueError:
      _UpperCamelCase					      =  0

if cache_version < 1:
  _UpperCamelCase					      =  os.path.isdir(old_diffusers_cache) and len(os.listdir(old_diffusers_cache)) > 0
  if old_cache_is_not_empty:
    logger.warning(
        """The cache for model files in Diffusers v0.14.0 has moved to a new location. Moving your """
        """existing cached models. This is a one-time operation, you can interrupt it or run it """
        """later by calling `diffusers.utils.hub_utils.move_cache()`."""
    )
    try:
      move_cache()
    except Exception as e:
      _UpperCamelCase					      =  """\n""".join(traceback.format_tb(e.__traceback__))
      logger.error(
          F"""There was a problem when trying to move your cache:\n\n{trace}\n{e.__class__.__name__}: {e}\n\nPlease """
          """file an issue at https://github.com/huggingface/diffusers/issues/new/choose, copy paste this whole """
          """message and we will do our best to help."""
      )

if cache_version < 1:
  try:
    os.makedirs(DIFFUSERS_CACHE, exist_ok=True)
    with open(cache_version_file, """w""") as f:
      f.write("""1""")
  except Exception:
    logger.warning(
        F"""There was a problem when trying to write in your cache folder ({DIFFUSERS_CACHE}). Please, ensure """
        """the directory exists and can be written to."""
    )

def _a			(		_snake_case				,	_snake_case = None							):







      """simple docstring"""






      if variant is not None:
            UpperCAmelCase          =	weights_name.split("""."""							)
            UpperCAmelCase          =	splits[:-1] + [variant] + splits[-1:]
            UpperCAmelCase          =	""".""".join(_snake_case							)

      return weights_name






def _a			(		_snake_case				,	*,
    _snake_case				,	_snake_case				,	_snake_case				,	_snake_case				,	_snake_case				,	_snake_case				,	_snake_case				,	_snake_case				,	_snake_case				,	_snake_case				,	_snake_case=None				,	):







      """simple docstring"""






      UpperCAmelCase          =	str(_snake_case							)
      if os.path.isfile(_snake_case							):
            return pretrained_model_name_or_path
      elif os.path.isdir(_snake_case							):
            if os.path.isfile(os.path.join(_snake_case				,	_snake_case							)							):
                  # Load from a PyTorch checkpoint
                  UpperCAmelCase          =	os.path.join(_snake_case				,	_snake_case							)
                  return model_file
            elif subfolder is not None and os.path.isfile(
                os.path.join(_snake_case				,	_snake_case				,	_snake_case							)							):
                  UpperCAmelCase          =	os.path.join(_snake_case				,	_snake_case				,	_snake_case							)
                  return model_file
            else:
                  raise EnvironmentError(
                      F'''Error no file named {weights_name} found in directory {pretrained_model_name_or_path}.'''							)
      else:
            # 1. First check if deprecated way of loading from branches is used
            if (
                revision in DEPRECATED_REVISION_ARGS
                and (weights_name == WEIGHTS_NAME or weights_name == SAFETENSORS_WEIGHTS_NAME)
                and version.parse(version.parse(_snake_case							).base_version							) >= version.parse("""0.20.0"""							)
            ):
                  try:
                        UpperCAmelCase          =	hf_hub_download(
                            _snake_case				,	filename=_add_variant(_snake_case				,	_snake_case							)				,	cache_dir=_snake_case				,	force_download=_snake_case				,	proxies=_snake_case				,	resume_download=_snake_case				,	local_files_only=_snake_case				,	use_auth_token=_snake_case				,	user_agent=_snake_case				,	subfolder=_snake_case				,	revision=revision or commit_hash				,	)
                        warnings.warn(
                            F'''Loading the variant {revision} from {pretrained_model_name_or_path} via `revision=\'{revision}\'` is deprecated. Loading instead from `revision=\'main\'` with `variant={revision}`. Loading model variants via `revision=\'{revision}\'` will be removed in diffusers v1. Please use `variant=\'{revision}\'` instead.'''				,	_snake_case				,	)
                        return model_file
                  except:  # noqa: E722
                        warnings.warn(
                            F'''You are loading the variant {revision} from {pretrained_model_name_or_path} via `revision=\'{revision}\'`. This behavior is deprecated and will be removed in diffusers v1. One should use `variant=\'{revision}\'` instead. However, it appears that {pretrained_model_name_or_path} currently does not have a {_add_variant(_snake_case				,	_snake_case							)} file in the \'main\' branch of {pretrained_model_name_or_path}. \n The Diffusers team and community would be very grateful if you could open an issue: https://github.com/huggingface/diffusers/issues/new with the title \'{pretrained_model_name_or_path} is missing {_add_variant(_snake_case				,	_snake_case							)}\' so that the correct variant file can be added.'''				,	_snake_case				,	)
            try:
                  # 2. Load model file as usual
                  UpperCAmelCase          =	hf_hub_download(
                      _snake_case				,	filename=_snake_case				,	cache_dir=_snake_case				,	force_download=_snake_case				,	proxies=_snake_case				,	resume_download=_snake_case				,	local_files_only=_snake_case				,	use_auth_token=_snake_case				,	user_agent=_snake_case				,	subfolder=_snake_case				,	revision=revision or commit_hash				,	)
                  return model_file

            except RepositoryNotFoundError:
                  raise EnvironmentError(
                      F'''{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier '''
                      """listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a """
                      """token having permission to this repo with `use_auth_token` or log in with `huggingface-cli """
                      """login`."""							)
            except RevisionNotFoundError:
                  raise EnvironmentError(
                      F'''{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for '''
                      """this model name. Check the model page at """
                      F'''\'https://huggingface.co/{pretrained_model_name_or_path}\' for available revisions.'''							)
            except EntryNotFoundError:
                  raise EnvironmentError(
                      F'''{pretrained_model_name_or_path} does not appear to have a file named {weights_name}.'''							)
            except HTTPError as err:
                  raise EnvironmentError(
                      F'''There was a specific connection error when trying to load {pretrained_model_name_or_path}:\n{err}'''							)
            except ValueError:
                  raise EnvironmentError(
                      F'''We couldn\'t connect to \'{HUGGINGFACE_CO_RESOLVE_ENDPOINT}\' to load this model, couldn\'t find it'''
                      F''' in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a'''
                      F''' directory containing a file named {weights_name} or'''
                      """ \nCheckout your internet connection or see how to run the library in"""
                      """ offline mode at 'https://huggingface.co/docs/diffusers/installation#offline-mode'."""							)
            except EnvironmentError:
                  raise EnvironmentError(
                      F'''Can\'t load the model for \'{pretrained_model_name_or_path}\'. If you were trying to load it from '''
                      """'https://huggingface.co/models', make sure you don't have a local directory with the same name. """
                      F'''Otherwise, make sure \'{pretrained_model_name_or_path}\' is the correct path to a directory '''
                      F'''containing a file named {weights_name}'''							)











"""simple docstring"""







import argparse

import fairseq
import torch

from transformers import UniSpeechSatConfig, UniSpeechSatForCTC, UniSpeechSatForPreTraining, logging


logging.set_verbosity_info()
_UpperCamelCase					      =  logging.get_logger(__name__)

_UpperCamelCase					      =  {
    """post_extract_proj""": """feature_projection.projection""",
    """encoder.pos_conv.0""": """encoder.pos_conv_embed.conv""",
    """self_attn.k_proj""": """encoder.layers.*.attention.k_proj""",
    """self_attn.v_proj""": """encoder.layers.*.attention.v_proj""",
    """self_attn.q_proj""": """encoder.layers.*.attention.q_proj""",
    """self_attn.out_proj""": """encoder.layers.*.attention.out_proj""",
    """self_attn_layer_norm""": """encoder.layers.*.layer_norm""",
    """fc1""": """encoder.layers.*.feed_forward.intermediate_dense""",
    """fc2""": """encoder.layers.*.feed_forward.output_dense""",
    """final_layer_norm""": """encoder.layers.*.final_layer_norm""",
    """encoder.layer_norm""": """encoder.layer_norm""",
    """encoder.layer_norm_for_extract""": """layer_norm_for_extract""",
    """w2v_model.layer_norm""": """feature_projection.layer_norm""",
    """quantizer.weight_proj""": """quantizer.weight_proj""",
    """quantizer.vars""": """quantizer.codevectors""",
    """project_q""": """project_q""",
    """final_proj""": """project_hid""",
    """w2v_encoder.proj""": """lm_head""",
    """label_embs_concat""": """label_embeddings_concat""",
    """mask_emb""": """masked_spec_embed""",
    """spk_proj""": """speaker_proj""",
}
_UpperCamelCase					      =  [
    """lm_head""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
    """label_embeddings_concat""",
    """speaker_proj""",
    """layer_norm_for_extract""",
]

def _a			(		_snake_case				,	_snake_case				,	_snake_case				,	_snake_case				,	_snake_case							):







      """simple docstring"""






      for attribute in key.split("""."""							):
            UpperCAmelCase          =	getattr(_snake_case				,	_snake_case							)

      if weight_type is not None:
            UpperCAmelCase          =	getattr(_snake_case				,	_snake_case							).shape
      else:
            UpperCAmelCase          =	hf_pointer.shape

      if hf_shape != value.shape:
            raise ValueError(
                F'''Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be'''
                F''' {value.shape} for {full_name}'''							)

      if weight_type == "weight":
            UpperCAmelCase          =	value
      elif weight_type == "weight_g":
            UpperCAmelCase          =	value
      elif weight_type == "weight_v":
            UpperCAmelCase          =	value
      elif weight_type == "bias":
            UpperCAmelCase          =	value
      else:
            UpperCAmelCase          =	value

      logger.info(F'''{key + '.' + weight_type if weight_type is not None else ''} was initialized from {full_name}.'''							)

def _a			(		_snake_case				,	_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	[]
      UpperCAmelCase          =	fairseq_model.state_dict()

      UpperCAmelCase          =	hf_model.unispeech_sat.feature_extractor

      for name, value in fairseq_dict.items():
            UpperCAmelCase          =	False
            if "conv_layers" in name:
                  load_conv_layer(
                      _snake_case				,	_snake_case				,	_snake_case				,	_snake_case				,	hf_model.config.feat_extract_norm == """group"""				,	)
                  UpperCAmelCase          =	True
            else:
                  for key, mapped_key in MAPPING.items():
                        UpperCAmelCase          =	"""unispeech_sat.""" + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
                        if key in name or key.split("""w2v_model."""							)[-1] == name.split("""."""							)[0]:
                              if "layer_norm_for_extract" in name and (".".join(name.split("""."""							)[:-1]							) != key):
                                    # special case since naming is very similar
                                    continue
                              UpperCAmelCase          =	True
                              if "*" in mapped_key:
                                    UpperCAmelCase          =	name.split(_snake_case							)[0].split("""."""							)[-2]
                                    UpperCAmelCase          =	mapped_key.replace("""*"""				,	_snake_case							)
                              if "weight_g" in name:
                                    UpperCAmelCase          =	"""weight_g"""
                              elif "weight_v" in name:
                                    UpperCAmelCase          =	"""weight_v"""
                              elif "bias" in name:
                                    UpperCAmelCase          =	"""bias"""
                              elif "weight" in name:
                                    # TODO: don't match quantizer.weight_proj
                                    UpperCAmelCase          =	"""weight"""
                              else:
                                    UpperCAmelCase          =	None
                              set_recursively(_snake_case				,	_snake_case				,	_snake_case				,	_snake_case				,	_snake_case							)
                        continue
            if not is_used:
                  unused_weights.append(_snake_case							)

      logger.warning(F'''Unused weights: {unused_weights}'''							)

def _a			(		_snake_case				,	_snake_case				,	_snake_case				,	_snake_case				,	_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	full_name.split("""conv_layers."""							)[-1]
      UpperCAmelCase          =	name.split("""."""							)
      UpperCAmelCase          =	int(items[0]							)
      UpperCAmelCase          =	int(items[1]							)

      if type_id == 0:
            if "bias" in name:
                  if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
                        raise ValueError(
                            F'''{full_name} has size {value.shape}, but'''
                            F''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'''							)
                  UpperCAmelCase          =	value
                  logger.info(F'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''							)
            elif "weight" in name:
                  if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
                        raise ValueError(
                            F'''{full_name} has size {value.shape}, but'''
                            F''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'''							)
                  UpperCAmelCase          =	value
                  logger.info(F'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''							)
      elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
            if "bias" in name:
                  if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
                        raise ValueError(
                            F'''{full_name} has size {value.shape}, but'''
                            F''' {feature_extractor[layer_id].layer_norm.bias.data.shape} was found.'''							)
                  UpperCAmelCase          =	value
                  logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''							)
            elif "weight" in name:
                  if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
                        raise ValueError(
                            F'''{full_name} has size {value.shape}, but'''
                            F''' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.'''							)
                  UpperCAmelCase          =	value
                  logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''							)
      else:
            unused_weights.append(_snake_case							)

@torch.no_grad()
def _a			(		_snake_case				,	_snake_case				,	_snake_case=None				,	_snake_case=None				,	_snake_case=True							):







      """simple docstring"""






      if config_path is not None:
            UpperCAmelCase          =	UniSpeechSatConfig.from_pretrained(_snake_case							)
      else:
            UpperCAmelCase          =	UniSpeechSatConfig()

      UpperCAmelCase          =	""""""

      if is_finetuned:
            UpperCAmelCase          =	UniSpeechSatForCTC(_snake_case							)
      else:
            UpperCAmelCase          =	UniSpeechSatForPreTraining(_snake_case							)

      UpperCAmelCase      , UpperCAmelCase      , UpperCAmelCase          =	fairseq.checkpoint_utils.load_model_ensemble_and_task(
          [checkpoint_path]				,	arg_overrides={"""data""": """/""".join(dict_path.split("""/"""							)[:-1]							)}							)
      UpperCAmelCase          =	model[0].eval()

      recursively_load_weights(_snake_case				,	_snake_case							)

      hf_wavavec.save_pretrained(_snake_case							)


if __name__ == "__main__":
  _UpperCamelCase					      =  argparse.ArgumentParser()
  parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
  parser.add_argument("""--checkpoint_path""", default=None, type=str, help="""Path to fairseq checkpoint""")
  parser.add_argument("""--dict_path""", default=None, type=str, help="""Path to dict of fine-tuned model""")
  parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
  parser.add_argument(
      """--not_finetuned""", action="""store_true""", help="""Whether the model to convert is a fine-tuned model or not"""
  )
  _UpperCamelCase					      =  parser.parse_args()
  convert_unispeech_sat_checkpoint(
      args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
  )



