

import math
from collections import defaultdict
from typing import List, Optional, Tuple, Union

import numpy as np
import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import randn_tensor
from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput



def  __lowerCAmelCase (  A    ,       A=0.999    ,       A="cosine"    ,       ):
	if alpha_transform_type == "cosine":

		def alpha_bar_fn(A      ):
			return math.cos((t + 0.008) / 1.008 * math.pi / 2      ) ** 2

	elif alpha_transform_type == "exp":

		def alpha_bar_fn(A      ):
			return math.exp(t * -12.0      )

	else:
		raise ValueError(F"Unsupported alpha_tranform_type: {alpha_transform_type}"      )

	UpperCAmelCase_    			=     []
	for i in range(A      ):
		UpperCAmelCase_    			=     i / num_diffusion_timesteps
		UpperCAmelCase_    			=     (i + 1) / num_diffusion_timesteps
		betas.append(min(1 - alpha_bar_fn(A      ) / alpha_bar_fn(A      )    ,       A      )      )
	return torch.tensor(A    ,       dtype=torch.floataa      )




class   __UpperCamelCase      (   lowercase       ,       lowercase  ):
	SCREAMING_SNAKE_CASE__      			=       [e.name for e in KarrasDiffusionSchedulers]
	SCREAMING_SNAKE_CASE__      			=       2




	@register_to_config
	def __init__(						self		:  List[Any]						,		lowerCAmelCase		:  int = 1_000						,		lowerCAmelCase		:  float = 0.00_085						,		lowerCAmelCase		:  float = 0.012						,		lowerCAmelCase		:  str = "linear"						,		lowerCAmelCase		:  Optional[Union[np.ndarray, List[float]]] = None						,		lowerCAmelCase		:  str = "epsilon"						,		lowerCAmelCase		:  str = "linspace"						,		lowerCAmelCase		:  int = 0						,		):


		'''simple docstring'''







		if trained_betas is not None:
			UpperCAmelCase_    			=     torch.tensor(lowerCAmelCase						,		dtype=torch.floataa							)
		elif beta_schedule == "linear":
			UpperCAmelCase_    			=     torch.linspace(lowerCAmelCase						,		lowerCAmelCase						,		lowerCAmelCase						,		dtype=torch.floataa							)
		elif beta_schedule == "scaled_linear":
			# this schedule is very specific to the latent diffusion model.
			UpperCAmelCase_    			=     (
			    torch.linspace(beta_start**0.5						,		beta_end**0.5						,		lowerCAmelCase						,		dtype=torch.floataa							) ** 2
			)
		elif beta_schedule == "squaredcos_cap_v2":
			# Glide cosine schedule
			UpperCAmelCase_    			=     betas_for_alpha_bar(lowerCAmelCase							)
		else:
			raise NotImplementedError(F"{beta_schedule} does is not implemented for {self.__class__}"							)

		UpperCAmelCase_    			=     1.0 - self.betas
		UpperCAmelCase_    			=     torch.cumprod(self.alphas						,		dim=0							)

		#  set all values
		self.set_timesteps(lowerCAmelCase						,		lowerCAmelCase						,		lowerCAmelCase							)




	def  __A   (						self		:  int						,		lowerCAmelCase		:  Tuple						,		lowerCAmelCase		:  Dict=None							):


		'''simple docstring'''







		if schedule_timesteps is None:
			UpperCAmelCase_    			=     self.timesteps

		UpperCAmelCase_    			=     (schedule_timesteps == timestep).nonzero()

		# The sigma index that is taken for the **very** first `step`
		# is always the second index (or the last index if there is only 1)
		# This way we can ensure we don't accidentally skip a sigma in
		# case we start in the middle of the denoising schedule (e.g. for image-to-image)
		if len(self._index_counter							) == 0:
			UpperCAmelCase_    			=     1 if len(lowerCAmelCase							) > 1 else 0
		else:
			UpperCAmelCase_    			=     timestep.cpu().item() if torch.is_tensor(lowerCAmelCase							) else timestep
			UpperCAmelCase_    			=     self._index_counter[timestep_int]

		return indices[pos].item()




	@property
	def  __A   (						self		:  List[Any]							):


		'''simple docstring'''







		if self.config.timestep_spacing in ["linspace", "trailing"]:
			return self.sigmas.max()

		return (self.sigmas.max() ** 2 + 1) ** 0.5




	def  __A   (						self		:  Optional[Any]						,		lowerCAmelCase		:  torch.FloatTensor						,		lowerCAmelCase		:  Union[float, torch.FloatTensor]						,		):


		'''simple docstring'''







		UpperCAmelCase_    			=     self.index_for_timestep(lowerCAmelCase							)

		if self.state_in_first_order:
			UpperCAmelCase_    			=     self.sigmas[step_index]
		else:
			UpperCAmelCase_    			=     self.sigmas_interpol[step_index - 1]

		UpperCAmelCase_    			=     sample / ((sigma**2 + 1) ** 0.5)
		return sample




	def  __A   (						self		:  Optional[Any]						,		lowerCAmelCase		:  int						,		lowerCAmelCase		:  Union[str, torch.device] = None						,		lowerCAmelCase		:  Optional[int] = None						,		):


		'''simple docstring'''







		UpperCAmelCase_    			=     num_inference_steps

		UpperCAmelCase_    			=     num_train_timesteps or self.config.num_train_timesteps

		# "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
		if self.config.timestep_spacing == "linspace":
			UpperCAmelCase_    			=     np.linspace(0						,		num_train_timesteps - 1						,		lowerCAmelCase						,		dtype=lowerCAmelCase							)[::-1].copy()
		elif self.config.timestep_spacing == "leading":
			UpperCAmelCase_    			=     num_train_timesteps // self.num_inference_steps
			# creates integer timesteps by multiplying by ratio
			# casting to int to avoid issues when num_inference_step is power of 3
			UpperCAmelCase_    			=     (np.arange(0						,		lowerCAmelCase							) * step_ratio).round()[::-1].copy().astype(lowerCAmelCase							)
			timesteps += self.config.steps_offset
		elif self.config.timestep_spacing == "trailing":
			UpperCAmelCase_    			=     num_train_timesteps / self.num_inference_steps
			# creates integer timesteps by multiplying by ratio
			# casting to int to avoid issues when num_inference_step is power of 3
			UpperCAmelCase_    			=     (np.arange(lowerCAmelCase						,		0						,		-step_ratio							)).round().copy().astype(lowerCAmelCase							)
			timesteps -= 1
		else:
			raise ValueError(
			    F"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'."							)

		UpperCAmelCase_    			=     np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5							)
		UpperCAmelCase_    			=     torch.from_numpy(np.log(lowerCAmelCase							)							).to(lowerCAmelCase							)

		UpperCAmelCase_    			=     np.interp(lowerCAmelCase						,		np.arange(0						,		len(lowerCAmelCase							)							)						,		lowerCAmelCase							)
		UpperCAmelCase_    			=     np.concatenate([sigmas, [0.0]]							).astype(np.floataa							)
		UpperCAmelCase_    			=     torch.from_numpy(lowerCAmelCase							).to(device=lowerCAmelCase							)

		# compute up and down sigmas
		UpperCAmelCase_    			=     sigmas.roll(-1							)
		UpperCAmelCase_    			=     0.0
		UpperCAmelCase_    			=     (sigmas_next**2 * (sigmas**2 - sigmas_next**2) / sigmas**2) ** 0.5
		UpperCAmelCase_    			=     (sigmas_next**2 - sigmas_up**2) ** 0.5
		UpperCAmelCase_    			=     0.0

		# compute interpolated sigmas
		UpperCAmelCase_    			=     sigmas.log().lerp(sigmas_down.log()						,		0.5							).exp()
		UpperCAmelCase_    			=     0.0

		# set sigmas
		UpperCAmelCase_    			=     torch.cat([sigmas[:1], sigmas[1:].repeat_interleave(2							), sigmas[-1:]]							)
		UpperCAmelCase_    			=     torch.cat(
		    [sigmas_interpol[:1], sigmas_interpol[1:].repeat_interleave(2							), sigmas_interpol[-1:]]							)
		UpperCAmelCase_    			=     torch.cat([sigmas_up[:1], sigmas_up[1:].repeat_interleave(2							), sigmas_up[-1:]]							)
		UpperCAmelCase_    			=     torch.cat([sigmas_down[:1], sigmas_down[1:].repeat_interleave(2							), sigmas_down[-1:]]							)

		if str(lowerCAmelCase							).startswith("mps"							):
			# mps does not support float64
			UpperCAmelCase_    			=     torch.from_numpy(lowerCAmelCase							).to(lowerCAmelCase						,		dtype=torch.floataa							)
		else:
			UpperCAmelCase_    			=     torch.from_numpy(lowerCAmelCase							).to(lowerCAmelCase							)

		UpperCAmelCase_    			=     self.sigma_to_t(lowerCAmelCase							).to(lowerCAmelCase						,		dtype=timesteps.dtype							)
		UpperCAmelCase_    			=     torch.stack((timesteps_interpol[:-2, None], timesteps[1:, None])						,		dim=-1							).flatten()

		UpperCAmelCase_    			=     torch.cat([timesteps[:1], interleaved_timesteps]							)

		UpperCAmelCase_    			=     None

		# for exp beta schedules, such as the one for `pipeline_shap_e.py`
		# we need an index counter
		UpperCAmelCase_    			=     defaultdict(lowerCAmelCase							)




	def  __A   (						self		:  Union[str, Any]						,		lowerCAmelCase		:  Dict							):


		'''simple docstring'''







		UpperCAmelCase_    			=     sigma.log()

		# get distribution
		UpperCAmelCase_    			=     log_sigma - self.log_sigmas[:, None]

		# get sigmas range
		UpperCAmelCase_    			=     dists.ge(0							).cumsum(dim=0							).argmax(dim=0							).clamp(max=self.log_sigmas.shape[0] - 2							)
		UpperCAmelCase_    			=     low_idx + 1

		UpperCAmelCase_    			=     self.log_sigmas[low_idx]
		UpperCAmelCase_    			=     self.log_sigmas[high_idx]

		# interpolate sigmas
		UpperCAmelCase_    			=     (low - log_sigma) / (low - high)
		UpperCAmelCase_    			=     w.clamp(0						,		1							)

		# transform interpolation to time range
		UpperCAmelCase_    			=     (1 - w) * low_idx + w * high_idx
		UpperCAmelCase_    			=     t.view(sigma.shape							)
		return t




	@property
	def  __A   (						self		:  Optional[Any]							):


		'''simple docstring'''







		return self.sample is None




	def  __A   (						self		:  Optional[int]						,		lowerCAmelCase		:  Union[torch.FloatTensor, np.ndarray]						,		lowerCAmelCase		:  Union[float, torch.FloatTensor]						,		lowerCAmelCase		:  Union[torch.FloatTensor, np.ndarray]						,		lowerCAmelCase		:  Optional[torch.Generator] = None						,		lowerCAmelCase		:  bool = True						,		):


		'''simple docstring'''







		UpperCAmelCase_    			=     self.index_for_timestep(lowerCAmelCase							)

		# advance index counter by 1
		UpperCAmelCase_    			=     timestep.cpu().item() if torch.is_tensor(lowerCAmelCase							) else timestep
		self._index_counter[timestep_int] += 1

		if self.state_in_first_order:
			UpperCAmelCase_    			=     self.sigmas[step_index]
			UpperCAmelCase_    			=     self.sigmas_interpol[step_index]
			UpperCAmelCase_    			=     self.sigmas_up[step_index]
			UpperCAmelCase_    			=     self.sigmas_down[step_index - 1]
		else:
			# 2nd order / KPDM2's method
			UpperCAmelCase_    			=     self.sigmas[step_index - 1]
			UpperCAmelCase_    			=     self.sigmas_interpol[step_index - 1]
			UpperCAmelCase_    			=     self.sigmas_up[step_index - 1]
			UpperCAmelCase_    			=     self.sigmas_down[step_index - 1]

		# currently only gamma=0 is supported. This usually works best anyways.
		# We can support gamma in the future but then need to scale the timestep before
		# passing it to the model which requires a change in API
		UpperCAmelCase_    			=     0
		UpperCAmelCase_    			=     sigma * (gamma + 1)  # Note: sigma_hat == sigma for now

		UpperCAmelCase_    			=     model_output.device
		UpperCAmelCase_    			=     randn_tensor(model_output.shape						,		dtype=model_output.dtype						,		device=lowerCAmelCase						,		generator=lowerCAmelCase							)

		# 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
		if self.config.prediction_type == "epsilon":
			UpperCAmelCase_    			=     sigma_hat if self.state_in_first_order else sigma_interpol
			UpperCAmelCase_    			=     sample - sigma_input * model_output
		elif self.config.prediction_type == "v_prediction":
			UpperCAmelCase_    			=     sigma_hat if self.state_in_first_order else sigma_interpol
			UpperCAmelCase_    			=     model_output * (-sigma_input / (sigma_input**2 + 1) ** 0.5) + (
			    sample / (sigma_input**2 + 1)
			)
		elif self.config.prediction_type == "sample":
			raise NotImplementedError("prediction_type not implemented yet: sample"							)
		else:
			raise ValueError(
			    F"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`"							)

		if self.state_in_first_order:
			# 2. Convert to an ODE derivative for 1st order
			UpperCAmelCase_    			=     (sample - pred_original_sample) / sigma_hat
			# 3. delta timestep
			UpperCAmelCase_    			=     sigma_interpol - sigma_hat

			# store for 2nd order step
			UpperCAmelCase_    			=     sample
			UpperCAmelCase_    			=     dt
			UpperCAmelCase_    			=     sample + derivative * dt
		else:
			# DPM-Solver-2
			# 2. Convert to an ODE derivative for 2nd order
			UpperCAmelCase_    			=     (sample - pred_original_sample) / sigma_interpol
			# 3. delta timestep
			UpperCAmelCase_    			=     sigma_down - sigma_hat

			UpperCAmelCase_    			=     self.sample
			UpperCAmelCase_    			=     None

			UpperCAmelCase_    			=     sample + derivative * dt
			UpperCAmelCase_    			=     prev_sample + noise * sigma_up

		if not return_dict:
			return (prev_sample,)

		return SchedulerOutput(prev_sample=lowerCAmelCase							)




	def  __A   (						self		:  Union[str, Any]						,		lowerCAmelCase		:  torch.FloatTensor						,		lowerCAmelCase		:  torch.FloatTensor						,		lowerCAmelCase		:  torch.FloatTensor						,		):


		'''simple docstring'''







		UpperCAmelCase_    			=     self.sigmas.to(device=original_samples.device						,		dtype=original_samples.dtype							)
		if original_samples.device.type == "mps" and torch.is_floating_point(lowerCAmelCase							):
			# mps does not support float64
			UpperCAmelCase_    			=     self.timesteps.to(original_samples.device						,		dtype=torch.floataa							)
			UpperCAmelCase_    			=     timesteps.to(original_samples.device						,		dtype=torch.floataa							)
		else:
			UpperCAmelCase_    			=     self.timesteps.to(original_samples.device							)
			UpperCAmelCase_    			=     timesteps.to(original_samples.device							)

		UpperCAmelCase_    			=     [self.index_for_timestep(lowerCAmelCase						,		lowerCAmelCase							) for t in timesteps]

		UpperCAmelCase_    			=     sigmas[step_indices].flatten()
		while len(sigma.shape							) < len(original_samples.shape							):
			UpperCAmelCase_    			=     sigma.unsqueeze(-1							)

		UpperCAmelCase_    			=     original_samples + noise * sigma
		return noisy_samples




	def __len__(						self		:  Tuple							):


		'''simple docstring'''







		return self.config.num_train_timesteps


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

				_a:      List[Any]								=     True
