
# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
# and https://github.com/hojonathanho/diffusion

from dataclasses import dataclass
from typing import Optional, Tuple, Union

import flax
import jax.numpy as jnp

from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils_flax import (
    CommonSchedulerState,
    FlaxKarrasDiffusionSchedulers,
    FlaxSchedulerMixin,
    FlaxSchedulerOutput,
    add_noise_common,
    get_velocity_common,
)






@flax.struct.dataclass
class 							SCREAMING_SNAKE_CASE__     :

							"""simple docstring"""






							_snake_case  		=	42
							_snake_case  		=	42

							# setable values
							_snake_case  		=	42
							_snake_case  		=	42
							_snake_case  		=	None



							@classmethod
							def      A__		(	cls  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> Union[str, Any]:





													'''simple docstring'''






													return cls(
													    common=SCREAMING_SNAKE_CASE_  ,						final_alpha_cumprod=SCREAMING_SNAKE_CASE_  ,						init_noise_sigma=SCREAMING_SNAKE_CASE_  ,						timesteps=SCREAMING_SNAKE_CASE_  ,						)






@dataclass
class 							SCREAMING_SNAKE_CASE__     (				SCREAMING_SNAKE_CASE_     ):

							"""simple docstring"""






							_snake_case  		=	42






class 							SCREAMING_SNAKE_CASE__     (				SCREAMING_SNAKE_CASE_ ,		SCREAMING_SNAKE_CASE_     ):

							"""simple docstring"""






							_snake_case  		=	[e.name for e in FlaxKarrasDiffusionSchedulers]

							_snake_case  		=	42



							@property
							def      A__		(	self      )-> List[Any]:





													'''simple docstring'''






													return True



							@register_to_config
							def __init__(	self  ,						SCREAMING_SNAKE_CASE_ = 1000  ,						SCREAMING_SNAKE_CASE_ = 0.0_0_0_1  ,						SCREAMING_SNAKE_CASE_ = 0.0_2  ,						SCREAMING_SNAKE_CASE_ = "linear"  ,						SCREAMING_SNAKE_CASE_ = None  ,						SCREAMING_SNAKE_CASE_ = True  ,						SCREAMING_SNAKE_CASE_ = 0  ,						SCREAMING_SNAKE_CASE_ = "epsilon"  ,						SCREAMING_SNAKE_CASE_ = jnp.floataa  ,						)-> str:





													'''simple docstring'''






													__UpperCamelCase  	=      dtype



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_ = None      )-> DDIMSchedulerState:





													'''simple docstring'''






													if common is None:
																			__UpperCamelCase  	=      CommonSchedulerState.create(self      )

													# At every step in ddim, we are looking into the previous alphas_cumprod
													# For the final step, there is no previous alphas_cumprod because we are already at 0
													# `set_alpha_to_one` decides whether we set this parameter simply to one or
													# whether we use the final alpha of the "non-previous" one.
													__UpperCamelCase  	=      (
													    jnp.array(1.0  ,						dtype=self.dtype      ) if self.config.set_alpha_to_one else common.alphas_cumprod[0]
													)

													# standard deviation of the initial noise distribution
													__UpperCamelCase  	=      jnp.array(1.0  ,						dtype=self.dtype      )

													__UpperCamelCase  	=      jnp.arange(0  ,						self.config.num_train_timesteps      ).round()[::-1]

													return DDIMSchedulerState.create(
													    common=SCREAMING_SNAKE_CASE_  ,						final_alpha_cumprod=SCREAMING_SNAKE_CASE_  ,						init_noise_sigma=SCREAMING_SNAKE_CASE_  ,						timesteps=SCREAMING_SNAKE_CASE_  ,						)



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_ = None      )-> jnp.ndarray:





													'''simple docstring'''






													return sample



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_ = ()      )-> DDIMSchedulerState:





													'''simple docstring'''






													__UpperCamelCase  	=      self.config.num_train_timesteps // num_inference_steps
													# creates integer timesteps by multiplying by ratio
													# rounding to avoid issues when num_inference_step is power of 3
													__UpperCamelCase  	=      (jnp.arange(0  ,						SCREAMING_SNAKE_CASE_      ) * step_ratio).round()[::-1] + self.config.steps_offset

													return state.replace(
													    num_inference_steps=SCREAMING_SNAKE_CASE_  ,						timesteps=SCREAMING_SNAKE_CASE_  ,						)



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_      )-> Optional[int]:





													'''simple docstring'''






													__UpperCamelCase  	=      state.common.alphas_cumprod[timestep]
													__UpperCamelCase  	=      jnp.where(
													    prev_timestep >= 0  ,						state.common.alphas_cumprod[prev_timestep]  ,						state.final_alpha_cumprod      )
													__UpperCamelCase  	=      1 - alpha_prod_t
													__UpperCamelCase  	=      1 - alpha_prod_t_prev

													__UpperCamelCase  	=      (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)

													return variance



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_ = 0.0  ,						SCREAMING_SNAKE_CASE_ = True  ,						)-> Union[FlaxDDIMSchedulerOutput, Tuple]:





													'''simple docstring'''






													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'''      )

													# See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf
													# Ideally, read DDIM paper in-detail understanding

													# Notation (<variable name> -> <name in paper>
													# - pred_noise_t -> e_theta(x_t, t)
													# - pred_original_sample -> f_theta(x_t, t) or x_0
													# - std_dev_t -> sigma_t
													# - eta -> η
													# - pred_sample_direction -> "direction pointing to x_t"
													# - pred_prev_sample -> "x_t-1"

													# 1. get previous step value (=t-1)
													__UpperCamelCase  	=      timestep - self.config.num_train_timesteps // state.num_inference_steps

													__UpperCamelCase  	=      state.common.alphas_cumprod
													__UpperCamelCase  	=      state.final_alpha_cumprod

													# 2. compute alphas, betas
													__UpperCamelCase  	=      alphas_cumprod[timestep]
													__UpperCamelCase  	=      jnp.where(prev_timestep >= 0  ,						alphas_cumprod[prev_timestep]  ,						SCREAMING_SNAKE_CASE_      )

													__UpperCamelCase  	=      1 - alpha_prod_t

													# 3. compute predicted original sample from predicted noise also called
													# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
													if self.config.prediction_type == "epsilon":
																			__UpperCamelCase  	=      (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
																			__UpperCamelCase  	=      model_output
													elif self.config.prediction_type == "sample":
																			__UpperCamelCase  	=      model_output
																			__UpperCamelCase  	=      (sample - alpha_prod_t ** 0.5 * pred_original_sample) / beta_prod_t ** 0.5
													elif self.config.prediction_type == "v_prediction":
																			__UpperCamelCase  	=      (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
																			__UpperCamelCase  	=      (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
													else:
																			raise ValueError(
																			    F"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"
																			    ''' `v_prediction`'''      )

													# 4. compute variance: "sigma_t(η)" -> see formula (16)
													# σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
													__UpperCamelCase  	=      self._get_variance(SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_      )
													__UpperCamelCase  	=      eta * variance ** 0.5

													# 5. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
													__UpperCamelCase  	=      (1 - alpha_prod_t_prev - std_dev_t**2) ** 0.5 * pred_epsilon

													# 6. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
													__UpperCamelCase  	=      alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction

													if not return_dict:
																			return (prev_sample, state)

													return FlaxDDIMSchedulerOutput(prev_sample=SCREAMING_SNAKE_CASE_  ,						state=SCREAMING_SNAKE_CASE_      )



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> jnp.ndarray:





													'''simple docstring'''






													return add_noise_common(state.common  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_      )



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> jnp.ndarray:





													'''simple docstring'''






													return get_velocity_common(state.common  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_      )



							def __len__(	self      )-> Union[str, Any]:





													'''simple docstring'''






													return self.config.num_train_timesteps









import json
import os
from collections.abc import Mapping
from typing import Dict, List, Optional, Union

import numpy as np
import torch

from .imports import is_safetensors_available


def 						_UpperCAmelCase							( snake_case      ,		snake_case      ,		snake_case      ,		snake_case=None     ):





					"""simple docstring"""

					_lowerCAmelCase         =       None
					# Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
					if str(weight.dtype     ) == "torch.bfloat16":
										# Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
										_lowerCAmelCase         =       weight.view(torch.intaa     )
										_lowerCAmelCase         =       """bfloat16"""
					_lowerCAmelCase         =       weight.cpu().numpy()
					_lowerCAmelCase         =       os.path.join(snake_case      ,		F'{weight_name}.dat'     )
					if index is not None:
										if dtype is None:
															_lowerCAmelCase         =       str(array.dtype     )
										_lowerCAmelCase         =       {"""dtype""": dtype, """shape""": list(array.shape     )}
					if array.ndim == 0:
										_lowerCAmelCase         =       array[None]
					_lowerCAmelCase         =       np.memmap(snake_case      ,		dtype=array.dtype      ,		mode="""w+"""      ,		shape=array.shape     )
					_lowerCAmelCase         =       array[:]
					file_array.flush()
					return index


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

					_lowerCAmelCase         =       tuple(weight_info["""shape"""]     )
					if shape == ():
										# NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
										_lowerCAmelCase         =       (1,)

					_lowerCAmelCase         =       weight_info["""dtype"""]
					if dtype == "bfloat16":
										# NumPy does not support bfloat16 so this was saved as a int16
										_lowerCAmelCase         =       """int16"""

					_lowerCAmelCase         =       np.memmap(snake_case      ,		dtype=snake_case      ,		shape=snake_case      ,		mode="""r"""     )

					if len(weight_info["""shape"""]     ) == 0:
										_lowerCAmelCase         =       weight[0]
					_lowerCAmelCase         =       torch.tensor(snake_case     )
					if weight_info["dtype"] == "bfloat16":
										_lowerCAmelCase         =       weight.view(torch.bfloataa     )

					return weight


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

					if index is None or len(snake_case     ) == 0:
										# Nothing to save
										return

					_lowerCAmelCase         =       os.path.join(snake_case      ,		"""index.json"""     )
					if os.path.isfile(snake_case     ):
										with open(snake_case      ,		"""r"""      ,		encoding="""utf-8"""     ) as f:
															_lowerCAmelCase         =       json.load(snake_case     )
					else:
										_lowerCAmelCase         =       {}
					current_index.update(snake_case     )

					with open(snake_case      ,		"""w"""      ,		encoding="""utf-8"""     ) as f:
										json.dump(snake_case      ,		snake_case      ,		indent=2     )


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

					os.makedirs(snake_case      ,		exist_ok=snake_case     )
					_lowerCAmelCase         =       {}
					for name, parameter in state_dict.items():
										_lowerCAmelCase         =       offload_weight(snake_case      ,		snake_case      ,		snake_case      ,		index=snake_case     )

					# Update index
					save_offload_index(snake_case      ,		snake_case     )





class    __lowerCAmelCase					(	lowerCamelCase__			):




				def __init__(				self		,							_snake_case		,							_snake_case						):



									"""simple docstring"""





									_lowerCAmelCase         =       dataset
									_lowerCAmelCase         =       prefix




				def __getitem__(				self		,							_snake_case						):



									"""simple docstring"""





									return self.dataset[F'{self.prefix}{key}']




				def __iter__(				self						):



									"""simple docstring"""





									return iter([key for key in self.dataset if key.startswith(self.prefix						)]						)




				def __len__(				self						):



									"""simple docstring"""





									return len(self.dataset						)





class    __lowerCAmelCase					(	lowerCamelCase__			):




				def __init__(				self		,							_snake_case = None		,							_snake_case = None		,							_snake_case = None		,							_snake_case=None		,							):



									"""simple docstring"""





									if state_dict is None and save_folder is None:
														raise ValueError("""Need either a `state_dict` or a `save_folder` containing offloaded weights."""						)

									_lowerCAmelCase         =       {} if state_dict is None else state_dict
									_lowerCAmelCase         =       save_folder
									if index is None and save_folder is not None:
														with open(os.path.join(_snake_case		,							"""index.json"""						)						) as f:
																			_lowerCAmelCase         =       json.load(_snake_case						)
									_lowerCAmelCase         =       {} if index is None else index
									_lowerCAmelCase         =       list(self.state_dict.keys()						)
									self.all_keys.extend([key for key in self.index if key not in self.all_keys]						)
									_lowerCAmelCase         =       device




				def __getitem__(				self		,							_snake_case						):



									"""simple docstring"""





									if key in self.state_dict:
														return self.state_dict[key]
									_lowerCAmelCase         =       self.index[key]
									if weight_info.get("""safetensors_file"""						) is not None:
														if not is_safetensors_available():
																			raise ImportError("""These offloaded weights require the use of safetensors: `pip install safetensors`."""						)

														from safetensors import safe_open

														_lowerCAmelCase         =       """cpu""" if self.device is None else self.device
														with safe_open(weight_info["""safetensors_file"""]		,							framework="""pt"""		,							device=_snake_case						) as f:
																			_lowerCAmelCase         =       f.get_tensor(weight_info.get("""weight_name"""		,							_snake_case						)						)

														if "dtype" in weight_info:
																			return tensor.to(getattr(_snake_case		,							weight_info["""dtype"""]						)						)
														else:
																			return tensor

									_lowerCAmelCase         =       os.path.join(self.save_folder		,							F'{key}.dat'						)
									return load_offloaded_weight(_snake_case		,							_snake_case						)




				def __iter__(				self						):



									"""simple docstring"""





									return iter(self.all_keys						)




				def __len__(				self						):



									"""simple docstring"""





									return len(self.all_keys						)


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

					_lowerCAmelCase         =       {}
					for module_name in submodule_names:
										# We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
										# submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
										result.update(
										    {
										        key: param
										        for key, param in state_dict.items()
										        if key == module_name or key.startswith(module_name + """."""     )
										    }     )
					return result



