






import importlib
import json
import os
from collections import OrderedDict
from typing import Dict, Optional, Union

# Build the list of all feature extractors
from ...configuration_utils import PretrainedConfig
from ...dynamic_module_utils import get_class_from_dynamic_module, resolve_trust_remote_code
from ...feature_extraction_utils import FeatureExtractionMixin
from ...utils import CONFIG_NAME, FEATURE_EXTRACTOR_NAME, get_file_from_repo, logging
from .auto_factory import _LazyAutoMapping
from .configuration_auto import (
    CONFIG_MAPPING_NAMES,
    AutoConfig,
    model_type_to_module_name,
    replace_list_option_in_docstrings,
)


__lowerCamelCase	:		int										=   logging.get_logger(__name__)

__lowerCamelCase	:		str										=   OrderedDict(
    [
        ("""audio-spectrogram-transformer""", """ASTFeatureExtractor"""),
        ("""beit""", """BeitFeatureExtractor"""),
        ("""chinese_clip""", """ChineseCLIPFeatureExtractor"""),
        ("""clap""", """ClapFeatureExtractor"""),
        ("""clip""", """CLIPFeatureExtractor"""),
        ("""clipseg""", """ViTFeatureExtractor"""),
        ("""conditional_detr""", """ConditionalDetrFeatureExtractor"""),
        ("""convnext""", """ConvNextFeatureExtractor"""),
        ("""cvt""", """ConvNextFeatureExtractor"""),
        ("""data2vec-audio""", """Wav2Vec2FeatureExtractor"""),
        ("""data2vec-vision""", """BeitFeatureExtractor"""),
        ("""deformable_detr""", """DeformableDetrFeatureExtractor"""),
        ("""deit""", """DeiTFeatureExtractor"""),
        ("""detr""", """DetrFeatureExtractor"""),
        ("""dinat""", """ViTFeatureExtractor"""),
        ("""donut-swin""", """DonutFeatureExtractor"""),
        ("""dpt""", """DPTFeatureExtractor"""),
        ("""encodec""", """EncodecFeatureExtractor"""),
        ("""flava""", """FlavaFeatureExtractor"""),
        ("""glpn""", """GLPNFeatureExtractor"""),
        ("""groupvit""", """CLIPFeatureExtractor"""),
        ("""hubert""", """Wav2Vec2FeatureExtractor"""),
        ("""imagegpt""", """ImageGPTFeatureExtractor"""),
        ("""layoutlmv2""", """LayoutLMv2FeatureExtractor"""),
        ("""layoutlmv3""", """LayoutLMv3FeatureExtractor"""),
        ("""levit""", """LevitFeatureExtractor"""),
        ("""maskformer""", """MaskFormerFeatureExtractor"""),
        ("""mctct""", """MCTCTFeatureExtractor"""),
        ("""mobilenet_v1""", """MobileNetV1FeatureExtractor"""),
        ("""mobilenet_v2""", """MobileNetV2FeatureExtractor"""),
        ("""mobilevit""", """MobileViTFeatureExtractor"""),
        ("""nat""", """ViTFeatureExtractor"""),
        ("""owlvit""", """OwlViTFeatureExtractor"""),
        ("""perceiver""", """PerceiverFeatureExtractor"""),
        ("""poolformer""", """PoolFormerFeatureExtractor"""),
        ("""regnet""", """ConvNextFeatureExtractor"""),
        ("""resnet""", """ConvNextFeatureExtractor"""),
        ("""segformer""", """SegformerFeatureExtractor"""),
        ("""sew""", """Wav2Vec2FeatureExtractor"""),
        ("""sew-d""", """Wav2Vec2FeatureExtractor"""),
        ("""speech_to_text""", """Speech2TextFeatureExtractor"""),
        ("""speecht5""", """SpeechT5FeatureExtractor"""),
        ("""swiftformer""", """ViTFeatureExtractor"""),
        ("""swin""", """ViTFeatureExtractor"""),
        ("""swinv2""", """ViTFeatureExtractor"""),
        ("""table-transformer""", """DetrFeatureExtractor"""),
        ("""timesformer""", """VideoMAEFeatureExtractor"""),
        ("""tvlt""", """TvltFeatureExtractor"""),
        ("""unispeech""", """Wav2Vec2FeatureExtractor"""),
        ("""unispeech-sat""", """Wav2Vec2FeatureExtractor"""),
        ("""van""", """ConvNextFeatureExtractor"""),
        ("""videomae""", """VideoMAEFeatureExtractor"""),
        ("""vilt""", """ViltFeatureExtractor"""),
        ("""vit""", """ViTFeatureExtractor"""),
        ("""vit_mae""", """ViTFeatureExtractor"""),
        ("""vit_msn""", """ViTFeatureExtractor"""),
        ("""wav2vec2""", """Wav2Vec2FeatureExtractor"""),
        ("""wav2vec2-conformer""", """Wav2Vec2FeatureExtractor"""),
        ("""wavlm""", """Wav2Vec2FeatureExtractor"""),
        ("""whisper""", """WhisperFeatureExtractor"""),
        ("""xclip""", """CLIPFeatureExtractor"""),
        ("""yolos""", """YolosFeatureExtractor"""),
    ]
)

