content: from partcrafter_src.utils.typing_utils import *
  
  import os
  from omegaconf import OmegaConf
  
  from torch import optim
  from torch.optim import lr_scheduler
  from diffusers.training_utils import *
  from diffusers.optimization import get_scheduler
  
  # https://github.com/huggingface/diffusers/pull/9812: fix `self.use_ema_warmup`
  class MyEMAModel(EMAModel):
      """
      Exponential Moving Average of models weights
      """
  
      def __init__(
          self,
          parameters: Iterable[torch.nn.Parameter],
          decay: float = 0.9999,
          min_decay: float = 0.0,
          update_after_step: int = 0,
          use_ema_warmup: bool = False,
          inv_gamma: Union[float, int] = 1.0,
          power: Union[float, int] = 2 / 3,
          foreach: bool = False,
          model_cls: Optional[Any] = None,
          model_config: Dict[str, Any] = None,
          **kwargs,
      ):
          """
          Args:
              parameters (Iterable[torch.nn.Parameter]): The parameters to track.
              decay (float): The decay factor for the exponential moving average.
              min_decay (float): The minimum decay factor for the exponential moving average.
              update_after_step (int): The number of steps to wait before starting to update the EMA weights.
              use_ema_warmup (bool): Whether to use EMA warmup.
              inv_gamma (float):
                  Inverse multiplicative factor of EMA warmup. Default: 1. Only used if `use_ema_warmup` is True.
              power (float): Exponential factor of EMA warmup. Default: 2/3. Only used if `use_ema_warmup` is True.
              foreach (bool): Use torch._foreach functions for updating shadow parameters. Should be faster.
              device (Optional[Union[str, torch.device]]): The device to store the EMA weights on. If None, the EMA
                          weights will be stored on CPU.
  
          @crowsonkb's notes on EMA Warmup:
              If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are good values for models you plan
              to train for a million or more steps (reaches decay factor 0.999 at 31.6K steps, 0.9999 at 1M steps),
              gamma=1, power=3/4 for models you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999
              at 215.4k steps).
          """
  
          if isinstance(parameters, torch.nn.Module):
              deprecation_message = (
                  "Passing a `torch.nn.Module` to `ExponentialMovingAverage` is deprecated. "
                  "Please pass the parameters of the module instead."
              )
              deprecate(
                  "passing a `torch.nn.Module` to `ExponentialMovingAverage`",
                  "1.0.0",
                  deprecation_message,
                  standard_warn=False,
              )
              parameters = parameters.parameters()
  
              # # set use_ema_warmup to True if a torch.nn.Module is passed for backwards compatibility
              # use_ema_warmup = True
  
          if kwargs.get("max_value", None) is not None:
              deprecation_message = "The `max_value` argument is deprecated. Please use `decay` instead."
              deprecate("max_value", "1.0.0", deprecation_message, standard_warn=False)
              decay = kwargs["max_value"]
  
          if kwargs.get("min_value", None) is not None:
              deprecation_message = "The `min_value` argument is deprecated. Please use `min_decay` instead."
              deprecate("min_value", "1.0.0", deprecation_message, standard_warn=False)
              min_decay = kwargs["min_value"]
  
          parameters = list(parameters)
          self.shadow_params = [p.clone().detach() for p in parameters]
  
          if kwargs.get("device", None) is not None:
              deprecation_message = "The `device` argument is deprecated. Please use `to` instead."
              deprecate("device", "1.0.0", deprecation_message, standard_warn=False)
              self.to(device=kwargs["device"])
  
          self.temp_stored_params = None
  
          self.decay = decay
          self.min_decay = min_decay
          self.update_after_step = update_after_step
          self.use_ema_warmup = use_ema_warmup
          self.inv_gamma = inv_gamma
          self.power = power
          self.optimization_step = 0
          self.cur_decay_value = None  # set in `step()`
          self.foreach = foreach
  
          self.model_cls = model_cls
          self.model_config = model_config
  
      def get_decay(self, optimization_step: int) -> float:
          """
          Compute the decay factor for the exponential moving average.
          """
          step = max(0, optimization_step - self.update_after_step - 1)
  
          if step <= 0:
              return 0.0
  
          if self.use_ema_warmup:
              cur_decay_value = 1 - (1 + step / self.inv_gamma) ** -self.power
          else:
              # cur_decay_value = (1 + step) / (10 + step)
              cur_decay_value = self.decay
  
          cur_decay_value = min(cur_decay_value, self.decay)
          # make sure decay is not smaller than min_decay
          cur_decay_value = max(cur_decay_value, self.min_decay)
          return cur_decay_value
  
  def get_configs(yaml_path: str, cli_configs: List[str]=[], **kwargs) -> DictConfig:
      yaml_configs = OmegaConf.load(yaml_path)
      cli_configs = OmegaConf.from_cli(cli_configs)
  
      configs = OmegaConf.merge(yaml_configs, cli_configs, kwargs)
      OmegaConf.resolve(configs)  # resolve ${...} placeholders
      return configs
  
  def get_optimizer(name: str, params: Parameter, **kwargs) -> Optimizer:
      if name == "adamw":
          return optim.AdamW(params=params, **kwargs)
      else:
          raise NotImplementedError(f"Not implemented optimizer: {name}")
  
  def get_lr_scheduler(name: str, optimizer: Optimizer, **kwargs) -> LRScheduler:
      if name == "one_cycle":
          return lr_scheduler.OneCycleLR(
              optimizer,
              max_lr=kwargs["max_lr"],
              total_steps=kwargs["total_steps"],
              pct_start=kwargs["pct_start"],
          )
      elif name == "cosine_warmup":
          return get_scheduler(
              "cosine", optimizer,
              num_warmup_steps=kwargs["num_warmup_steps"],
              num_training_steps=kwargs["total_steps"],
          )
      elif name == "constant_warmup":
          return get_scheduler(
              "constant_with_warmup", optimizer,
              num_warmup_steps=kwargs["num_warmup_steps"],
              num_training_steps=kwargs["total_steps"],
          )
      elif name == "constant":
          return lr_scheduler.LambdaLR(optimizer=optimizer, lr_lambda=lambda _: 1)
      elif name == "linear_decay":
          return lr_scheduler.LambdaLR(
              optimizer=optimizer,
              lr_lambda=lambda epoch: max(0., 1. - epoch / kwargs["total_epochs"]),
          )
      else:
          raise NotImplementedError(f"Not implemented lr scheduler: {name}")
  
  def save_experiment_params(
      args: Namespace, 
      configs: DictConfig, 
      save_dir: str
  ) -> Dict[str, Any]:
      params = OmegaConf.merge(configs, {"args": {k: str(v) for k, v in vars(args).items()}})
      OmegaConf.save(params, os.path.join(save_dir, "params.yaml"))
      return dict(params)
  
  
  def save_model_architecture(model: Module, save_dir: str) -> None:
      num_buffers = sum(b.numel() for b in model.buffers())
      num_params = sum(p.numel() for p in model.parameters())
      num_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
      message = f"Number of buffers: {num_buffers}\n" +\
          f"Number of trainable / all parameters: {num_trainable_params} / {num_params}\n\n" +\
          f"Model architecture:\n{model}"
  
      with open(os.path.join(save_dir, "model.txt"), "w") as f:
          f.write(message)