except ImportError:
				_a:      int								=     False


# Cache this result has it's a C FFI call which can be pretty time-consuming
_a:      str								=     torch.distributed.is_available()



def  __lowerCAmelCase (  A      ):
	# Check we're not importing a "pkg_name" directory somewhere but the actual library by trying to grab the version
	UpperCAmelCase_    			=     importlib.util.find_spec(A      ) is not None
	if package_exists:
		try:
			UpperCAmelCase_    			=     importlib.metadata.metadata(A      )
			return True
		except importlib.metadata.PackageNotFoundError:
			return False



def  __lowerCAmelCase (  ):
	return _torch_distributed_available



def  __lowerCAmelCase (  ):
	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  __lowerCAmelCase (  ):
	return importlib.metadata.version("oneccl_bind_pt"      )



def  __lowerCAmelCase (  ):
	return _is_package_available("transformer_engine"      )



@lru_cache
def  __lowerCAmelCase (  A=True      ):
	# Due to bugs on the amp series GPUs, we disable torch-xla on them
	if torch.cuda.is_available():
		return False
	if _tpu_available and check_device:
		try:
			# Will raise a RuntimeError if no XLA configuration is found
			UpperCAmelCase_    			=     xm.xla_device()
			return True
		except RuntimeError:
			return False
	return _tpu_available



