




"""simple docstring"""







import math
import warnings
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union

import numpy as np
import torch
from scipy import integrate

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin






@dataclass
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->LMSDiscrete
class  SCREAMING_SNAKE_CASE__   (   _A   ):
						_a        =  42
						_a        =  None





def 							lowercase   (snake_case__				:	str ,    snake_case__				:	List[str]=0.999 ,    snake_case__				:	Union[str, Any]="cosine" ,    )       ->					str:


				'''simple docstring'''






				if alpha_transform_type == "cosine":

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

				elif alpha_transform_type == "exp":

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

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

				lowerCAmelCase		=			[]
				for i in range(A_	):
								lowerCAmelCase		=			i / num_diffusion_timesteps
								lowerCAmelCase		=			(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  SCREAMING_SNAKE_CASE__   (   _A  ,			_A   ):
						_a        =  [e.name for e in KarrasDiffusionSchedulers]
						_a        =  1


						@register_to_config
						def __init__(							self			:     Tuple  ,  lowerCAmelCase			:     int = 1000  ,  lowerCAmelCase			:     float = 0.0001  ,  lowerCAmelCase			:     float = 0.02  ,  lowerCAmelCase			:     str = "linear"  ,  lowerCAmelCase			:     Optional[Union[np.ndarray, List[float]]] = None  ,  lowerCAmelCase			:     Optional[bool] = False  ,  lowerCAmelCase			:     str = "epsilon"  ,  lowerCAmelCase			:     str = "linspace"  ,  lowerCAmelCase			:     int = 0  ,  ):






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

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

										lowerCAmelCase		=			np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5   )
										lowerCAmelCase		=			np.concatenate([sigmas[::-1], [0.0]]   ).astype(np.floataa   )
										lowerCAmelCase		=			torch.from_numpy(UpperCamelCase__   )

										# setable values
										lowerCAmelCase		=			None
										lowerCAmelCase		=			use_karras_sigmas
										self.set_timesteps(UpperCamelCase__  ,  UpperCamelCase__   )
										lowerCAmelCase		=			[]
										lowerCAmelCase		=			False


						@property
						def 					__lowercase	(							self			:     Tuple   ):






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

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


						def 					__lowercase	(							self			:     Dict  ,  lowerCAmelCase			:     torch.FloatTensor  ,  lowerCAmelCase			:     Union[float, torch.FloatTensor]   ):






										if isinstance(UpperCamelCase__  ,  torch.Tensor   ):
														lowerCAmelCase		=			timestep.to(self.timesteps.device   )
										lowerCAmelCase		=			(self.timesteps == timestep).nonzero().item()
										lowerCAmelCase		=			self.sigmas[step_index]
										lowerCAmelCase		=			sample / ((sigma**2 + 1) ** 0.5)
										lowerCAmelCase		=			True
										return sample


						def 					__lowercase	(							self			:     List[Any]  ,  lowerCAmelCase			:     int  ,  lowerCAmelCase			:     Optional[int]  ,  lowerCAmelCase			:     Any   ):






										def lms_derivative(lowerCAmelCase			:     Tuple   ):
														lowerCAmelCase		=			1.0
														for k in range(UpperCamelCase__   ):
																		if current_order == k:
																						continue
																		prod *= (tau - self.sigmas[t - k]) / (self.sigmas[t - current_order] - self.sigmas[t - k])
														return prod

										lowerCAmelCase		=			integrate.quad(UpperCamelCase__  ,  self.sigmas[t]  ,  self.sigmas[t + 1]  ,  epsrel=1e-4   )[0]

										return integrated_coeff


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

										# "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
										if self.config.timestep_spacing == "linspace":
														lowerCAmelCase		=			np.linspace(0  ,  self.config.num_train_timesteps - 1  ,  UpperCamelCase__  ,  dtype=UpperCamelCase__   )[
														    ::-1
														].copy()
										elif self.config.timestep_spacing == "leading":
														lowerCAmelCase		=			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
														lowerCAmelCase		=			(np.arange(0  ,  UpperCamelCase__   ) * step_ratio).round()[::-1].copy().astype(UpperCamelCase__   )
														timesteps += self.config.steps_offset
										elif self.config.timestep_spacing == "trailing":
														lowerCAmelCase		=			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
														lowerCAmelCase		=			(np.arange(self.config.num_train_timesteps  ,  0  ,  -step_ratio   )).round().copy().astype(UpperCamelCase__   )
														timesteps -= 1
										else:
														raise ValueError(
														    f'''{self.config.timestep_spacing} is not supported. Please make sure to choose one of \'linspace\', \'leading\' or \'trailing\'.'''   )

										lowerCAmelCase		=			np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5   )
										lowerCAmelCase		=			np.log(UpperCamelCase__   )
										lowerCAmelCase		=			np.interp(UpperCamelCase__  ,  np.arange(0  ,  len(UpperCamelCase__   )   )  ,  UpperCamelCase__   )

										if self.use_karras_sigmas:
														lowerCAmelCase		=			self._convert_to_karras(in_sigmas=UpperCamelCase__   )
														lowerCAmelCase		=			np.array([self._sigma_to_t(UpperCamelCase__  ,  UpperCamelCase__   ) for sigma in sigmas]   )

										lowerCAmelCase		=			np.concatenate([sigmas, [0.0]]   ).astype(np.floataa   )

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

										lowerCAmelCase		=			[]


						def 					__lowercase	(							self			:     Tuple  ,  lowerCAmelCase			:     Dict  ,  lowerCAmelCase			:     List[Any]   ):
										lowerCAmelCase		=			np.log(UpperCamelCase__   )

										# get distribution
										lowerCAmelCase		=			log_sigma - log_sigmas[:, np.newaxis]

										# get sigmas range
										lowerCAmelCase		=			np.cumsum((dists >= 0)  ,  axis=0   ).argmax(axis=0   ).clip(max=log_sigmas.shape[0] - 2   )
										lowerCAmelCase		=			low_idx + 1

										lowerCAmelCase		=			log_sigmas[low_idx]
										lowerCAmelCase		=			log_sigmas[high_idx]

										# interpolate sigmas
										lowerCAmelCase		=			(low - log_sigma) / (low - high)
										lowerCAmelCase		=			np.clip(UpperCamelCase__  ,  0  ,  1   )

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


						def 					__lowercase	(							self			:     Union[str, Any]  ,  lowerCAmelCase			:     torch.FloatTensor   ):
										lowerCAmelCase		=			in_sigmas[-1].item()
										lowerCAmelCase		=			in_sigmas[0].item()

										lowerCAmelCase		=			7.0  # 7.0 is the value used in the paper
										lowerCAmelCase		=			np.linspace(0  ,  1  ,  self.num_inference_steps   )
										lowerCAmelCase		=			sigma_min ** (1 / rho)
										lowerCAmelCase		=			sigma_max ** (1 / rho)
										lowerCAmelCase		=			(max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
										return sigmas


						def 					__lowercase	(							self			:     List[str]  ,  lowerCAmelCase			:     torch.FloatTensor  ,  lowerCAmelCase			:     Union[float, torch.FloatTensor]  ,  lowerCAmelCase			:     torch.FloatTensor  ,  lowerCAmelCase			:     int = 4  ,  lowerCAmelCase			:     bool = True  ,  ):






										if not self.is_scale_input_called:
														warnings.warn(
														    """The `scale_model_input` function should be called before `step` to ensure correct denoising. """
														    """See `StableDiffusionPipeline` for a usage example."""   )

										if isinstance(UpperCamelCase__  ,  torch.Tensor   ):
														lowerCAmelCase		=			timestep.to(self.timesteps.device   )
										lowerCAmelCase		=			(self.timesteps == timestep).nonzero().item()
										lowerCAmelCase		=			self.sigmas[step_index]

										# 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
										if self.config.prediction_type == "epsilon":
														lowerCAmelCase		=			sample - sigma * model_output
										elif self.config.prediction_type == "v_prediction":
														# * c_out + input * c_skip
														lowerCAmelCase		=			model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1))
										elif self.config.prediction_type == "sample":
														lowerCAmelCase		=			model_output
										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
										lowerCAmelCase		=			(sample - pred_original_sample) / sigma
										self.derivatives.append(UpperCamelCase__   )
										if len(self.derivatives   ) > order:
														self.derivatives.pop(0   )

										# 3. Compute linear multistep coefficients
										lowerCAmelCase		=			min(step_index + 1  ,  UpperCamelCase__   )
										lowerCAmelCase		=			[self.get_lms_coefficient(UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__   ) for curr_order in range(UpperCamelCase__   )]

										# 4. Compute previous sample based on the derivatives path
										lowerCAmelCase		=			sample + sum(
										    coeff * derivative for coeff, derivative in zip(UpperCamelCase__  ,  reversed(self.derivatives   )   )   )

										if not return_dict:
														return (prev_sample,)

										return LMSDiscreteSchedulerOutput(prev_sample=UpperCamelCase__  ,  pred_original_sample=UpperCamelCase__   )


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

										lowerCAmelCase		=			[(schedule_timesteps == t).nonzero().item() for t in timesteps]

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

										lowerCAmelCase		=			original_samples + noise * sigma
										return noisy_samples




						def __len__(							self			:     str   ):
										return self.config.num_train_timesteps






