



"""simple docstring"""






import os

import torch

from ..logging import get_logger
from .constants import FSDP_PYTORCH_VERSION, MODEL_NAME, OPTIMIZER_NAME
from .versions import is_torch_version


if is_torch_version(">=", FSDP_PYTORCH_VERSION):
		import torch.distributed.checkpoint as dist_cp
		from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner
		from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
		from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
		from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType


_lowercase				:				int   		=		get_logger(__name__)






def   snake_case__     (	__lowerCamelCase       :					Any							,	__lowerCamelCase       :					Optional[Any]							,	__lowerCamelCase       :					Tuple							,	__lowerCamelCase       :					Union[str, Any]							,	__lowerCamelCase       :					Union[str, Any]=0			):
			"""simple docstring"""
			os.makedirs(__lowerCamelCase							,	exist_ok=__lowerCamelCase			)
			with FSDP.state_dict_type(
			    __lowerCamelCase							,	fsdp_plugin.state_dict_type							,	fsdp_plugin.state_dict_config							,	fsdp_plugin.optim_state_dict_config			):
						lowerCamelCase__	:	Union[str, Any]				=model.state_dict()
						if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
									lowerCamelCase__	:	List[str]				=f'''{MODEL_NAME}.bin''' if model_index == 0 else f'''{MODEL_NAME}_{model_index}.bin'''
									lowerCamelCase__	:	Tuple				=os.path.join(__lowerCamelCase							,	__lowerCamelCase			)
									if accelerator.process_index == 0:
												logger.info(f'''Saving model to {output_model_file}'''			)
												torch.save(__lowerCamelCase							,	__lowerCamelCase			)
												logger.info(f'''Model saved to {output_model_file}'''			)
						elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
									lowerCamelCase__	:	Union[str, Any]				=(
									    f'''{MODEL_NAME}_rank{accelerator.process_index}.bin'''
									    if model_index == 0
									    else f'''{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'''
									)
									lowerCamelCase__	:	str				=os.path.join(__lowerCamelCase							,	__lowerCamelCase			)
									logger.info(f'''Saving model to {output_model_file}'''			)
									torch.save(__lowerCamelCase							,	__lowerCamelCase			)
									logger.info(f'''Model saved to {output_model_file}'''			)
						elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
									lowerCamelCase__	:	str				=os.path.join(__lowerCamelCase							,	f'''{MODEL_NAME}_{model_index}'''			)
									os.makedirs(__lowerCamelCase							,	exist_ok=__lowerCamelCase			)
									logger.info(f'''Saving model to {ckpt_dir}'''			)
									lowerCamelCase__	:	List[Any]				={'''model''': state_dict}

									dist_cp.save_state_dict(
									    state_dict=__lowerCamelCase							,	storage_writer=dist_cp.FileSystemWriter(__lowerCamelCase			)							,	planner=DefaultSavePlanner()							,	)
									logger.info(f'''Model saved to {ckpt_dir}'''			)






def   snake_case__     (	__lowerCamelCase       :					Union[str, Any]							,	__lowerCamelCase       :					Any							,	__lowerCamelCase       :					Dict							,	__lowerCamelCase       :					Any							,	__lowerCamelCase       :					str=0			):
			"""simple docstring"""
			accelerator.wait_for_everyone()
			with FSDP.state_dict_type(
			    __lowerCamelCase							,	fsdp_plugin.state_dict_type							,	fsdp_plugin.state_dict_config							,	fsdp_plugin.optim_state_dict_config			):
						if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
									if type(__lowerCamelCase			) != FSDP and accelerator.process_index != 0:
												if not fsdp_plugin.sync_module_states:
															raise ValueError(
															    '''Set the `sync_module_states` flag to `True` so that model states are synced across processes when '''
															    '''initializing FSDP object'''			)
												return
									lowerCamelCase__	:	List[str]				=f'''{MODEL_NAME}.bin''' if model_index == 0 else f'''{MODEL_NAME}_{model_index}.bin'''
									lowerCamelCase__	:	Any				=os.path.join(__lowerCamelCase							,	__lowerCamelCase			)
									logger.info(f'''Loading model from {input_model_file}'''			)
									lowerCamelCase__	:	str				=torch.load(__lowerCamelCase			)
									logger.info(f'''Model loaded from {input_model_file}'''			)
						elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
									lowerCamelCase__	:	Tuple				=(
									    f'''{MODEL_NAME}_rank{accelerator.process_index}.bin'''
									    if model_index == 0
									    else f'''{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'''
									)
									lowerCamelCase__	:	Optional[Any]				=os.path.join(__lowerCamelCase							,	__lowerCamelCase			)
									logger.info(f'''Loading model from {input_model_file}'''			)
									lowerCamelCase__	:	int				=torch.load(__lowerCamelCase			)
									logger.info(f'''Model loaded from {input_model_file}'''			)
						elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
									lowerCamelCase__	:	List[Any]				=(
									    os.path.join(__lowerCamelCase							,	f'''{MODEL_NAME}_{model_index}'''			)
									    if f'''{MODEL_NAME}''' not in input_dir
									    else input_dir
									)
									logger.info(f'''Loading model from {ckpt_dir}'''			)
									lowerCamelCase__	:	int				={'''model''': model.state_dict()}
									dist_cp.load_state_dict(
									    state_dict=__lowerCamelCase							,	storage_reader=dist_cp.FileSystemReader(__lowerCamelCase			)							,	planner=DefaultLoadPlanner()							,	)
									lowerCamelCase__	:	Union[str, Any]				=state_dict['''model''']
									logger.info(f'''Model loaded from {ckpt_dir}'''			)
						model.load_state_dict(__lowerCamelCase			)






