




'''simple docstring'''



import os
import random
from pathlib import Path
from typing import List

import numpy as np
import torch
from torch.cuda.amp import GradScaler

from .utils import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    get_pretty_name,
    is_tpu_available,
    is_xpu_available,
    save,
)


if is_tpu_available(check_device=False):
       import torch_xla.core.xla_model as xm

from .logging import get_logger
from .state import PartialState


UpperCAmelCase_       :  str      =		get_logger(__name__)







def 			_UpperCamelCase (_lowerCamelCase							:     str   ,	_lowerCamelCase							:     List[dict]   ,	_lowerCamelCase							:     list   ,	_lowerCamelCase							:     list   ,	_lowerCamelCase							:     int   ,	_lowerCamelCase							:     GradScaler = None   ,	)->   int:




     '''simple docstring'''





     for i, state in enumerate(_lowerCamelCase	):
          __snake_case    					=  f'''{MODEL_NAME}.bin''' if i == 0 else f'''{MODEL_NAME}_{i}.bin'''
          __snake_case    					=  os.path.join(_lowerCamelCase   ,	_lowerCamelCase	)
          save(_lowerCamelCase   ,	_lowerCamelCase	)
          logger.info(f'''Model weights saved in {output_model_file}'''	)
     # Optimizer states
     for i, opt in enumerate(_lowerCamelCase	):
          __snake_case    					=  opt.state_dict()
          __snake_case    					=  f'''{OPTIMIZER_NAME}.bin''' if i == 0 else f'''{OPTIMIZER_NAME}_{i}.bin'''
          __snake_case    					=  os.path.join(_lowerCamelCase   ,	_lowerCamelCase	)
          save(_lowerCamelCase   ,	_lowerCamelCase	)
          logger.info(f'''Optimizer state saved in {output_optimizer_file}'''	)
     # Scheduler states
     for i, scheduler in enumerate(_lowerCamelCase	):
          __snake_case    					=  scheduler.state_dict()
          __snake_case    					=  f'''{SCHEDULER_NAME}.bin''' if i == 0 else f'''{SCHEDULER_NAME}_{i}.bin'''
          __snake_case    					=  os.path.join(_lowerCamelCase   ,	_lowerCamelCase	)
          save(_lowerCamelCase   ,	_lowerCamelCase	)
          logger.info(f'''Scheduler state saved in {output_scheduler_file}'''	)
     # GradScaler state
     if scaler is not None:
          __snake_case    					=  scaler.state_dict()
          __snake_case    					=  os.path.join(_lowerCamelCase   ,	_lowerCamelCase	)
          torch.save(_lowerCamelCase   ,	_lowerCamelCase	)
          logger.info(f'''Gradient scaler state saved in {output_scaler_file}'''	)
     # Random number generator states
     __snake_case    					=  {}
     __snake_case    					=  f'''{RNG_STATE_NAME}_{process_index}.pkl'''
     __snake_case    					=  random.getstate()
     __snake_case    					=  np.random.get_state()
     __snake_case    					=  torch.get_rng_state()
     if is_xpu_available():
          __snake_case    					=  torch.xpu.get_rng_state_all()
     else:
          __snake_case    					=  torch.cuda.get_rng_state_all()
     if is_tpu_available():
          __snake_case    					=  xm.get_rng_state()
     __snake_case    					=  os.path.join(_lowerCamelCase   ,	_lowerCamelCase	)
     torch.save(_lowerCamelCase   ,	_lowerCamelCase	)
     logger.info(f'''Random states saved in {output_states_file}'''	)
     return output_dir