"""simple docstring"""







import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding






class  SCREAMING_SNAKE_CASE__   (   _a   ):
						_a        =  ['image_processor', 'tokenizer']
						_a        =  'VideoMAEImageProcessor'
						_a        =  ('CLIPTokenizer', 'CLIPTokenizerFast')


						def __init__(							self			:     Union[str, Any]  ,  lowerCAmelCase			:     Optional[Any]=None  ,  lowerCAmelCase			:     Tuple=None  ,  **lowerCAmelCase			:     List[str]   ):
										lowerCAmelCase		=			None
										if "feature_extractor" in kwargs:
														warnings.warn(
														    """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
														    """ instead."""  ,  lowerCAmelCase  ,  )
														lowerCAmelCase		=			kwargs.pop("""feature_extractor"""   )

										lowerCAmelCase		=			image_processor if image_processor is not None else feature_extractor
										if image_processor is None:
														raise ValueError("""You need to specify an `image_processor`."""   )
										if tokenizer is None:
														raise ValueError("""You need to specify a `tokenizer`."""   )

										super().__init__(lowerCAmelCase  ,  lowerCAmelCase   )
										lowerCAmelCase		=			self.image_processor


						def __call__(							self			:     Union[str, Any]  ,  lowerCAmelCase			:     str=None  ,  lowerCAmelCase			:     Tuple=None  ,  lowerCAmelCase			:     int=None  ,  **lowerCAmelCase			:     str   ):

										if text is None and videos is None:
														raise ValueError("""You have to specify either text or videos. Both cannot be none."""   )

										if text is not None:
														lowerCAmelCase		=			self.tokenizer(lowerCAmelCase  ,  return_tensors=lowerCAmelCase  ,  **lowerCAmelCase   )

										if videos is not None:
														lowerCAmelCase		=			self.image_processor(lowerCAmelCase  ,  return_tensors=lowerCAmelCase  ,  **lowerCAmelCase   )

										if text is not None and videos is not None:
														lowerCAmelCase		=			image_features.pixel_values
														return encoding
										elif text is not None:
														return encoding
										else:
														return BatchEncoding(data=dict(**lowerCAmelCase   )  ,  tensor_type=lowerCAmelCase   )


						def 					__lowercase	(							self			:     str  ,  *lowerCAmelCase			:     List[str]  ,  **lowerCAmelCase			:     int   ):
										return self.tokenizer.batch_decode(*lowerCAmelCase  ,  **lowerCAmelCase   )


						def 					__lowercase	(							self			:     List[str]  ,  *lowerCAmelCase			:     Any  ,  **lowerCAmelCase			:     List[Any]   ):
										return self.tokenizer.decode(*lowerCAmelCase  ,  **lowerCAmelCase   )


						@property
						def 					__lowercase	(							self			:     Dict   ):
										return ["input_ids", "attention_mask", "position_ids", "pixel_values"]


						@property
						def 					__lowercase	(							self			:     List[str]   ):
										warnings.warn(
										    """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""  ,  lowerCAmelCase  ,  )
										return self.image_processor_class




						@property
						def 					__lowercase	(							self			:     Tuple   ):
										warnings.warn(
										    """`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."""  ,  lowerCAmelCase  ,  )
										return self.image_processor