def  __lowerCAmelCase (  ):
	return _is_package_available("deepspeed"      )



def  __lowerCAmelCase (  A=False      ):
	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  __lowerCAmelCase (  ):
	UpperCAmelCase_    			=     _is_package_available("bitsandbytes"      )
	if package_exists:
		UpperCAmelCase_    			=     version.parse(importlib.metadata.version("bitsandbytes"      )      )
		return compare_versions(A    ,       ">="    ,       "0.39.0"      )
	return False



def  __lowerCAmelCase (  ):
	UpperCAmelCase_    			=     _is_package_available("bitsandbytes"      )
	if package_exists:
		UpperCAmelCase_    			=     version.parse(importlib.metadata.version("bitsandbytes"      )      )
		return compare_versions(A    ,       ">="    ,       "0.37.2"      )
	return False



def  __lowerCAmelCase (  ):
	return _is_package_available("bitsandbytes"      )



def  __lowerCAmelCase (  ):
	if strtobool(os.environ.get("ACCELERATE_USE_MEGATRON_LM"    ,       "False"      )      ) == 1:
		UpperCAmelCase_    			=     importlib.util.find_spec("megatron"      ) is not None
		if package_exists:
			try:
				UpperCAmelCase_    			=     parse(importlib.metadata.version("megatron-lm"      )      )
				return compare_versions(A    ,       ">="    ,       "2.2.0"      )
			except Exception as e:
				warnings.warn(F"Parse Megatron version failed. Exception:{e}"      )
				return False