__lowerCamelCase	:		Dict										=   _LazyAutoMapping(CONFIG_MAPPING_NAMES, FEATURE_EXTRACTOR_MAPPING_NAMES)







def 			A__							(      _a  :					Optional[Any]	):





	'''simple docstring'''







	for module_name, extractors in FEATURE_EXTRACTOR_MAPPING_NAMES.items():
		if class_name in extractors:
			snake_case__  :    List[Any]							       =model_type_to_module_name(_UpperCAmelCase	)

			snake_case__  :    Optional[int]							       =importlib.import_module(f".{module_name}"		,  """transformers.models"""	)
			try:
				return getattr(_UpperCAmelCase		,  _UpperCAmelCase	)
			except AttributeError:
				continue

	for _, extractor in FEATURE_EXTRACTOR_MAPPING._extra_content.items():
		if getattr(_UpperCAmelCase		,  """__name__"""		,  _UpperCAmelCase	) == class_name:
			return extractor

    # We did not fine the class, but maybe it's because a dep is missing. In that case, the class will be in the main
    # init and we return the proper dummy to get an appropriate error message.
	snake_case__  :    Union[str, Any]							       =importlib.import_module("""transformers"""	)
	if hasattr(_UpperCAmelCase		,  _UpperCAmelCase	):
		return getattr(_UpperCAmelCase		,  _UpperCAmelCase	)

	return None







def 			A__							(      _a  :					List[Any]		,  _a  :					str = None		,  _a  :					Tuple = False		,  _a  :					Any = False		,  _a  :					int = None		,  _a  :					Dict = None		,  _a  :					List[str] = None		,  _a  :					Optional[int] = False		,  **_a  :					Tuple		,  ):





	'''simple docstring'''







	snake_case__  :    str							       =get_file_from_repo(
	    _UpperCAmelCase		,  _UpperCAmelCase		,  cache_dir=_UpperCAmelCase		,  force_download=_UpperCAmelCase		,  resume_download=_UpperCAmelCase		,  proxies=_UpperCAmelCase		,  use_auth_token=_UpperCAmelCase		,  revision=_UpperCAmelCase		,  local_files_only=_UpperCAmelCase		,  )
	if resolved_config_file is None:
		logger.info(
		    """Could not locate the feature extractor configuration file, will try to use the model config instead."""	)
		return {}

	with open(_UpperCAmelCase		,  encoding="""utf-8"""	) as reader:
		return json.load(_UpperCAmelCase	)