def   snake_case__     (	__lowerCamelCase       :					str							,	__lowerCamelCase       :					Optional[int]							,	__lowerCamelCase       :					List[str]							,	__lowerCamelCase       :					List[Any]							,	__lowerCamelCase       :					List[Any]							,	__lowerCamelCase       :					Dict=0			):
			"""simple docstring"""
			os.makedirs(__lowerCamelCase							,	exist_ok=__lowerCamelCase			)
			with FSDP.state_dict_type(
			    __lowerCamelCase							,	fsdp_plugin.state_dict_type							,	fsdp_plugin.state_dict_config							,	fsdp_plugin.optim_state_dict_config			):
						lowerCamelCase__	:	Optional[int]				=FSDP.optim_state_dict(__lowerCamelCase							,	__lowerCamelCase			)
						if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
									if accelerator.process_index == 0:
												lowerCamelCase__	:	int				=(
												    f'''{OPTIMIZER_NAME}.bin''' if optimizer_index == 0 else f'''{OPTIMIZER_NAME}_{optimizer_index}.bin'''
												)
												lowerCamelCase__	:	Optional[int]				=os.path.join(__lowerCamelCase							,	__lowerCamelCase			)
												logger.info(f'''Saving Optimizer state to {output_optimizer_file}'''			)
												torch.save(__lowerCamelCase							,	__lowerCamelCase			)
												logger.info(f'''Optimizer state saved in {output_optimizer_file}'''			)
						else:
									lowerCamelCase__	:	int				=os.path.join(__lowerCamelCase							,	f'''{OPTIMIZER_NAME}_{optimizer_index}'''			)
									os.makedirs(__lowerCamelCase							,	exist_ok=__lowerCamelCase			)
									logger.info(f'''Saving Optimizer state to {ckpt_dir}'''			)
									dist_cp.save_state_dict(
									    state_dict={'''optimizer''': optim_state}							,	storage_writer=dist_cp.FileSystemWriter(__lowerCamelCase			)							,	planner=DefaultSavePlanner()							,	)
									logger.info(f'''Optimizer state saved in {ckpt_dir}'''			)






