code: 
  
  
  
  
  '''simple docstring'''
  
  
  
  
  
  
  
  from dataclasses import dataclass
  from typing import Optional, Tuple, Union
  
  import flax
  import jax.numpy as jnp
  from scipy import integrate
  
  from ..configuration_utils import ConfigMixin, register_to_config
  from .scheduling_utils_flax import (
      CommonSchedulerState,
      FlaxKarrasDiffusionSchedulers,
      FlaxSchedulerMixin,
      FlaxSchedulerOutput,
      broadcast_to_shape_from_left,
  )
  
  
  
  
  
  @flax.struct.dataclass
  class   lowerCamelCase_							:
  
  
  
  
  					'''simple docstring'''
  
  
  
  
  
  
  					lowercase__					=				42
  
  					# setable values
  					lowercase__					=				42
  					lowercase__					=				42
  					lowercase__					=				42
  					lowercase__					=				None
  
  					# running values
  					lowercase__					=				None
  
  
  
  
  
  
  					@classmethod
  					def        __SCREAMING_SNAKE_CASE			(					cls,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase				):
  							return cls(common=__UpperCamelCase,						init_noise_sigma=__UpperCamelCase,						timesteps=__UpperCamelCase,						sigmas=__UpperCamelCase				)
  
  
  
  
  
  @dataclass
  class   lowerCamelCase_							( UpperCAmelCase_							):
  
  
  
  
  					'''simple docstring'''
  
  
  
  
  
  
  					lowercase__					=				42
  
  
  
  
  
  class   lowerCamelCase_							( UpperCAmelCase_, UpperCAmelCase_							):
  
  
  
  
  					'''simple docstring'''
  
  
  
  
  
  
  					lowercase__					=				[e.name for e in FlaxKarrasDiffusionSchedulers]
  
  					lowercase__					=				42
  
  
  
  
  
  
  					@property
  					def        __SCREAMING_SNAKE_CASE			(					self				):
  							return True
  
  
  
  
  
  
  					@register_to_config
  					def __init__(					self,						__UpperCamelCase = 1_000,						__UpperCamelCase = 0.00_01,						__UpperCamelCase = 0.02,						__UpperCamelCase = "linear",						__UpperCamelCase = None,						__UpperCamelCase = "epsilon",						__UpperCamelCase = jnp.floataa,						):
  							_snake_case			 = dtype
  
  
  
  
  
  
  					def        __SCREAMING_SNAKE_CASE			(					self,						__UpperCamelCase = None				):
  							if common is None:
  									_snake_case			 = CommonSchedulerState.create(self				)
  
  							_snake_case			 = jnp.arange(0,						self.config.num_train_timesteps				).round()[::-1]
  							_snake_case			 = ((1 - common.alphas_cumprod) / common.alphas_cumprod) ** 0.5
  
  							# standard deviation of the initial noise distribution
  							_snake_case			 = sigmas.max()
  
  							return LMSDiscreteSchedulerState.create(
  							    common=__UpperCamelCase,						init_noise_sigma=__UpperCamelCase,						timesteps=__UpperCamelCase,						sigmas=__UpperCamelCase,						)
  
  
  
  
  
  
  					def        __SCREAMING_SNAKE_CASE			(					self,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase				):
  							((_snake_case) ,					)			 = jnp.where(state.timesteps == timestep,						size=1				)
  							_snake_case			 = step_index[0]
  
  							_snake_case			 = state.sigmas[step_index]
  							_snake_case			 = sample / ((sigma**2 + 1) ** 0.5)
  							return sample
  
  
  
  
  
  
  					def        __SCREAMING_SNAKE_CASE			(					self,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase				):
  
  							def lms_derivative(__UpperCamelCase				):
  									_snake_case			 = 1.0
  									for k in range(__UpperCamelCase				):
  											if current_order == k:
  													continue
  											prod *= (tau - state.sigmas[t - k]) / (state.sigmas[t - current_order] - state.sigmas[t - k])
  									return prod
  
  							_snake_case			 = integrate.quad(__UpperCamelCase,						state.sigmas[t],						state.sigmas[t + 1],						epsrel=1E-4				)[0]
  
  							return integrated_coeff
  
  
  
  
  
  
  					def        __SCREAMING_SNAKE_CASE			(					self,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase = ()				):
  							_snake_case			 = jnp.linspace(self.config.num_train_timesteps - 1,						0,						__UpperCamelCase,						dtype=self.dtype				)
  
  							_snake_case			 = jnp.floor(__UpperCamelCase				).astype(jnp.intaa				)
  							_snake_case			 = jnp.ceil(__UpperCamelCase				).astype(jnp.intaa				)
  
  							_snake_case			 = jnp.mod(__UpperCamelCase,						1.0				)
  
  							_snake_case			 = ((1 - state.common.alphas_cumprod) / state.common.alphas_cumprod) ** 0.5
  							_snake_case			 = (1 - frac) * sigmas[low_idx] + frac * sigmas[high_idx]
  							_snake_case			 = jnp.concatenate([sigmas, jnp.array([0.0],						dtype=self.dtype				)]				)
  
  							_snake_case			 = timesteps.astype(jnp.intaa				)
  
  							# initial running values
  							_snake_case			 = jnp.zeros((0,) + shape,						dtype=self.dtype				)
  
  							return state.replace(
  							    timesteps=__UpperCamelCase,						sigmas=__UpperCamelCase,						num_inference_steps=__UpperCamelCase,						derivatives=__UpperCamelCase,						)
  
  
  
  
  
  
  					def        __SCREAMING_SNAKE_CASE			(					self,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase = 4,						__UpperCamelCase = True,						):
  							if state.num_inference_steps is None:
  									raise ValueError(
  									    "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"				)
  
  							_snake_case			 = state.sigmas[timestep]
  
  							# 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
  							if self.config.prediction_type == "epsilon":
  									_snake_case			 = sample - sigma * model_output
  							elif self.config.prediction_type == "v_prediction":
  									# * c_out + input * c_skip
  									_snake_case			 = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1))
  							else:
  									raise ValueError(
  									    f'''prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`'''				)
  
  							# 2. Convert to an ODE derivative
  							_snake_case			 = (sample - pred_original_sample) / sigma
  							_snake_case			 = state.replace(derivatives=jnp.append(state.derivatives,						__UpperCamelCase				)				)
  							if len(state.derivatives				) > order:
  									_snake_case			 = state.replace(derivatives=jnp.delete(state.derivatives,						0				)				)
  
  							# 3. Compute linear multistep coefficients
  							_snake_case			 = min(timestep + 1,						__UpperCamelCase				)
  							_snake_case			 = [self.get_lms_coefficient(__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase				) for curr_order in range(__UpperCamelCase				)]
  
  							# 4. Compute previous sample based on the derivatives path
  							_snake_case			 = sample + sum(
  							    coeff * derivative for coeff, derivative in zip(__UpperCamelCase,						reversed(state.derivatives				)				)				)
  
  							if not return_dict:
  									return (prev_sample, state)
  
  							return FlaxLMSSchedulerOutput(prev_sample=__UpperCamelCase,						state=__UpperCamelCase				)
  
  
  
  
  
  
  					def        __SCREAMING_SNAKE_CASE			(					self,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase,						):
  							_snake_case			 = state.sigmas[timesteps].flatten()
  							_snake_case			 = broadcast_to_shape_from_left(__UpperCamelCase,						noise.shape				)
  
  							_snake_case			 = original_samples + noise * sigma
  
  							return noisy_samples
  
  
  
  
  
  
  					def __len__(					self				):
  							return self.config.num_train_timesteps
  
  
  
  
  
  
