






# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
# and https://github.com/hojonathanho/diffusion
import math
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union

import numpy as np
import torch

from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.schedulers.scheduling_utils import SchedulerMixin
from diffusers.utils import BaseOutput, deprecate


@dataclass
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DDIM
class    _lowerCAmelCase			(		lowercase_      ):



				"""simple docstring"""



				_lowercase				:   torch.FloatTensor
				_lowercase				:   Optional[torch.FloatTensor]      =  None


def _UpperCAmelCase   (	a						:	Optional[int]  ,					a						:	Optional[int]=0.999  ,					a						:	str="cosine"  ,					):
							if alpha_transform_type == "cosine":

														def alpha_bar_fn(a						:	List[Any] ):
																					return math.cos((t + 0.008) / 1.008 * math.pi / 2 ) ** 2

							elif alpha_transform_type == "exp":

														def alpha_bar_fn(a						:	Union[str, Any] ):
																					return math.exp(t * -12.0 )

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

							snake_case__	  =       []
							for i in range(a ):
														snake_case__	  =       i / num_diffusion_timesteps
														snake_case__	  =       (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    _lowerCAmelCase			(		lowercase_     ,      lowercase_      ):



				"""simple docstring"""



				_lowercase				:   Optional[Any]      =  1



				@register_to_config
				def __init__(      self    :    Optional[int]			,		UpperCamelCase__    :    int = 1_0_0_0			,		UpperCamelCase__    :    float = 0.00_01			,		UpperCamelCase__    :    float = 0.02			,		UpperCamelCase__    :    str = "linear"			,		UpperCamelCase__    :    Optional[Union[np.ndarray, List[float]]] = None			,		UpperCamelCase__    :    bool = True			,		UpperCamelCase__    :    bool = True			,		UpperCamelCase__    :    int = 0			,		UpperCamelCase__    :    str = "epsilon"			,		UpperCamelCase__    :    float = 1.0			,		**UpperCamelCase__    :    Any			,		):






											'''simple docstring'''







											if kwargs.get("""set_alpha_to_one"""			,		UpperCamelCase__) is not None:
																		snake_case__	  =       (
																		    """The `set_alpha_to_one` argument is deprecated. Please use `set_alpha_to_zero` instead."""
																		)
																		deprecate("""set_alpha_to_one"""			,		"""1.0.0"""			,		UpperCamelCase__			,		standard_warn=UpperCamelCase__)
																		snake_case__	  =       kwargs["""set_alpha_to_one"""]
											if trained_betas is not None:
																		snake_case__	  =       torch.tensor(UpperCamelCase__			,		dtype=torch.floataa)
											elif beta_schedule == "linear":
																		snake_case__	  =       torch.linspace(UpperCamelCase__			,		UpperCamelCase__			,		UpperCamelCase__			,		dtype=torch.floataa)
											elif beta_schedule == "scaled_linear":
																		# this schedule is very specific to the latent diffusion model.
																		snake_case__	  =       (
																		    torch.linspace(beta_start**0.5			,		beta_end**0.5			,		UpperCamelCase__			,		dtype=torch.floataa) ** 2
																		)
											elif beta_schedule == "squaredcos_cap_v2":
																		# Glide cosine schedule
																		snake_case__	  =       betas_for_alpha_bar(UpperCamelCase__)
											else:
																		raise NotImplementedError(F'''{beta_schedule} does is not implemented for {self.__class__}''')

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

											# At every step in inverted ddim, we are looking into the next alphas_cumprod
											# For the final step, there is no next alphas_cumprod, and the index is out of bounds
											# `set_alpha_to_zero` decides whether we set this parameter simply to zero
											# in this case, self.step() just output the predicted noise
											# or whether we use the final alpha of the "non-previous" one.
											snake_case__	  =       torch.tensor(0.0) if set_alpha_to_zero else self.alphas_cumprod[-1]

											# standard deviation of the initial noise distribution
											snake_case__	  =       1.0

											# setable values
											snake_case__	  =       None
											snake_case__	  =       torch.from_numpy(np.arange(0			,		UpperCamelCase__).copy().astype(np.intaa))



				def 					__magic_name__							(      self    :    Union[str, Any]			,		UpperCamelCase__    :    torch.FloatTensor			,		UpperCamelCase__    :    Optional[int] = None):






											'''simple docstring'''







											return sample



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






											'''simple docstring'''







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

											snake_case__	  =       num_inference_steps
											snake_case__	  =       self.config.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
											snake_case__	  =       (np.arange(0			,		UpperCamelCase__) * step_ratio).round().copy().astype(np.intaa)
											snake_case__	  =       torch.from_numpy(UpperCamelCase__).to(UpperCamelCase__)
											self.timesteps += self.config.steps_offset



				def 					__magic_name__							(      self    :    List[Any]			,		UpperCamelCase__    :    torch.FloatTensor			,		UpperCamelCase__    :    int			,		UpperCamelCase__    :    torch.FloatTensor			,		UpperCamelCase__    :    float = 0.0			,		UpperCamelCase__    :    bool = False			,		UpperCamelCase__    :    Optional[torch.FloatTensor] = None			,		UpperCamelCase__    :    bool = True			,		):






											'''simple docstring'''







											snake_case__	  =       timestep + self.config.num_train_timesteps // self.num_inference_steps

											# 2. compute alphas, betas
											# change original implementation to exactly match noise levels for analogous forward process
											snake_case__	  =       self.alphas_cumprod[timestep]
											snake_case__	  =       (
											    self.alphas_cumprod[prev_timestep]
											    if prev_timestep < self.config.num_train_timesteps
											    else self.final_alpha_cumprod
											)

											snake_case__	  =       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":
																		snake_case__	  =       (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
																		snake_case__	  =       model_output
											elif self.config.prediction_type == "sample":
																		snake_case__	  =       model_output
																		snake_case__	  =       (sample - alpha_prod_t ** 0.5 * pred_original_sample) / beta_prod_t ** 0.5
											elif self.config.prediction_type == "v_prediction":
																		snake_case__	  =       (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
																		snake_case__	  =       (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. Clip or threshold "predicted x_0"
											if self.config.clip_sample:
																		snake_case__	  =       pred_original_sample.clamp(
																		    -self.config.clip_sample_range			,		self.config.clip_sample_range)

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

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

											if not return_dict:
																		return (prev_sample, pred_original_sample)
											return DDIMSchedulerOutput(prev_sample=UpperCamelCase__			,		pred_original_sample=UpperCamelCase__)



				def __len__(      self    :    Tuple):






											'''simple docstring'''







											return self.config.num_train_timesteps








from typing import List, Optional, Union

import numpy as np
import tensorflow as tf

from .utils import logging


a__       =     logging.get_logger(__name__)

def _UpperCAmelCase   (	a						:	Union[tf.Tensor, np.ndarray] ):
							if isinstance(a  ,					np.ndarray ):
														return list(tensor.shape )

							snake_case__	  =       tf.shape(a )

							if tensor.shape == tf.TensorShape(a ):
														return dynamic

							snake_case__	  =       tensor.shape.as_list()

							return [dynamic[i] if s is None else s for i, s in enumerate(a )]

def _UpperCAmelCase   (	a						:	tf.Tensor  ,					a						:	Optional[int] = None  ,					a						:	Optional[str] = None ):
							return tf.nn.softmax(logits=logits + 1e-9  ,					axis=a  ,					name=a )

def _UpperCAmelCase   (	a						:	Optional[int]  ,					a						:	Union[str, Any]  ,					a						:	Dict  ,					a						:	int=1e-5  ,					a						:	Dict=-1 ):
							# This is a very simplified functional layernorm, designed to duplicate
							# the functionality of PyTorch nn.functional.layer_norm when this is needed to port
							# models in Transformers.

							if weight.shape.rank != 1 or bias.shape.rank != 1 or not isinstance(a  ,					a ):
														raise NotImplementedError("""Only 1D weight and bias tensors are supported for now, with only a single axis.""" )

							# Get mean and variance on the axis to be normalized
							snake_case__  , snake_case__	  =       tf.nn.moments(a  ,					axes=[axis]  ,					keepdims=a )

							if axis != -1:
														# Reshape scale and weight to have the same rank as inputs, but with 1 dimensions
														# on every dimension except axis
														snake_case__	  =       [1] * inputs.shape.rank
														snake_case__	  =       shape_list(a )[axis]
														snake_case__	  =       tf.reshape(a  ,					a )
														snake_case__	  =       tf.reshape(a  ,					a )

							# Compute layer normalization using the batch_normalization
							# function.
							snake_case__	  =       tf.nn.batch_normalization(
							    a  ,					a  ,					a  ,					offset=a  ,					scale=a  ,					variance_epsilon=a  ,					)
							return outputs

def _UpperCAmelCase   (	a						:	Optional[int]  ,					a						:	Dict=0  ,					a						:	List[Any]=-1 ):
							# Replicates the behavior of torch.flatten in TF

							# If end_dim or start_dim is negative, count them from the end
							if end_dim < 0:
														end_dim += input.shape.rank
							if start_dim < 0:
														start_dim += input.shape.rank

							if start_dim == end_dim:
														return input

							snake_case__	  =       tf.shape(a )
							snake_case__	  =       tf.math.reduce_prod(in_shape[start_dim : end_dim + 1] )
							snake_case__	  =       tf.concat([in_shape[:start_dim], [flattened_dim], in_shape[end_dim + 1 :]]  ,					axis=0 )
							return tf.reshape(a  ,					a )

def _UpperCAmelCase   (	a						:	tf.Tensor ):
							if not isinstance(a  ,					tf.Tensor ):
														snake_case__	  =       tf.convert_to_tensor(a )  # Catches stray NumPy inputs
							if encoder_attention_mask.shape.rank == 3:
														snake_case__	  =       encoder_attention_mask[:, None, :, :]
							if encoder_attention_mask.shape.rank == 2:
														snake_case__	  =       encoder_attention_mask[:, None, None, :]
							# T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
							# Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
							# /transformer/transformer_layers.py#L270
							# encoder_extended_attention_mask = (encoder_extended_attention_mask ==
							# encoder_extended_attention_mask.transpose(-1, -2))
							snake_case__	  =       (
							    tf.cast(1  ,					encoder_attention_mask.dtype ) - encoder_extended_attention_mask
							) * encoder_extended_attention_mask.dtype.min

							return encoder_extended_attention_mask

def _UpperCAmelCase   (	a						:	tf.Tensor  ,					a						:	int  ,					a						:	str = "input_ids" ):
							tf.debugging.assert_less(
							    a  ,					tf.cast(a  ,					dtype=tensor.dtype )  ,					message=(
							        F'''The maximum value of {tensor_name} ({tf.math.reduce_max(a )}) must be smaller than the embedding '''
							        F'''layer\'s input dimension ({embed_dim}). The likely cause is some problem at tokenization time.'''
							    )  ,					)

def _UpperCAmelCase   (	a						:	str  ,					a						:	Tuple  ,					a						:	Optional[int] ):
							snake_case__	  =       6_4512
							# Check that no item in `data` is larger than `HDF5_OBJECT_HEADER_LIMIT`
							# because in that case even chunking the array would not make the saving
							# possible.
							snake_case__	  =       [x for x in data if len(a ) > HDF5_OBJECT_HEADER_LIMIT]

							# Expecting this to never be true.
							if bad_attributes:
														raise RuntimeError(
														    """The following attributes cannot be saved to HDF5 file because """
														    F'''they are larger than {HDF5_OBJECT_HEADER_LIMIT} '''
														    F'''bytes: {bad_attributes}''' )

							snake_case__	  =       np.asarray(a )

							snake_case__	  =       1
							snake_case__	  =       np.array_split(a  ,					a )

							# This will never loop forever thanks to the test above.
							while any(x.nbytes > HDF5_OBJECT_HEADER_LIMIT for x in chunked_data ):
														num_chunks += 1
														snake_case__	  =       np.array_split(a  ,					a )

							if num_chunks > 1:
														for chunk_id, chunk_data in enumerate(a ):
																					snake_case__	  =       chunk_data
							else:
														snake_case__	  =       data

def _UpperCAmelCase   (	a						:	List[Any]  ,					a						:	Optional[int] ):
							if name in group.attrs:
														snake_case__	  =       [n.decode("""utf8""" ) if hasattr(a  ,					"""decode""" ) else n for n in group.attrs[name]]
							else:
														snake_case__	  =       []
														snake_case__	  =       0
														while "%s%d" % (name, chunk_id) in group.attrs:
																					data.extend(
																					    [n.decode("""utf8""" ) if hasattr(a  ,					"""decode""" ) else n for n in group.attrs["""%s%d""" % (name, chunk_id)]] )
																					chunk_id += 1
							return data



def _UpperCAmelCase   (	a						:	Optional[int] ):

							def _expand_single_ad_tensor(a						:	Any ):
														if isinstance(a  ,					tf.Tensor ) and t.shape.rank == 1:
																					return tf.expand_dims(a  ,					axis=-1 )
														return t

							return tf.nest.map_structure(_expand_single_ad_tensor  ,					a )