def   snake_case__     (	__lowerCamelCase       :					Any							,	__lowerCamelCase       :					Dict							,	__lowerCamelCase       :					List[Any]							,	__lowerCamelCase       :					int							,	__lowerCamelCase       :					int							,	__lowerCamelCase       :					Optional[Any]=0			):
			"""simple docstring"""
			accelerator.wait_for_everyone()
			with FSDP.state_dict_type(
			    __lowerCamelCase							,	fsdp_plugin.state_dict_type							,	fsdp_plugin.state_dict_config							,	fsdp_plugin.optim_state_dict_config			):
						if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
									lowerCamelCase__	:	Tuple				=None
									# below check should work but currently it isn't working (mostly opytorch issue),
									# in the meantime disabling it at the cost of excess memory usage
									# if accelerator.process_index == 0 or not fsdp_plugin.optim_state_dict_config.rank0_only:
									lowerCamelCase__	:	Dict				=(
									    f'''{OPTIMIZER_NAME}.bin''' if optimizer_index == 0 else f'''{OPTIMIZER_NAME}_{optimizer_index}.bin'''
									)
									lowerCamelCase__	:	List[Any]				=os.path.join(__lowerCamelCase							,	__lowerCamelCase			)
									logger.info(f'''Loading Optimizer state from {input_optimizer_file}'''			)
									lowerCamelCase__	:	Optional[int]				=torch.load(__lowerCamelCase			)
									logger.info(f'''Optimizer state loaded from {input_optimizer_file}'''			)
						else:
									lowerCamelCase__	:	str				=(
									    os.path.join(__lowerCamelCase							,	f'''{OPTIMIZER_NAME}_{optimizer_index}'''			)
									    if f'''{OPTIMIZER_NAME}''' not in input_dir
									    else input_dir
									)
									logger.info(f'''Loading Optimizer from {ckpt_dir}'''			)
									lowerCamelCase__	:	Union[str, Any]				=load_sharded_optimizer_state_dict(
									    model_state_dict=model.state_dict()							,	optimizer_key='''optimizer'''							,	storage_reader=dist_cp.FileSystemReader(__lowerCamelCase			)							,	)
									lowerCamelCase__	:	str				=optim_state['''optimizer''']
									logger.info(f'''Optimizer loaded from {ckpt_dir}'''			)
						lowerCamelCase__	:	List[str]				=FSDP.optim_state_dict_to_load(__lowerCamelCase							,	__lowerCamelCase							,	__lowerCamelCase			)
						optimizer.load_state_dict(__lowerCamelCase			)









"""simple docstring"""






from typing import TYPE_CHECKING

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


_lowercase				:				Optional[Any]   		=		{
    "configuration_clip": [
        "CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "CLIPConfig",
        "CLIPOnnxConfig",
        "CLIPTextConfig",
        "CLIPVisionConfig",
    ],
    "processing_clip": ["CLIPProcessor"],
    "tokenization_clip": ["CLIPTokenizer"],
}

try:
		if not is_tokenizers_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase				:				str   		=		["CLIPTokenizerFast"]

try:
		if not is_vision_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase				:				Any   		=		["CLIPFeatureExtractor"]
		_lowercase				:				int   		=		["CLIPImageProcessor"]

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase				:				Optional[Any]   		=		[
		    "CLIP_PRETRAINED_MODEL_ARCHIVE_LIST",
		    "CLIPModel",
		    "CLIPPreTrainedModel",
		    "CLIPTextModel",
		    "CLIPTextModelWithProjection",
		    "CLIPVisionModel",
		    "CLIPVisionModelWithProjection",
		]

try:
		if not is_tf_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase				:				Dict   		=		[
		    "TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST",
		    "TFCLIPModel",
		    "TFCLIPPreTrainedModel",
		    "TFCLIPTextModel",
		    "TFCLIPVisionModel",
		]

try:
		if not is_flax_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase				:				Union[str, Any]   		=		[
		    "FlaxCLIPModel",
		    "FlaxCLIPPreTrainedModel",
		    "FlaxCLIPTextModel",
		    "FlaxCLIPTextPreTrainedModel",
		    "FlaxCLIPVisionModel",
		    "FlaxCLIPVisionPreTrainedModel",
		]


if TYPE_CHECKING:
		from .configuration_clip import (
		    CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    CLIPConfig,
		    CLIPOnnxConfig,
		    CLIPTextConfig,
		    CLIPVisionConfig,
		)
		from .processing_clip import CLIPProcessor
		from .tokenization_clip import CLIPTokenizer

		try:
				if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .tokenization_clip_fast import CLIPTokenizerFast

		try:
				if not is_vision_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .feature_extraction_clip import CLIPFeatureExtractor
				from .image_processing_clip import CLIPImageProcessor

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_clip import (
				    CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
				    CLIPModel,
				    CLIPPreTrainedModel,
				    CLIPTextModel,
				    CLIPTextModelWithProjection,
				    CLIPVisionModel,
				    CLIPVisionModelWithProjection,
				)

		try:
				if not is_tf_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_tf_clip import (
				    TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFCLIPModel,
				    TFCLIPPreTrainedModel,
				    TFCLIPTextModel,
				    TFCLIPVisionModel,
				)

		try:
				if not is_flax_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_flax_clip import (
				    FlaxCLIPModel,
				    FlaxCLIPPreTrainedModel,
				    FlaxCLIPTextModel,
				    FlaxCLIPTextPreTrainedModel,
				    FlaxCLIPVisionModel,
				    FlaxCLIPVisionPreTrainedModel,
				)


else:
		import sys

		_lowercase				:				Union[str, Any]   		=		_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)