style_context: 
  
  
  
  
  '''simple docstring'''
  
  
  
  
  
  
  
  import importlib
  import importlib.metadata
  import os
  import warnings
  from distutils.util import strtobool
  from functools import lru_cache
  
  import torch
  from packaging import version
  from packaging.version import parse
  
  from .environment import parse_flag_from_env
  from .versions import compare_versions, is_torch_version
  
  
  try:
  						import torch_xla.core.xla_model as xm  # noqa: F401
  
  						_lowerCamelCase							:		int										=						True
  except ImportError:
  						_lowerCamelCase							:		List[str]										=						False
  
  
  # Cache this result has it's a C FFI call which can be pretty time-consuming
  _lowerCamelCase							:		List[Any]										=						torch.distributed.is_available()
  
  
  
  
  
  
  def 	_lowercase	(			lowercase_						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		_snake_case			 = importlib.util.find_spec(lowercase_						) is not None
  		if package_exists:
  				try:
  						_snake_case			 = importlib.metadata.metadata(lowercase_						)
  						return True
  				except importlib.metadata.PackageNotFoundError:
  						return False
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _torch_distributed_available
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		try:
  				pass
  		except ImportError:
  				print(
  				    "Intel(R) oneCCL Bindings for PyTorch* is required to run DDP on Intel(R) GPUs, but it is not"
  				    " detected. If you see \"ValueError: Invalid backend: 'ccl'\" error, please install Intel(R) oneCCL"
  				    " Bindings for PyTorch*."						)
  		return (
  		    importlib.util.find_spec("torch_ccl"						) is not None
  		    or importlib.util.find_spec("oneccl_bindings_for_pytorch"						) is not None
  		)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return importlib.metadata.version("oneccl_bind_pt"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("transformer_engine"						)
  
  
  
  
  
  
  @lru_cache
  def 	_lowercase	(			lowercase_=True						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if torch.cuda.is_available():
  				return False
  		if _tpu_available and check_device:
  				try:
  						# Will raise a RuntimeError if no XLA configuration is found
  						_snake_case			 = xm.xla_device()
  						return True
  				except RuntimeError:
  						return False
  		return _tpu_available
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("deepspeed"						)
  
  
  
  
  
  
  def 	_lowercase	(			lowercase_=False						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if is_tpu_available():
  				return not ignore_tpu
  		if torch.cuda.is_available():
  				return torch.cuda.is_bfaa_supported()
  		if is_npu_available():
  				return False
  		return True
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		_snake_case			 = _is_package_available("bitsandbytes"						)
  		if package_exists:
  				_snake_case			 = version.parse(importlib.metadata.version("bitsandbytes"						)						)
  				return compare_versions(lowercase_		,							">="		,							"0.39.0"						)
  		return False
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		_snake_case			 = _is_package_available("bitsandbytes"						)
  		if package_exists:
  				_snake_case			 = version.parse(importlib.metadata.version("bitsandbytes"						)						)
  				return compare_versions(lowercase_		,							">="		,							"0.37.2"						)
  		return False
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("bitsandbytes"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if strtobool(os.environ.get("ACCELERATE_USE_MEGATRON_LM"		,							"False"						)						) == 1:
  				_snake_case			 = importlib.util.find_spec("megatron"						) is not None
  				if package_exists:
  						try:
  								_snake_case			 = parse(importlib.metadata.version("megatron-lm"						)						)
  								return compare_versions(lowercase_		,							">="		,							"2.2.0"						)
  						except Exception as e:
  								warnings.warn(f'''Parse Megatron version failed. Exception:{e}'''						)
  								return False
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("safetensors"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("transformers"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("datasets"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("aim"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("tensorboard"						) or _is_package_available("tensorboardX"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("wandb"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("comet_ml"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("boto3"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if _is_package_available("rich"						):
  				if "ACCELERATE_DISABLE_RICH" in os.environ:
  						warnings.warn(
  						    "`ACCELERATE_DISABLE_RICH` is deprecated and will be removed in v0.22.0 and deactivated by default. Please use `ACCELERATE_ENABLE_RICH` if you wish to use `rich`."						)
  						return not parse_flag_from_env("ACCELERATE_DISABLE_RICH"		,							lowercase_						)
  				return parse_flag_from_env("ACCELERATE_ENABLE_RICH"		,							lowercase_						)
  		return False
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("sagemaker"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("tqdm"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("mlflow"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return is_torch_version(">="		,							"1.12"						) and torch.backends.mps.is_available() and torch.backends.mps.is_built()
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		def get_major_and_minor_from_version(lowercase_						):
  				return str(version.parse(lowercase_						).major						) + "." + str(version.parse(lowercase_						).minor						)
  
  		_snake_case			 = importlib.metadata.version("torch"						)
  		if importlib.util.find_spec("intel_extension_for_pytorch"						) is None:
  				return False
  		_snake_case			 = "N/A"
  		try:
  				_snake_case			 = importlib.metadata.version("intel_extension_for_pytorch"						)
  		except importlib.metadata.PackageNotFoundError:
  				return False
  		_snake_case			 = get_major_and_minor_from_version(_torch_version						)
  		_snake_case			 = get_major_and_minor_from_version(_ipex_version						)
  		if torch_major_and_minor != ipex_major_and_minor:
  				warnings.warn(
  				    f'''Intel Extension for PyTorch {ipex_major_and_minor} needs to work with PyTorch {ipex_major_and_minor}.*,'''
  				    f''' but PyTorch {_torch_version} is found. Please switch to the matching version and run again.'''						)
  				return False
  		return True
  
  
  
  
  
  
  @lru_cache
  def 	_lowercase	(			lowercase_=False						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if importlib.util.find_spec("torch"						) is None or importlib.util.find_spec("torch_npu"						) is None:
  				return False
  
  		import torch
  		import torch_npu  # noqa: F401
  
  		if check_device:
  				try:
  						# Will raise a RuntimeError if no NPU is found
  						_snake_case			 = torch.npu.device_count()
  						return torch.npu.is_available()
  				except RuntimeError:
  						return False
  		return hasattr(lowercase_		,							"npu"						) and torch.npu.is_available()
  
  
  
  
  
  
  @lru_cache
  def 	_lowercase	(			lowercase_=False						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if not parse_flag_from_env("ACCELERATE_USE_XPU"		,							default=lowercase_						):
  				return False
  		if is_ipex_available():
  				import torch
  
  				if is_torch_version("<="		,							"1.12"						):
  						return False
  		else:
  				return False
  
  		import intel_extension_for_pytorch  # noqa: F401
  
  		if check_device:
  				try:
  						# Will raise a RuntimeError if no XPU  is found
  						_snake_case			 = torch.xpu.device_count()
  						return torch.xpu.is_available()
  				except RuntimeError:
  						return False
  		return hasattr(lowercase_		,							"xpu"						) and torch.xpu.is_available()
  
  
  
  
  
  