def  __lowerCAmelCase (  ):
	return _is_package_available("safetensors"      )



def  __lowerCAmelCase (  ):
	return _is_package_available("transformers"      )



def  __lowerCAmelCase (  ):
	return _is_package_available("datasets"      )



def  __lowerCAmelCase (  ):
	return _is_package_available("aim"      )



def  __lowerCAmelCase (  ):
	return _is_package_available("tensorboard"      ) or _is_package_available("tensorboardX"      )



def  __lowerCAmelCase (  ):
	return _is_package_available("wandb"      )



def  __lowerCAmelCase (  ):
	return _is_package_available("comet_ml"      )



def  __lowerCAmelCase (  ):
	return _is_package_available("boto3"      )



def  __lowerCAmelCase (  ):
	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"    ,       A      )
		return parse_flag_from_env("ACCELERATE_ENABLE_RICH"    ,       A      )
	return False



def  __lowerCAmelCase (  ):
	return _is_package_available("sagemaker"      )



def  __lowerCAmelCase (  ):
	return _is_package_available("tqdm"      )



def  __lowerCAmelCase (  ):
	return _is_package_available("mlflow"      )



def  __lowerCAmelCase (  ):
	return is_torch_version(">="    ,       "1.12"      ) and torch.backends.mps.is_available() and torch.backends.mps.is_built()



