




import argparse
import json
import os

import torch
from transformers.file_utils import has_file

from diffusers import UNetaDConditionModel, UNetaDModel


UpperCamelCase__ :   str						=							False
UpperCamelCase__ :   List[str]						=							True
UpperCamelCase__ :   Optional[Any]						=							False


if __name__ == "__main__":
						UpperCamelCase__ :   Union[str, Any]						=							argparse.ArgumentParser()

						parser.add_argument(
						    '''--repo_path''',
						    default=None,
						    type=str,
						    required=True,
						    help='''The config json file corresponding to the architecture.''',
						)

						parser.add_argument('''--dump_path''', default=None, type=str, required=True, help='''Path to the output model.''')

						UpperCamelCase__ :   Optional[Any]						=							parser.parse_args()

						UpperCamelCase__ :   Tuple						=							{
						    '''image_size''': '''sample_size''',
						    '''num_res_blocks''': '''layers_per_block''',
						    '''block_channels''': '''block_out_channels''',
						    '''down_blocks''': '''down_block_types''',
						    '''up_blocks''': '''up_block_types''',
						    '''downscale_freq_shift''': '''freq_shift''',
						    '''resnet_num_groups''': '''norm_num_groups''',
						    '''resnet_act_fn''': '''act_fn''',
						    '''resnet_eps''': '''norm_eps''',
						    '''num_head_channels''': '''attention_head_dim''',
						}

						UpperCamelCase__ :   Optional[int]						=							{
						    '''time_steps''': '''time_proj''',
						    '''mid''': '''mid_block''',
						    '''downsample_blocks''': '''down_blocks''',
						    '''upsample_blocks''': '''up_blocks''',
						}

						UpperCamelCase__ :   Tuple						=							'''''' if has_file(args.repo_path, '''config.json''') else '''unet'''

						with open(os.path.join(args.repo_path, subfolder, '''config.json'''), '''r''', encoding='''utf-8''') as reader:
												UpperCamelCase__ :   Optional[int]						=							reader.read()
												UpperCamelCase__ :   Union[str, Any]						=							json.loads(text)

						if do_only_config:
												for key in config_parameters_to_change.keys():
																		config.pop(key, None)

						if has_file(args.repo_path, '''config.json'''):
												UpperCamelCase__ :   Any						=							UNetaDModel(**config)
						else:
												UpperCamelCase__ :   Tuple						=							UNetaDConditionModel if '''ldm-text2im-large-256''' in args.repo_path else UNetaDModel
												UpperCamelCase__ :   List[str]						=							class_name(**config)

						if do_only_config:
												model.save_config(os.path.join(args.repo_path, subfolder))

						UpperCamelCase__ :   Union[str, Any]						=							dict(model.config)

						if do_only_renaming:
												for key, value in config_parameters_to_change.items():
																		if key in config:
																								UpperCamelCase__ :   List[str]						=							config[key]
																								del config[key]

												UpperCamelCase__ :   Dict						=							[k.replace('''UNetRes''', '''''') for k in config['''down_block_types''']]
												UpperCamelCase__ :   Dict						=							[k.replace('''UNetRes''', '''''') for k in config['''up_block_types''']]

						if do_only_weights:
												UpperCamelCase__ :   int						=							torch.load(os.path.join(args.repo_path, subfolder, '''diffusion_pytorch_model.bin'''))

												UpperCamelCase__ :   List[Any]						=							{}
												for param_key, param_value in state_dict.items():
																		if param_key.endswith('''.op.bias''') or param_key.endswith('''.op.weight'''):
																								continue
																		UpperCamelCase__ :   str						=							False
																		for key, new_key in key_parameters_to_change.items():
																								if not has_changed and param_key.split('''.''')[0] == key:
																														UpperCamelCase__ :   Tuple						=							param_value
																														UpperCamelCase__ :   Dict						=							True
																		if not has_changed:
																								UpperCamelCase__ :   str						=							param_value

												model.load_state_dict(new_state_dict)
												model.save_pretrained(os.path.join(args.repo_path, subfolder))









from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCamelCase__ :   Optional[Any]						=							{
    '''configuration_whisper''': ['''WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''WhisperConfig''', '''WhisperOnnxConfig'''],
    '''feature_extraction_whisper''': ['''WhisperFeatureExtractor'''],
    '''processing_whisper''': ['''WhisperProcessor'''],
    '''tokenization_whisper''': ['''WhisperTokenizer'''],
}

try:
						if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						UpperCamelCase__ :   str						=							['''WhisperTokenizerFast''']

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						UpperCamelCase__ :   List[Any]						=							[
						    '''WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST''',
						    '''WhisperForConditionalGeneration''',
						    '''WhisperModel''',
						    '''WhisperPreTrainedModel''',
						    '''WhisperForAudioClassification''',
						]

try:
						if not is_tf_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						UpperCamelCase__ :   Tuple						=							[
						    '''TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST''',
						    '''TFWhisperForConditionalGeneration''',
						    '''TFWhisperModel''',
						    '''TFWhisperPreTrainedModel''',
						]

try:
						if not is_flax_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						UpperCamelCase__ :   List[str]						=							[
						    '''FlaxWhisperForConditionalGeneration''',
						    '''FlaxWhisperModel''',
						    '''FlaxWhisperPreTrainedModel''',
						    '''FlaxWhisperForAudioClassification''',
						]


if TYPE_CHECKING:
						from .configuration_whisper import WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP, WhisperConfig, WhisperOnnxConfig
						from .feature_extraction_whisper import WhisperFeatureExtractor
						from .processing_whisper import WhisperProcessor
						from .tokenization_whisper import WhisperTokenizer

						try:
												if not is_tokenizers_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_whisper_fast import WhisperTokenizerFast

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_whisper import (
												    WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
												    WhisperForAudioClassification,
												    WhisperForConditionalGeneration,
												    WhisperModel,
												    WhisperPreTrainedModel,
												)

						try:
												if not is_tf_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_tf_whisper import (
												    TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
												    TFWhisperForConditionalGeneration,
												    TFWhisperModel,
												    TFWhisperPreTrainedModel,
												)

						try:
												if not is_flax_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_flax_whisper import (
												    FlaxWhisperForAudioClassification,
												    FlaxWhisperForConditionalGeneration,
												    FlaxWhisperModel,
												    FlaxWhisperPreTrainedModel,
												)

else:
						import sys

						UpperCamelCase__ :   Optional[Any]						=							_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



