code: 
  
  
  
  
  
  '''simple docstring'''
  
  from dataclasses import dataclass
  from typing import List, Optional, Tuple, Union
  
  import numpy as np
  import torch
  
  from ..configuration_utils import ConfigMixin, register_to_config
  from ..utils import BaseOutput, logging, randn_tensor
  from .scheduling_utils import SchedulerMixin
  
  
  __lowerCAmelCase							:			Dict				=logging.get_logger(__name__)  # pylint: disable=invalid-name
  
  
  
  
  
  
  @dataclass
  class      UpperCAmelCase			( UpperCamelCase__       ):
   __lowercase     	=       42
  
  
  
  
  
  
  class      UpperCAmelCase			( UpperCamelCase__    , UpperCamelCase__       ):
   __lowercase     	=       1
   @register_to_config
   def __init__(       self   :Dict		, lowercase_   :int = 40		, lowercase_   :float = 0.0_0_2		, lowercase_   :float = 8_0.0		, lowercase_   :float = 0.5		, lowercase_   :float = 1.0		, lowercase_   :float = 7.0		, lowercase_   :bool = True		, )->							Any:
       # standard deviation of the initial noise distribution
       A__       =    sigma_max
  
       A__       =    np.linspace(0		, 1		, lowercase_							)
       A__       =    self._convert_to_karras(lowercase_							)
       A__       =    self.sigma_to_t(lowercase_							)
  
       # setable values
       A__       =    None
       A__       =    torch.from_numpy(lowercase_							)
       A__       =    torch.from_numpy(lowercase_							)
       A__       =    False
       A__       =    False
   def   UpperCAmelCase_							(       self   :str		, lowercase_   :Union[str, Any]		, lowercase_   :Optional[int]=None							)->							Union[str, Any]:
       if schedule_timesteps is None:
           A__       =    self.timesteps
  
       A__       =    (schedule_timesteps == timestep).nonzero()
       return indices.item()
   def   UpperCAmelCase_							(       self   :Any		, lowercase_   :torch.FloatTensor		, lowercase_   :Union[float, torch.FloatTensor]							)->							torch.FloatTensor:
       # Get sigma corresponding to timestep
       if isinstance(lowercase_		, torch.Tensor							):
           A__       =    timestep.to(self.timesteps.device							)
       A__       =    self.index_for_timestep(lowercase_							)
       A__       =    self.sigmas[step_idx]
  
       A__       =    sample / ((sigma**2 + self.config.sigma_data**2) ** 0.5)
  
       A__       =    True
       return sample
   def   UpperCAmelCase_							(       self   :Optional[Any]		, lowercase_   :Union[float, np.ndarray]							)->							Dict:
       if not isinstance(lowercase_		, np.ndarray							):
           A__       =    np.array(lowercase_		, dtype=np.floataa							)
  
       A__       =    10_00 * 0.2_5 * np.log(sigmas + 1E-44							)
  
       return timesteps
   def   UpperCAmelCase_							(       self   :Tuple		, lowercase_   :Optional[int] = None		, lowercase_   :Union[str, torch.device] = None		, lowercase_   :Optional[List[int]] = None		, )->							Optional[int]:
       if num_inference_steps is None and timesteps is None:
           raise ValueError("Exactly one of `num_inference_steps` or `timesteps` must be supplied."							)
  
       if num_inference_steps is not None and timesteps is not None:
           raise ValueError("Can only pass one of `num_inference_steps` or `timesteps`."							)
  
       # Follow DDPMScheduler custom timesteps logic
       if timesteps is not None:
           for i in range(1		, len(lowercase_							)							):
               if timesteps[i] >= timesteps[i - 1]:
                   raise ValueError("`timesteps` must be in descending order."							)
  
           if timesteps[0] >= self.config.num_train_timesteps:
               raise ValueError(
                   F"`timesteps` must start before `self.config.train_timesteps`:"
                   F" {self.config.num_train_timesteps}."							)
  
           A__       =    np.array(lowercase_		, dtype=np.intaa							)
           A__       =    True
       else:
           if num_inference_steps > self.config.num_train_timesteps:
               raise ValueError(
                   F"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"
                   F" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"
                   F" maximal {self.config.num_train_timesteps} timesteps."							)
  
           A__       =    num_inference_steps
  
           A__       =    self.config.num_train_timesteps // self.num_inference_steps
           A__       =    (np.arange(0		, lowercase_							) * step_ratio).round()[::-1].copy().astype(np.intaa							)
           A__       =    False
  
       # Map timesteps to Karras sigmas directly for multistep sampling
       # See https://github.com/openai/consistency_models/blob/main/cm/karras_diffusion.py#L675
       A__       =    self.config.num_train_timesteps
       A__       =    timesteps[::-1].copy()
       A__       =    ramp / (num_train_timesteps - 1)
       A__       =    self._convert_to_karras(lowercase_							)
       A__       =    self.sigma_to_t(lowercase_							)
  
       A__       =    np.concatenate([sigmas, [self.sigma_min]]							).astype(np.floataa							)
       A__       =    torch.from_numpy(lowercase_							).to(device=lowercase_							)
  
       if str(lowercase_							).startswith("mps"							):
           # mps does not support float64
           A__       =    torch.from_numpy(lowercase_							).to(lowercase_		, dtype=torch.floataa							)
       else:
           A__       =    torch.from_numpy(lowercase_							).to(device=lowercase_							)
   def   UpperCAmelCase_							(       self   :Optional[int]		, lowercase_   :str							)->							str:
       A__       =    self.config.sigma_min
       A__       =    self.config.sigma_max
  
       A__       =    self.config.rho
       A__       =    sigma_min ** (1 / rho)
       A__       =    sigma_max ** (1 / rho)
       A__       =    (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
       return sigmas
   def   UpperCAmelCase_							(       self   :Optional[int]		, lowercase_   :Optional[Any]							)->							Optional[int]:
       A__       =    self.config.sigma_data
  
       A__       =    sigma_data**2 / (sigma**2 + sigma_data**2)
       A__       =    sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
       return c_skip, c_out
   def   UpperCAmelCase_							(       self   :Optional[int]		, lowercase_   :int							)->							Any:
       A__       =    self.config.sigma_min
       A__       =    self.config.sigma_data
  
       A__       =    sigma_data**2 / ((sigma - sigma_min) ** 2 + sigma_data**2)
       A__       =    (sigma - sigma_min) * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
       return c_skip, c_out
   def   UpperCAmelCase_							(       self   :str		, lowercase_   :torch.FloatTensor		, lowercase_   :Union[float, torch.FloatTensor]		, lowercase_   :torch.FloatTensor		, lowercase_   :Optional[torch.Generator] = None		, lowercase_   :bool = True		, )->							Union[CMStochasticIterativeSchedulerOutput, Tuple]:
  
       if (
           isinstance(lowercase_		, lowercase_							)
           or isinstance(lowercase_		, torch.IntTensor							)
           or isinstance(lowercase_		, torch.LongTensor							)
       ):
           raise ValueError(
               (
                   "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
                   F" `{self.__class__}.step()` is not supported. Make sure to pass"
                   " one of the `scheduler.timesteps` as a timestep."
               )		, )
  
       if not self.is_scale_input_called:
           logger.warning(
               "The `scale_model_input` function should be called before `step` to ensure correct denoising. "
               "See `StableDiffusionPipeline` for a usage example."							)
  
       if isinstance(lowercase_		, torch.Tensor							):
           A__       =    timestep.to(self.timesteps.device							)
  
       A__       =    self.config.sigma_min
       A__       =    self.config.sigma_max
  
       A__       =    self.index_for_timestep(lowercase_							)
  
       # sigma_next corresponds to next_t in original implementation
       A__       =    self.sigmas[step_index]
       if step_index + 1 < self.config.num_train_timesteps:
           A__       =    self.sigmas[step_index + 1]
       else:
           # Set sigma_next to sigma_min
           A__       =    self.sigmas[-1]
  
       # Get scalings for boundary conditions
       A__,			A__       =    self.get_scalings_for_boundary_condition(lowercase_							)
  
       # 1. Denoise model output using boundary conditions
       A__       =    c_out * model_output + c_skip * sample
       if self.config.clip_denoised:
           A__       =    denoised.clamp(-1		, 1							)
  
       # 2. Sample z ~ N(0, s_noise^2 * I)
       # Noise is not used for onestep sampling.
       if len(self.timesteps							) > 1:
           A__       =    randn_tensor(
               model_output.shape		, dtype=model_output.dtype		, device=model_output.device		, generator=lowercase_							)
       else:
           A__       =    torch.zeros_like(lowercase_							)
       A__       =    noise * self.config.s_noise
  
       A__       =    sigma_next.clamp(min=lowercase_		, max=lowercase_							)
  
       # 3. Return noisy sample
       # tau = sigma_hat, eps = sigma_min
       A__       =    denoised + z * (sigma_hat**2 - sigma_min**2) ** 0.5
  
       if not return_dict:
           return (prev_sample,)
  
       return CMStochasticIterativeSchedulerOutput(prev_sample=lowercase_							)
   def   UpperCAmelCase_							(       self   :List[str]		, lowercase_   :torch.FloatTensor		, lowercase_   :torch.FloatTensor		, lowercase_   :torch.FloatTensor		, )->							torch.FloatTensor:
       # Make sure sigmas and timesteps have the same device and dtype as original_samples
       A__       =    self.sigmas.to(device=original_samples.device		, dtype=original_samples.dtype							)
       if original_samples.device.type == "mps" and torch.is_floating_point(lowercase_							):
           # mps does not support float64
           A__       =    self.timesteps.to(original_samples.device		, dtype=torch.floataa							)
           A__       =    timesteps.to(original_samples.device		, dtype=torch.floataa							)
       else:
           A__       =    self.timesteps.to(original_samples.device							)
           A__       =    timesteps.to(original_samples.device							)
  
       A__       =    [(schedule_timesteps == t).nonzero().item() for t in timesteps]
  
       A__       =    sigmas[step_indices].flatten()
       while len(sigma.shape							) < len(original_samples.shape							):
           A__       =    sigma.unsqueeze(-1							)
  
       A__       =    original_samples + noise * sigma
       return noisy_samples
  
  
  
   def __len__(       self   :Any							)->							Any:
       return self.config.num_train_timesteps
  
  
  
  
style_context: 
  
  
  
  
  
  '''simple docstring'''
  
  import collections
  
  from .utils import ExplicitEnum, is_torch_available, logging
  
  
  if is_torch_available():
    import torch
  
  
  __lowerCAmelCase							:			Optional[int]				=logging.get_logger(__name__)
  
  
  
  
  
  
  class      UpperCAmelCase			:
   def __init__(       self   :Optional[int]		, lowercase_   :Any		, lowercase_   :Any=21		, lowercase_   :Any=[]		, lowercase_   :Dict=None							)->							Tuple:
       A__       =    model
       A__       =    trace_batch_nums
       A__       =    abort_after_batch_num
  
       # keep a LIFO buffer of frames to dump as soon as inf/nan is encountered to give context to the problem emergence
       A__       =    collections.deque([]		, lowercase_							)
       A__       =    []
       A__       =    0
       A__       =    0
       A__       =    False
       A__       =    "                 "
  
       self.analyse_model()
  
       self.register_forward_hook()
   def   UpperCAmelCase_							(       self   :Any		, lowercase_   :Any=None							)->							Optional[Any]:
       if frame is not None:
           self.expand_frame(lowercase_							)
       self.frames.append("\n".join(self.frame							)							)
       A__       =    []  # start a new frame
   def   UpperCAmelCase_							(       self   :Union[str, Any]		, lowercase_   :Dict							)->							List[str]:
       self.frame.append(lowercase_							)
   def   UpperCAmelCase_							(       self   :int							)->							Dict:
       print("\n".join(self.frames							)							)
       A__       =    []
   def   UpperCAmelCase_							(       self   :str							)->							Union[str, Any]:
       A__       =    []
   def   UpperCAmelCase_							(       self   :List[str]							)->							Optional[Any]:
       print(F"\nDetected inf/nan during batch_number={self.batch_number}"							)
       print(F"Last {len(self.frames							)} forward frames:"							)
       print(F"{'abs min':8} {'abs max':8} metadata"							)
       print("\n".join(self.frames							)							)
       print("\n\n"							)
       A__       =    []
   def   UpperCAmelCase_							(       self   :Any							)->							Optional[Any]:
       # extract the fully qualified module names, to be able to report at run time. e.g.:
       # encoder.block.2.layer.0.SelfAttention.o
       #
       # for shared weights only the first shared module name will be registered
       A__       =    {m: name for name, m in self.model.named_modules()}
       # self.longest_module_name = max(len(v) for v in self.module_names.values())
   def   UpperCAmelCase_							(       self   :Tuple		, lowercase_   :Dict		, lowercase_   :Any							)->							str:
       if torch.is_tensor(lowercase_							):
           self.expand_frame(get_abs_min_max(lowercase_		, lowercase_							)							)
           if detect_overflow(lowercase_		, lowercase_							):
               A__       =    True
       elif var is None:
           self.expand_frame(F"{'None':>17} {ctx}"							)
       else:
           self.expand_frame(F"{'not a tensor':>17} {ctx}"							)
   def   UpperCAmelCase_							(       self   :Union[str, Any]							)->							str:
       self.expand_frame(F"\n\n{self.prefix} *** Starting batch number={self.batch_number} ***"							)
       self.expand_frame(F"{'abs min':8} {'abs max':8} metadata"							)
   def   UpperCAmelCase_							(       self   :int							)->							int:
       self.expand_frame(F"{self.prefix} *** Finished batch number={self.batch_number-1} ***\n\n"							)
   def   UpperCAmelCase_							(       self   :Optional[int]		, lowercase_   :Dict		, lowercase_   :Dict		, lowercase_   :str							)->							int:
       self.expand_frame(F"{self.prefix} {self.module_names[module]} {module.__class__.__name__}"							)
  
       # params
       for name, p in module.named_parameters(recurse=lowercase_							):
           self.analyse_variable(lowercase_		, lowercase_							)
  
       # inputs
       if isinstance(lowercase_		, lowercase_							):
           for i, x in enumerate(lowercase_							):
               self.analyse_variable(lowercase_		, F"input[{i}]"							)
       else:
           self.analyse_variable(lowercase_		, "input"							)
  
       # outputs
       if isinstance(lowercase_		, lowercase_							):
           for i, x in enumerate(lowercase_							):
               # possibly a tuple of tuples
               if isinstance(lowercase_		, lowercase_							):
                   for j, y in enumerate(lowercase_							):
                       self.analyse_variable(lowercase_		, F"output[{i}][{j}]"							)
               else:
                   self.analyse_variable(lowercase_		, F"output[{i}]"							)
       else:
           self.analyse_variable(lowercase_		, "output"							)
  
       self.save_frame()
   def   UpperCAmelCase_							(       self   :Any							)->							int:
       self.model.apply(self._register_forward_hook							)
   def   UpperCAmelCase_							(       self   :Tuple		, lowercase_   :Optional[int]							)->							Tuple:
       module.register_forward_hook(self.forward_hook							)
  
  
  
   def   UpperCAmelCase_							(       self   :List[str]		, lowercase_   :int		, lowercase_   :Optional[Any]		, lowercase_   :Dict							)->							Union[str, Any]:
       # - input is a tuple of packed inputs (could be non-Tensors)
       # - output could be a Tensor or a tuple of Tensors and non-Tensors
  
       A__       =    False
  
       A__       =    True if self.batch_number in self.trace_batch_nums else False
       if trace_mode:
           self.reset_saved_frames()
  
       if self.total_calls == 0:
           self.batch_start_frame()
       self.total_calls += 1
  
       # count batch numbers - the very first forward hook of the batch will be called when the
       # batch completes - i.e. it gets called very last - we know this batch has finished
       if module == self.model:
           self.batch_number += 1
           A__       =    True
  
       self.create_frame(lowercase_		, lowercase_		, lowercase_							)
  
       # if last_frame_of_batch:
       #     self.batch_end_frame()
  
       if trace_mode:
           self.trace_frames()
  
       if last_frame_of_batch:
           self.batch_start_frame()
  
       if self.detected_overflow and not trace_mode:
           self.dump_saved_frames()
  
           # now we can abort, as it's pointless to continue running
           raise ValueError(
               "DebugUnderflowOverflow: inf/nan detected, aborting as there is no point running further. "
               "Please scroll up above this traceback to see the activation values prior to this event."							)
  
       # abort after certain batch if requested to do so
       if self.abort_after_batch_num is not None and self.batch_number > self.abort_after_batch_num:
           raise ValueError(
               F"DebugUnderflowOverflow: aborting after {self.batch_number} batches due to"
               F" `abort_after_batch_num={self.abort_after_batch_num}` arg"							)
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      Dict		, _lowerCamelCase      :      int ):
      A__       =    var.abs()
      return F"{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}"
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      Tuple		, _lowerCamelCase      :      Optional[int] ):
      A__       =    False
      if torch.isnan(_lowerCamelCase ).any().item():
          A__       =    True
          print(F"{ctx} has nans" )
      if torch.isinf(_lowerCamelCase ).any().item():
          A__       =    True
          print(F"{ctx} has infs" )
  
      # if needed to monitor large elements can enable the following
      if 0:  # and detected:
          A__       =    var[torch.ge(var.abs()		, 1_00 )]
          if naaa.numel() > 0:
              print(F"{ctx}:  n100={naaa.numel()}" )
          A__       =    var[torch.ge(var.abs()		, 10_00 )]
          if naaaa.numel() > 0:
              print(F"{ctx}: n1000={naaaa.numel()}" )
          A__       =    var[torch.ge(var.abs()		, 1_00_00 )]
          if naaaaa.numel() > 0:
              print(F"{ctx}: n10000={naaaaa.numel()}" )
  
      if 0:
          print(F"min={var.min():9.2e} max={var.max():9.2e}" )
  
      if 0:
          print(F"min={var.min():9.2e} max={var.max():9.2e} var={var.var():9.2e} mean={var.mean():9.2e} ({ctx})" )
  
      return detected
  
  
  
  
  
  
  class      UpperCAmelCase			( UpperCamelCase__       ):
   __lowercase     	=       """underflow_overflow"""
   __lowercase     	=       """tpu_metrics_debug"""
  
  
  
  