def  __lowerCAmelCase (  ):
	def get_major_and_minor_from_version(A      ):
		return str(version.parse(A      ).major      ) + "." + str(version.parse(A      ).minor      )

	UpperCAmelCase_    			=     importlib.metadata.version("torch"      )
	if importlib.util.find_spec("intel_extension_for_pytorch"      ) is None:
		return False
	UpperCAmelCase_    			=     "N/A"
	try:
		UpperCAmelCase_    			=     importlib.metadata.version("intel_extension_for_pytorch"      )
	except importlib.metadata.PackageNotFoundError:
		return False
	UpperCAmelCase_    			=     get_major_and_minor_from_version(_torch_version      )
	UpperCAmelCase_    			=     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  __lowerCAmelCase (  A=False      ):
	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
			UpperCAmelCase_    			=     torch.npu.device_count()
			return torch.npu.is_available()
		except RuntimeError:
			return False
	return hasattr(A    ,       "npu"      ) and torch.npu.is_available()




@lru_cache
def  __lowerCAmelCase (  A=False      ):
	if not parse_flag_from_env("ACCELERATE_USE_XPU"    ,       default=A      ):
		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
			UpperCAmelCase_    			=     torch.xpu.device_count()
			return torch.xpu.is_available()
		except RuntimeError:
			return False
	return hasattr(A    ,       "xpu"      ) and torch.xpu.is_available()