def 			_UpperCamelCase (_lowerCamelCase							:     Union[str, Any]   ,	_lowerCamelCase							:     List[str]   ,	_lowerCamelCase							:     Optional[int]   ,	_lowerCamelCase							:     Tuple   ,	_lowerCamelCase							:     Any   ,	_lowerCamelCase							:     Dict=None   ,	_lowerCamelCase							:     List[Any]=None   ,	**_lowerCamelCase							:     str   ,	)->   int:




     '''simple docstring'''





     if map_location not in [None, "cpu", "on_device"]:
          raise TypeError(
              '''Unsupported optimizer map location passed, please choose one of `None`, `\'cpu\'`, or `\'on_device\'`'''	)
     if map_location is None:
          __snake_case    					=  '''cpu'''
     elif map_location == "on_device":
          __snake_case    					=  PartialState().device
     # Model states
     for i, model in enumerate(_lowerCamelCase	):
          __snake_case    					=  f'''{MODEL_NAME}.bin''' if i == 0 else f'''{MODEL_NAME}_{i}.bin'''
          __snake_case    					=  os.path.join(_lowerCamelCase   ,	_lowerCamelCase	)
          models[i].load_state_dict(torch.load(_lowerCamelCase   ,	map_location=_lowerCamelCase	)   ,	**_lowerCamelCase	)
     logger.info('''All model weights loaded successfully'''	)

     # Optimizer states
     for i, opt in enumerate(_lowerCamelCase	):
          __snake_case    					=  f'''{OPTIMIZER_NAME}.bin''' if i == 0 else f'''{OPTIMIZER_NAME}_{i}.bin'''
          __snake_case    					=  os.path.join(_lowerCamelCase   ,	_lowerCamelCase	)
          __snake_case    					=  torch.load(_lowerCamelCase   ,	map_location=_lowerCamelCase	)
          optimizers[i].load_state_dict(_lowerCamelCase	)
     logger.info('''All optimizer states loaded successfully'''	)

     # Scheduler states
     for i, scheduler in enumerate(_lowerCamelCase	):
          __snake_case    					=  f'''{SCHEDULER_NAME}.bin''' if i == 0 else f'''{SCHEDULER_NAME}_{i}.bin'''
          __snake_case    					=  os.path.join(_lowerCamelCase   ,	_lowerCamelCase	)
          scheduler.load_state_dict(torch.load(_lowerCamelCase	)	)
     logger.info('''All scheduler states loaded successfully'''	)

     # GradScaler state
     if scaler is not None:
          __snake_case    					=  os.path.join(_lowerCamelCase   ,	_lowerCamelCase	)
          scaler.load_state_dict(torch.load(_lowerCamelCase	)	)
          logger.info('''GradScaler state loaded successfully'''	)

     # Random states
     try:
          __snake_case    					=  torch.load(os.path.join(_lowerCamelCase   ,	f'''{RNG_STATE_NAME}_{process_index}.pkl'''	)	)
          random.setstate(states['''random_state''']	)
          np.random.set_state(states['''numpy_random_seed''']	)
          torch.set_rng_state(states['''torch_manual_seed''']	)
          if is_xpu_available():
               torch.xpu.set_rng_state_all(states['''torch_xpu_manual_seed''']	)
          else:
               torch.cuda.set_rng_state_all(states['''torch_cuda_manual_seed''']	)
          if is_tpu_available():
               xm.set_rng_state(states['''xm_seed''']	)
          logger.info('''All random states loaded successfully'''	)
     except Exception:
          logger.info('''Could not load random states'''	)







def 			_UpperCamelCase (_lowerCamelCase							:     Optional[int]   ,	_lowerCamelCase							:     str   ,	_lowerCamelCase							:     int = 0	)->   str:




     '''simple docstring'''





     __snake_case    					=  Path(_lowerCamelCase	) / f'''custom_checkpoint_{index}.pkl'''
     logger.info(f'''Saving the state of {get_pretty_name(_lowerCamelCase	)} to {save_location}'''	)
     torch.save(obj.state_dict()   ,	_lowerCamelCase	)







def 			_UpperCamelCase (_lowerCamelCase							:     Tuple   ,	_lowerCamelCase							:     Union[str, Any]   ,	_lowerCamelCase							:     int = 0	)->   Optional[int]:




     '''simple docstring'''





     __snake_case    					=  f'''{path}/custom_checkpoint_{index}.pkl'''
     logger.info(f'''Loading the state of {get_pretty_name(_lowerCamelCase	)} from {load_location}'''	)
     obj.load_state_dict(torch.load(_lowerCamelCase   ,	map_location='''cpu'''	)	)












'''simple docstring'''



from ..utils import deprecate
from .attention_processor import (  # noqa: F401
    Attention,
    AttentionProcessor,
    AttnAddedKVProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRALinearLayer,
    LoRAXFormersAttnProcessor,
    SlicedAttnAddedKVProcessor,
    SlicedAttnProcessor,
    XFormersAttnProcessor,
)
from .attention_processor import AttnProcessor as AttnProcessorRename  # noqa: F401


deprecate(
    '''cross_attention''',
    '''0.20.0''',
    '''Importing from cross_attention is deprecated. Please import from diffusers.models.attention_processor instead.''',
    standard_warn=False,
)


UpperCAmelCase_       :  Optional[int]      =		AttentionProcessor





class     lowerCAmelCase   ( __lowerCAmelCase):


       def __init__(   self    ,  *__SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )      ->  Optional[int]:



            '''simple docstring'''




            __snake_case    					=  F'''{self.__class__.__name__} is deprecated and will be removed in `0.20.0`. Please use `from diffusers.models.attention_processor import {"".join(self.__class__.__name__.split("Cross"  )  )} instead.'''
            deprecate('''cross_attention'''    ,  '''0.20.0'''    ,  __SCREAMING_SNAKE_CASE    ,  standard_warn=__SCREAMING_SNAKE_CASE  )
            super().__init__(*__SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )





class     lowerCAmelCase   ( __lowerCAmelCase):


       def __init__(   self    ,  *__SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )      ->  Any:



            '''simple docstring'''




            __snake_case    					=  F'''{self.__class__.__name__} is deprecated and will be removed in `0.20.0`. Please use `from diffusers.models.attention_processor import {"".join(self.__class__.__name__.split("Cross"  )  )} instead.'''
            deprecate('''cross_attention'''    ,  '''0.20.0'''    ,  __SCREAMING_SNAKE_CASE    ,  standard_warn=__SCREAMING_SNAKE_CASE  )
            super().__init__(*__SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )





class     lowerCAmelCase   ( __lowerCAmelCase):


       def __init__(   self    ,  *__SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )      ->  Tuple:



            '''simple docstring'''




            __snake_case    					=  F'''{self.__class__.__name__} is deprecated and will be removed in `0.20.0`. Please use `from diffusers.models.attention_processor import {"".join(self.__class__.__name__.split("Cross"  )  )} instead.'''
            deprecate('''cross_attention'''    ,  '''0.20.0'''    ,  __SCREAMING_SNAKE_CASE    ,  standard_warn=__SCREAMING_SNAKE_CASE  )
            super().__init__(*__SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )





class     lowerCAmelCase   ( __lowerCAmelCase):


       def __init__(   self    ,  *__SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )      ->  Optional[Any]:



            '''simple docstring'''




            __snake_case    					=  F'''{self.__class__.__name__} is deprecated and will be removed in `0.20.0`. Please use `from diffusers.models.attention_processor import {"".join(self.__class__.__name__.split("Cross"  )  )} instead.'''
            deprecate('''cross_attention'''    ,  '''0.20.0'''    ,  __SCREAMING_SNAKE_CASE    ,  standard_warn=__SCREAMING_SNAKE_CASE  )
            super().__init__(*__SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )





class     lowerCAmelCase   ( __lowerCAmelCase):


       def __init__(   self    ,  *__SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )      ->  Optional[int]:



            '''simple docstring'''




            __snake_case    					=  F'''{self.__class__.__name__} is deprecated and will be removed in `0.20.0`. Please use `from diffusers.models.attention_processor import {"".join(self.__class__.__name__.split("Cross"  )  )} instead.'''
            deprecate('''cross_attention'''    ,  '''0.20.0'''    ,  __SCREAMING_SNAKE_CASE    ,  standard_warn=__SCREAMING_SNAKE_CASE  )
            super().__init__(*__SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )





class     lowerCAmelCase   ( __lowerCAmelCase):


       def __init__(   self    ,  *__SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )      ->  Tuple:



            '''simple docstring'''




            __snake_case    					=  F'''{self.__class__.__name__} is deprecated and will be removed in `0.20.0`. Please use `from diffusers.models.attention_processor import {"".join(self.__class__.__name__.split("Cross"  )  )} instead.'''
            deprecate('''cross_attention'''    ,  '''0.20.0'''    ,  __SCREAMING_SNAKE_CASE    ,  standard_warn=__SCREAMING_SNAKE_CASE  )
            super().__init__(*__SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )





class     lowerCAmelCase   ( __lowerCAmelCase):


       def __init__(   self    ,  *__SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )      ->  int:



            '''simple docstring'''




            __snake_case    					=  F'''{self.__class__.__name__} is deprecated and will be removed in `0.20.0`. Please use `from diffusers.models.attention_processor import {"".join(self.__class__.__name__.split("Cross"  )  )} instead.'''
            deprecate('''cross_attention'''    ,  '''0.20.0'''    ,  __SCREAMING_SNAKE_CASE    ,  standard_warn=__SCREAMING_SNAKE_CASE  )
            super().__init__(*__SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )





class     lowerCAmelCase   ( __lowerCAmelCase):


       def __init__(   self    ,  *__SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )      ->  List[str]:



            '''simple docstring'''




            __snake_case    					=  F'''{self.__class__.__name__} is deprecated and will be removed in `0.20.0`. Please use `from diffusers.models.attention_processor import {"".join(self.__class__.__name__.split("Cross"  )  )} instead.'''
            deprecate('''cross_attention'''    ,  '''0.20.0'''    ,  __SCREAMING_SNAKE_CASE    ,  standard_warn=__SCREAMING_SNAKE_CASE  )
            super().__init__(*__SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  )