class  _lowercase   :
						def __init__(     self   ):
							raise EnvironmentError(
							    """AutoFeatureExtractor is designed to be instantiated """
							    """using the `AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path)` method."""   )
						@classmethod
						@replace_list_option_in_docstrings(__SCREAMING_SNAKE_CASE   )
						def 						lowercase__					(     cls				,			a				,			**a   ):
							snake_case__  :    Dict							       =kwargs.pop("""config"""				,			__SCREAMING_SNAKE_CASE   )
							snake_case__  :    List[str]							       =kwargs.pop("""trust_remote_code"""				,			__SCREAMING_SNAKE_CASE   )
							snake_case__  :    List[str]							       =True

							snake_case__       ,   snake_case__  :    Tuple							       =FeatureExtractionMixin.get_feature_extractor_dict(__SCREAMING_SNAKE_CASE				,			**__SCREAMING_SNAKE_CASE   )
							snake_case__  :    Optional[int]							       =config_dict.get("""feature_extractor_type"""				,			__SCREAMING_SNAKE_CASE   )
							snake_case__  :    Any							       =None
							if "AutoFeatureExtractor" in config_dict.get("""auto_map"""				,			{}   ):
								snake_case__  :    Tuple							       =config_dict["""auto_map"""]["""AutoFeatureExtractor"""]

							# If we don't find the feature extractor class in the feature extractor config, let's try the model config.
							if feature_extractor_class is None and feature_extractor_auto_map is None:
								if not isinstance(__SCREAMING_SNAKE_CASE				,			__SCREAMING_SNAKE_CASE   ):
									snake_case__  :    Optional[Any]							       =AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE				,			**__SCREAMING_SNAKE_CASE   )
								# It could be in `config.feature_extractor_type``
								snake_case__  :    Any							       =getattr(__SCREAMING_SNAKE_CASE				,			"""feature_extractor_type"""				,			__SCREAMING_SNAKE_CASE   )
								if hasattr(__SCREAMING_SNAKE_CASE				,			"""auto_map"""   ) and "AutoFeatureExtractor" in config.auto_map:
									snake_case__  :    Optional[int]							       =config.auto_map["""AutoFeatureExtractor"""]

							if feature_extractor_class is not None:
								snake_case__  :    Optional[int]							       =feature_extractor_class_from_name(__SCREAMING_SNAKE_CASE   )

							snake_case__  :    Optional[Any]							       =feature_extractor_auto_map is not None
							snake_case__  :    Optional[Any]							       =feature_extractor_class is not None or type(__SCREAMING_SNAKE_CASE   ) in FEATURE_EXTRACTOR_MAPPING
							snake_case__  :    Tuple							       =resolve_trust_remote_code(
							    __SCREAMING_SNAKE_CASE				,			__SCREAMING_SNAKE_CASE				,			__SCREAMING_SNAKE_CASE				,			__SCREAMING_SNAKE_CASE   )

							if has_remote_code and trust_remote_code:
								snake_case__  :    Optional[Any]							       =get_class_from_dynamic_module(
								    __SCREAMING_SNAKE_CASE				,			__SCREAMING_SNAKE_CASE				,			**__SCREAMING_SNAKE_CASE   )
								snake_case__  :    Optional[int]							       =kwargs.pop("""code_revision"""				,			__SCREAMING_SNAKE_CASE   )
								if os.path.isdir(__SCREAMING_SNAKE_CASE   ):
									feature_extractor_class.register_for_auto_class()
								return feature_extractor_class.from_dict(__SCREAMING_SNAKE_CASE				,			**__SCREAMING_SNAKE_CASE   )
							elif feature_extractor_class is not None:
								return feature_extractor_class.from_dict(__SCREAMING_SNAKE_CASE				,			**__SCREAMING_SNAKE_CASE   )
							# Last try: we use the FEATURE_EXTRACTOR_MAPPING.
							elif type(__SCREAMING_SNAKE_CASE   ) in FEATURE_EXTRACTOR_MAPPING:
								snake_case__  :    str							       =FEATURE_EXTRACTOR_MAPPING[type(__SCREAMING_SNAKE_CASE   )]
								return feature_extractor_class.from_dict(__SCREAMING_SNAKE_CASE				,			**__SCREAMING_SNAKE_CASE   )

							raise ValueError(
							    F"Unrecognized feature extractor in {pretrained_model_name_or_path}. Should have a "
							    F"`feature_extractor_type` key in its {FEATURE_EXTRACTOR_NAME} of {CONFIG_NAME}, or one of the following "
							    F"`model_type` keys in its {CONFIG_NAME}: {', '.join(c for c in FEATURE_EXTRACTOR_MAPPING_NAMES.keys()   )}"   )



						@staticmethod
						def 						lowercase__					(     a				,			a   ):
							FEATURE_EXTRACTOR_MAPPING.register(__SCREAMING_SNAKE_CASE				,			__SCREAMING_SNAKE_CASE   )













import argparse

import torch

from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_from_original_stable_diffusion_ckpt


if __name__ == "__main__":
       __lowerCamelCase	:		List[str]										=   argparse.ArgumentParser()

       parser.add_argument(
           """--checkpoint_path""", default=None, type=str, required=True, help="""Path to the checkpoint to convert."""
       )
       # !wget https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml
       parser.add_argument(
           """--original_config_file""",
           default=None,
           type=str,
           help="""The YAML config file corresponding to the original architecture.""",
       )
       parser.add_argument(
           """--num_in_channels""",
           default=None,
           type=int,
           help="""The number of input channels. If `None` number of input channels will be automatically inferred.""",
       )
       parser.add_argument(
           """--scheduler_type""",
           default="""pndm""",
           type=str,
           help="""Type of scheduler to use. Should be one of ['pndm', 'lms', 'ddim', 'euler', 'euler-ancestral', 'dpm']""",
       )
       parser.add_argument(
           """--pipeline_type""",
           default=None,
           type=str,
           help=(
               """The pipeline type. One of 'FrozenOpenCLIPEmbedder', 'FrozenCLIPEmbedder', 'PaintByExample'"""
               """. If `None` pipeline will be automatically inferred."""
           ),
       )
       parser.add_argument(
           """--image_size""",
           default=None,
           type=int,
           help=(
               """The image size that the model was trained on. Use 512 for Stable Diffusion v1.X and Stable Siffusion v2"""
               """ Base. Use 768 for Stable Diffusion v2."""
           ),
       )
       parser.add_argument(
           """--prediction_type""",
           default=None,
           type=str,
           help=(
               """The prediction type that the model was trained on. Use 'epsilon' for Stable Diffusion v1.X and Stable"""
               """ Diffusion v2 Base. Use 'v_prediction' for Stable Diffusion v2."""
           ),
       )
       parser.add_argument(
           """--extract_ema""",
           action="""store_true""",
           help=(
               """Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights"""
               """ or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield"""
               """ higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning."""
           ),
       )
       parser.add_argument(
           """--upcast_attention""",
           action="""store_true""",
           help=(
               """Whether the attention computation should always be upcasted. This is necessary when running stable"""
               """ diffusion 2.1."""
           ),
       )
       parser.add_argument(
           """--from_safetensors""",
           action="""store_true""",
           help="""If `--checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch.""",
       )
       parser.add_argument(
           """--to_safetensors""",
           action="""store_true""",
           help="""Whether to store pipeline in safetensors format or not.""",
       )
       parser.add_argument("""--dump_path""", default=None, type=str, required=True, help="""Path to the output model.""")
       parser.add_argument("""--device""", type=str, help="""Device to use (e.g. cpu, cuda:0, cuda:1, etc.)""")
       parser.add_argument(
           """--stable_unclip""",
           type=str,
           default=None,
           required=False,
           help="""Set if this is a stable unCLIP model. One of 'txt2img' or 'img2img'.""",
       )
       parser.add_argument(
           """--stable_unclip_prior""",
           type=str,
           default=None,
           required=False,
           help="""Set if this is a stable unCLIP txt2img model. Selects which prior to use. If `--stable_unclip` is set to `txt2img`, the karlo prior (https://huggingface.co/kakaobrain/karlo-v1-alpha/tree/main/prior) is selected by default.""",
       )
       parser.add_argument(
           """--clip_stats_path""",
           type=str,
           help="""Path to the clip stats file. Only required if the stable unclip model's config specifies `model.params.noise_aug_config.params.clip_stats_path`.""",
           required=False,
       )
       parser.add_argument(
           """--controlnet""", action="""store_true""", default=None, help="""Set flag if this is a controlnet checkpoint."""
       )
       parser.add_argument("""--half""", action="""store_true""", help="""Save weights in half precision.""")
       parser.add_argument(
           """--vae_path""",
           type=str,
           default=None,
           required=False,
           help="""Set to a path, hub id to an already converted vae to not convert it again.""",
       )
       __lowerCamelCase	:		Optional[Any]										=   parser.parse_args()

       __lowerCamelCase	:		Optional[int]										=   download_from_original_stable_diffusion_ckpt(
           checkpoint_path=args.checkpoint_path,
           original_config_file=args.original_config_file,
           image_size=args.image_size,
           prediction_type=args.prediction_type,
           model_type=args.pipeline_type,
           extract_ema=args.extract_ema,
           scheduler_type=args.scheduler_type,
           num_in_channels=args.num_in_channels,
           upcast_attention=args.upcast_attention,
           from_safetensors=args.from_safetensors,
           device=args.device,
           stable_unclip=args.stable_unclip,
           stable_unclip_prior=args.stable_unclip_prior,
           clip_stats_path=args.clip_stats_path,
           controlnet=args.controlnet,
           vae_path=args.vae_path,
       )

       if args.half:
              pipe.to(torch_dtype=torch.floataa)

       if args.controlnet:
              # only save the controlnet model
              pipe.controlnet.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)
       else:
              pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)





