o Gi0 @sddlZddlmZmZddlZddlZddlmZm Z ddl m Z m Z ddl mZddlmZmZmZe r;ddlZ dd ed ed eddejfddZGdddeeZdS)N)CallableLiteral) ConfigMixinregister_to_config) deprecateis_scipy_available) randn_tensor)KarrasDiffusionSchedulersSchedulerMixinSchedulerOutput+?cosinenum_diffusion_timestepsmax_betaalpha_transform_type)rexplaplacereturncCs|dkr dd}n|dkrdd}n|dkrdd}ntd|g}t|D]}||}|d |}|td |||||q(tj|tjd S) a> Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of (1-beta) over time from t = [0,1]. Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up to that part of the diffusion process. Args: num_diffusion_timesteps (`int`): The number of betas to produce. max_beta (`float`, defaults to `0.999`): The maximum beta to use; use values lower than 1 to avoid numerical instability. alpha_transform_type (`str`, defaults to `"cosine"`): The type of noise schedule for `alpha_bar`. Choose from `cosine`, `exp`, or `laplace`. Returns: `torch.Tensor`: The betas used by the scheduler to step the model outputs. rcSs t|ddtjddS)NgMb?gT㥛 ?r)mathcospitr\/home/ubuntu/.local/lib/python3.10/site-packages/diffusers/schedulers/scheduling_sasolver.py alpha_bar_fn=s z)betas_for_alpha_bar..alpha_bar_fnrc SsPdtdd|tddtd|d}t|}t|d|S)Ngr ?rgư>)rcopysignlogfabsrsqrt)rlmbsnrrrrrBs4 rcSst|dS)Ng()rrrrrrrIsz"Unsupported alpha_transform_type: r dtype) ValueErrorrangeappendmintorchtensorfloat32)rrrrbetasit1t2rrrbetas_for_alpha_bar#s     "r2c0@seZdZdZddeDZdZeddddd d d d d d d dddd d d d ded d ddfde dedede de j e eBd Bde de de ded Bdededed e d!ed"ed#ed$ed%ed&ed'ed(e d Bd)e d*e f.d+d,Zed-d.Zed/d0Zdrd1e fd2d3Zdsd4e d5e ejBfd6d7Zd8ejd9ejfd:d;Zdd?Zd@ejd9ejfdAdBZd@ejd4e d9ejfdCdDZ Edtd@ejd4e dFedGed9ejf dHdIZd dJdKejd8ejd9ejfdLdMZdNdOZ dPdQZ!dRdSZ"dTdUZ#dKejd8ejdVejdWe dXejd9ejf dYdZZ$d[ejd\ejd]ejd^ejdWe dXejd9ejfd_d`Z% dudae ejBdbejd Bd9e fdcddZ&dedfZ' dvdKejdae d8ejdged9e(e)Bf dhdiZ*d8ejd9ejfdjdkZ+dlejdVejdmej,d9ejfdndoZ-dpdqZ.d S)wSASolverScheduleru `SASolverScheduler` is a fast dedicated high-order solver for diffusion SDEs. This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic methods the library implements for all schedulers such as loading and saving. Args: num_train_timesteps (`int`, defaults to 1000): The number of diffusion steps to train the model. beta_start (`float`, defaults to 0.0001): The starting `beta` value of inference. beta_end (`float`, defaults to 0.02): The final `beta` value. beta_schedule (`str`, defaults to `"linear"`): The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from `linear`, `scaled_linear`, or `squaredcos_cap_v2`. trained_betas (`np.ndarray`, *optional*): Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. predictor_order (`int`, defaults to 2): The predictor order which can be `1` or `2` or `3` or '4'. It is recommended to use `predictor_order=2` for guided sampling, and `predictor_order=3` for unconditional sampling. corrector_order (`int`, defaults to 2): The corrector order which can be `1` or `2` or `3` or '4'. It is recommended to use `corrector_order=2` for guided sampling, and `corrector_order=3` for unconditional sampling. prediction_type (`str`, defaults to `epsilon`, *optional*): Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process), `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen Video](https://huggingface.co/papers/2210.02303) paper). tau_func (`Callable`, *optional*): Stochasticity during the sampling. Default in init is `lambda t: 1 if t >= 200 and t <= 800 else 0`. SA-Solver will sample from vanilla diffusion ODE if tau_func is set to `lambda t: 0`. SA-Solver will sample from vanilla diffusion SDE if tau_func is set to `lambda t: 1`. For more details, please check https://huggingface.co/papers/2309.05019 thresholding (`bool`, defaults to `False`): Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such as Stable Diffusion. dynamic_thresholding_ratio (`float`, defaults to 0.995): The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. sample_max_value (`float`, defaults to 1.0): The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `algorithm_type="dpmsolver++"`. algorithm_type (`str`, defaults to `data_prediction`): Algorithm type for the solver; can be `data_prediction` or `noise_prediction`. It is recommended to use `data_prediction` with `solver_order=2` for guided sampling like in Stable Diffusion. lower_order_final (`bool`, defaults to `True`): Whether to use lower-order solvers in the final steps. Default = True. use_karras_sigmas (`bool`, *optional*, defaults to `False`): Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, the sigmas are determined according to a sequence of noise levels {σi}. use_exponential_sigmas (`bool`, *optional*, defaults to `False`): Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. use_beta_sigmas (`bool`, *optional*, defaults to `False`): Whether to use beta sigmas for step sizes in the noise schedule during the sampling process. Refer to [Beta Sampling is All You Need](https://huggingface.co/papers/2407.12173) for more information. lambda_min_clipped (`float`, defaults to `-inf`): Clipping threshold for the minimum value of `lambda(t)` for numerical stability. This is critical for the cosine (`squaredcos_cap_v2`) noise schedule. variance_type (`str`, *optional*): Set to "learned" or "learned_range" for diffusion models that predict variance. If set, the model's output contains the predicted Gaussian variance. timestep_spacing (`str`, defaults to `"linspace"`): The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. steps_offset (`int`, defaults to 0): An offset added to the inference steps, as required by some model families. cCsg|]}|jqSr)name).0errr szSASolverScheduler.r ig-C6?g{Gz?linearNrepsilonFgףp= ??data_predictionTinflinspacernum_train_timesteps beta_startbeta_end beta_schedule trained_betaspredictor_ordercorrector_orderprediction_typetau_func thresholdingdynamic_thresholding_ratiosample_max_valuealgorithm_typelower_order_finaluse_karras_sigmasuse_exponential_sigmasuse_beta_sigmasuse_flow_sigmas flow_shiftlambda_min_clipped variance_typetimestep_spacing steps_offsetcCs|jjr ts tdt|jj|jj|jjgdkrtd|dur,tj |tj d|_ n:|dkrz,SASolverScheduler.__init__..r;cpu))configrNr ImportErrorsumrMrLr'r+r,r-r.r=r2NotImplementedError __class__alphascumprodalphas_cumprodr"alpha_tsigma_tr lambda_tsigmasinit_noise_sigmanum_inference_stepsnpcopy from_numpy timestepsmax timestep_list model_outputsrF predict_x0lower_order_nums last_sample _step_index _begin_indexto)selfr>r?r@rArBrCrDrErFrGrHrIrJrKrLrMrNrOrPrQrRrSrTrprrr__init__s\  &   zSASolverScheduler.__init__cC|jS)zg The index counter for current timestep. It will increase 1 after each scheduler step. )rwrzrrr step_indexzSASolverScheduler.step_indexcCr|)zq The index for the first timestep. It should be set from pipeline with `set_begin_index` method. rxr}rrr begin_indexrzSASolverScheduler.begin_indexrcCs ||_dS)z Sets the begin index for the scheduler. This function should be run from pipeline before the inference. Args: begin_index (`int`, defaults to `0`): The begin index for the scheduler. Nr)rzrrrrset_begin_indexs z!SASolverScheduler.set_begin_indexrldevicec sttjdgjj}jj|}jj dkr;t d|d|d ddddd t j}nXjj dkri||d}t d|d| ddddd t j}|jj7}n*jj dkrjj|}t |d|  t j}|d8}n tjj dt djjd }t |jjrt | }j||d }t fd d |D }t ||ddgt j}nΈjjrt | }j||d }t fd d |D}t ||ddgt j}njjr6t | }j||d }t fdd |D}t ||ddgt j}nnjjr}t ddjj|d}d|}t jj|djjd|dd }|jj }t ||ddgt j}n't |t dt!||}djdjdd } t || ggt j}t"|_#t"|j$|tjd_%t!|_&dgt'jj(jj)d_*d_+d_,d_-d_.j#$d_#dS)a Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: num_inference_steps (`int`): The number of diffusion steps used when generating samples with a pre-trained model. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. rr=r NrZleadingtrailingzY is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'.r) in_sigmasrlcg|]}|qSr _sigma_to_tr5sigma log_sigmasrzrrr78r]z3SASolverScheduler.set_timesteps..crrrrrrrr7=r]crrrrrrrr7Br]r:)rr&r^)/r+ searchsortedfliprir_rQr>numpyitemrSrmr=roundrnastypeint64arangerTr'arrayrfr rL_convert_to_karras concatenater-rM_convert_to_exponentialrN_convert_to_betarOrPinterplenrorjryrprlrqrCrDrsrurvrwrx) rzrlr clipped_idx last_timesteprp step_ratiorjrd sigma_lastrrr set_timesteps sl  6  6            2   zSASolverScheduler.set_timestepssamplercCs|j}|j^}}}|tjtjfvr|}|||t|}| }tj ||j j dd}tj |d|j jd}|d}t || ||}|j||g|R}||}|S)az Apply dynamic thresholding to the predicted sample. "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing pixels from saturation at each step. We find that dynamic thresholding results in significantly better photorealism as well as better image-text alignment, especially when using very large guidance weights." https://huggingface.co/papers/2205.11487 Args: sample (`torch.Tensor`): The predicted sample to be thresholded. Returns: `torch.Tensor`: The thresholded sample. r rW)r*rq)r&shaper+r-float64floatreshapermprodabsquantiler_rHclamprI unsqueezery)rzrr& batch_sizechannelsremaining_dims abs_samplesrrr_threshold_sample_s    z#SASolverScheduler._threshold_samplec Cstt|d}||ddtjf}tj|dkddjddj|jddd}|d}||}||}||||} t| dd} d| || |} | |j} | S)a Convert sigma values to corresponding timestep values through interpolation. Args: sigma (`np.ndarray`): The sigma value(s) to convert to timestep(s). log_sigmas (`np.ndarray`): The logarithm of the sigma schedule used for interpolation. Returns: `np.ndarray`: The interpolated timestep value(s) corresponding to the input sigma(s). g|=Nr)axisr)rqr ) rmr maximumnewaxiscumsumargmaxcliprr) rzrr log_sigmadistslow_idxhigh_idxlowhighwrrrrrs, zSASolverScheduler._sigma_to_tcCs@|jjrd|}|}||fSd|ddd}||}||fS)a( Convert sigma values to alpha_t and sigma_t values. Args: sigma (`torch.Tensor`): The sigma value(s) to convert. Returns: `tuple[torch.Tensor, torch.Tensor]`: A tuple containing (alpha_t, sigma_t) values. r rr)r_rO)rzrrgrhrrr_sigma_to_alpha_sigma_ts z)SASolverScheduler._sigma_to_alpha_sigma_trc Cst|jdr |jj}nd}t|jdr|jj}nd}|dur |n|d}|dur,|n|d}d}tdd|}|d|}|d|}|||||} | S)a Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative Models](https://huggingface.co/papers/2206.00364). Args: in_sigmas (`torch.Tensor`): The input sigma values to be converted. num_inference_steps (`int`): The number of inference steps to generate the noise schedule for. Returns: `torch.Tensor`: The converted sigma values following the Karras noise schedule. sigma_minN sigma_maxrZrg@r )hasattrr_rrrrmr=) rzrrlrrrhoramp min_inv_rho max_inv_rhorjrrrrs      z$SASolverScheduler._convert_to_karrascCst|jdr |jj}nd}t|jdr|jj}nd}|dur |n|d}|dur,|n|d}ttt |t ||}|S)a Construct an exponential noise schedule. Args: in_sigmas (`torch.Tensor`): The input sigma values to be converted. num_inference_steps (`int`): The number of inference steps to generate the noise schedule for. Returns: `torch.Tensor`: The converted sigma values following an exponential schedule. rNrrZr) rr_rrrrmrr=rr )rzrrlrrrjrrrrs     z)SASolverScheduler._convert_to_exponential333333?alphabetac st|jdr |jjndt|jdr|jjnddur n|ddur,n|dtfddfddd tdd |DD}|S) a Construct a beta noise schedule as proposed in [Beta Sampling is All You Need](https://huggingface.co/papers/2407.12173). Args: in_sigmas (`torch.Tensor`): The input sigma values to be converted. num_inference_steps (`int`): The number of inference steps to generate the noise schedule for. alpha (`float`, *optional*, defaults to `0.6`): The alpha parameter for the beta distribution. beta (`float`, *optional*, defaults to `0.6`): The beta parameter for the beta distribution. Returns: `torch.Tensor`: The converted sigma values following a beta distribution schedule. rNrrZrcsg|] }|qSrr)r5ppf)rrrrr75sz6SASolverScheduler._convert_to_beta..csg|] }tjj|qSr)scipystatsrr)r5timestep)rrrrr77sr )rr_rrrrmrr=)rzrrlrrrjr)rrrrrrs       z"SASolverScheduler._convert_to_betar model_outputc Os"t|dkr |dn|dd}|dur#t|dkr|d}ntd|dur-tddd|j|j}||\}}|jjd vr|jj d kr_|jj d vrV|dddd f}||||} n5|jj d krh|} n,|jj dkrw||||} n|jj dkr|j|j}|||} n td|jj d|jj r| | } | S|jjdvr|jj d kr|jj d vr|dddd f} n+|} n(|jj d kr||||} n|jj dkr||||} n td|jj d|jj r |j ||j|}}||| |} | | } ||| |} | SdS)a= Convert the model output to the corresponding type the data_prediction/noise_prediction algorithm needs. Noise_prediction is designed to discretize an integral of the noise prediction model, and data_prediction is designed to discretize an integral of the data prediction model. > [!TIP] > The algorithm and model type are decoupled. You can use either data_prediction or noise_prediction for both > noise prediction and data prediction models. Args: model_output (`torch.Tensor`): The direct output from the learned diffusion model. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. Returns: `torch.Tensor`: The converted model output. rrNr /missing `sample` as a required keyword argumentrp1.0.0zPassing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`)r;r9)learned learned_ranger v_predictionflow_predictionzprediction_type given as zd must be one of `epsilon`, `sample`, `v_prediction`, or `flow_prediction` for the SASolverScheduler.)rYzQ must be one of `epsilon`, `sample`, or `v_prediction` for the SASolverScheduler.)rpopr'rrjr~rr_rJrErRrGrrgrh) rzrrargskwargsrrrgrhx0_predr9rrrconvert_model_output?sd                  z&SASolverScheduler.convert_model_outputcCs|dvsJd|dkrt| t||dS|dkr4t| |dt|||dS|dkrYt| |dd|dt|||dd|dS|dkrt| |dd|dd|dt|||dd|dd|dSdS) zd Calculate the integral of exp(-x) * x^order dx from interval_start to interval_end rr rr)order is only supported for 0, 1, 2 and 3rr rrNr+r)rzorderinterval_start interval_endrrr%get_coefficients_exponential_negatives,     z7SASolverScheduler.get_coefficients_exponential_negativecCst|dvsJdd|d|}d|d|}|dkr1t|dt|| d|dS|dkrRt||d|dt|| d|ddS|dkrt||dd|d|dd|dt|| d|ddS|dkrt||dd|dd|d|dd|dd|dt|| d|ddSd S) zl Calculate the integral of exp(x(1+tau^2)) * x^order dx from interval_start to interval_end rrr rrrrNr)rzrrrtauinterval_end_covinterval_start_covrrr%get_coefficients_exponential_positivesL( z7SASolverScheduler.get_coefficients_exponential_positivec Cs$|dvsJ|t|dksJ|dkrdggS|dkrJd|d|d|d |d|dgd|d|d|d |d|dggS|dkr|d|d|d|d}|d|d|d|d}|d|d|d|d}d||d |d||d|d|gd||d |d||d|d|gd||d |d||d|d|ggS|dkr|d|d|d|d|d|d}|d|d|d|d|d|d}|d|d|d|d|d|d}|d|d|d|d|d|d}d||d |d|d||d|d|d|d|d|d||d |d|d|gd||d |d|d||d|d|d|d|d|d||d |d|d|gd||d |d|d||d|d|d|d|d|d||d |d|d|gd||d |d|d||d|d|d|d|d|d||d |d|d|ggSdS)zB Calculate the coefficient of lagrange polynomial rr rrrN)r)rzr lambda_list denominator1 denominator2 denominator3 denominator4rrrlagrange_polynomial_coefficients        z1SASolverScheduler.lagrange_polynomial_coefficientc Cs|dvsJ|t|ksJdg}||d|}t|D];}d} t|D]-} |jr@| ||| ||d| |||7} q&| ||| ||d| ||7} q&|| qt||ksdJd|S)N)r rrrz4the length of lambda list must be equal to the orderr rz3the length of coefficients does not match the order)rrr(rtrrr)) rzrrrrr coefficientslagrange_coefficientr/ coefficientjrrrget_coefficients_fnKs"     z%SASolverScheduler.get_coefficients_fnnoiserrc Os t|dkr |dn|dd}|dur#t|dkr|d}ntd|dur6t|dkr2|d}ntd|durIt|dkrE|d}ntd |dur\t|d krX|d }ntd |durftdd d |j} |j|jd|j|j} } || \} } || \} } t | t | }t | t | }t |}||}g}t |D] }|j|}||j|\}}t |t |}| |q| |||||}|}|jrm|dkrm|j|jd}||\}}t |t |}|ddtd|d||dd|d|ddtd|d| d|dd||7<|ddtd|d||dd|d|ddtd|d| d|dd||8<t |D]>}|jr|d|d| t|d |||| |d 7}qq|d|d | ||| |d 7}qq|jr| tdtd|d||}n|| ttd|d|}|jrt|d || | |||}n | | |||}||j}|S)ag One step for the SA-Predictor. Args: model_output (`torch.Tensor`): The direct output from the learned diffusion model at the current timestep. prev_timestep (`int`): The previous discrete timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. order (`int`): The order of SA-Predictor at this timestep. Returns: `torch.Tensor`: The sample tensor at the previous timestep. r prev_timestepNr rrz.missing `noise` as a required keyword argumentr.missing `order` as a required keyword argumentr,missing `tau` as a required keyword argumentrzPassing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`r:rrr'rrsrjr~rr+r zeros_liker(r)rrtrr"ryr&) rzrrrrrrrrmodel_output_listrhsigma_s0rgalpha_s0ri lambda_s0 gradient_parthrr/sialpha_sisigma_si lambda_sigradient_coefficientsx temp_sigma temp_alpha_s temp_sigma_s temp_lambda_s noise_partx_trrr!stochastic_adams_bashforth_update_s               FF   0*$* z3SASolverScheduler.stochastic_adams_bashforth_updatethis_model_outputrv last_noise this_samplec Ost|dkr |dn|dd} |dur#t|dkr|d}ntd|dur6t|dkr2|d}ntd|durIt|dkrE|d}ntd |dur\t|d krX|d }ntd |durot|d krk|d }ntd | durytddd|j} |j|j|j|jd} } || \} } || \}} t | t | }t |t | }t |}||}g}t |D] }|j|}||j|\}}t |t |}| |q| |g}| |||||}|}|jrd|dkrd|ddtd|d||d|d|ddtd|d| d|dd|7<|ddtd|d||d|d|ddtd|d| d|dd|8<t |D]>}|jr|d|d| t|d |||||d 7}qh|d|d | ||||d 7}qh|jr| tdtd|d||}n|| ttd|d|}|jrt|d || | |||}n | ||||}||j}|S)a One step for the SA-Corrector. Args: this_model_output (`torch.Tensor`): The model outputs at `x_t`. this_timestep (`int`): The current timestep `t`. last_sample (`torch.Tensor`): The generated sample before the last predictor `x_{t-1}`. this_sample (`torch.Tensor`): The generated sample after the last predictor `x_{t}`. order (`int`): The order of SA-Corrector at this step. Returns: `torch.Tensor`: The corrected sample tensor at the current timestep. r this_timestepNr z4missing `last_sample` as a required keyword argumentrz3missing `last_noise` as a required keyword argumentrz4missing `this_sample` as a required keyword argumentrrrrzPassing `this_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`r:rr)rzrrvrrrrrrrrrhrrgrrirrrrr/rrrrmodel_prev_listrrr r rrrstochastic_adams_moulton_updates                  FF   0*$* z1SASolverScheduler.stochastic_adams_moulton_updaterschedule_timestepscCsd|dur|j}||k}t|dkrt|jd}|St|dkr*|d}|S|d}|S)a Find the index for a given timestep in the schedule. Args: timestep (`int` or `torch.Tensor`): The timestep for which to find the index. schedule_timesteps (`torch.Tensor`, *optional*): The timestep schedule to search in. If `None`, uses `self.timesteps`. Returns: `int`: The index of the timestep in the schedule. Nrr )rpnonzerorr)rzrrindex_candidatesr~rrrindex_for_timestepds    z$SASolverScheduler.index_for_timestepcCs@|jdurt|tjr||jj}|||_dS|j |_dS)z Initialize the step_index counter for the scheduler. Args: timestep (`int` or `torch.Tensor`): The current timestep for which to initialize the step index. N) r isinstancer+Tensorryrprrrwrx)rzrrrr_init_step_indexs  z"SASolverScheduler._init_step_index return_dictcCs|jdur td|jdur|||jdko|jdu}|j||d}|r<||jd}|j||j|j ||j |d}t t |j j|j jddD]} |j| d|j| <|j| d|j| <qK||jd<||jd<t|j||j|jd} |j jrt|j jt|j|j} t|j jt|j|jd} n|j j} |j j} t| |jd|_t| |jd |_ |jdksJ|j dksJ||_| |_ ||jd}|j||| |j|d } |jt |j j|j jdkr|jd7_|jd7_|s| fSt| d S) a Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the SA-Solver. Args: model_output (`torch.Tensor`): The direct output from learned diffusion model. timestep (`int`): The current discrete timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. generator (`torch.Generator`, *optional*): A random number generator. return_dict (`bool`): Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`. Returns: [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`: If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a tuple is returned where the first element is the sample tensor. NzaNumber of inference steps is 'None', you need to run 'set_timesteps' after creating the schedulerrrrZ)rrvrrrrr ) generatorrr&r)rrrrr) prev_sample)rlr'r~rrvrrFrrrrthis_corrector_orderr(rqr_rCrDrsr rrr&rKr*rrpruthis_predictor_orderr rwr )rzrrrrr use_correctormodel_output_convert current_taur/rr rrrrrstepsl   "     zSASolverScheduler.stepcOs|S)a? Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep. Args: sample (`torch.Tensor`): The input sample. Returns: `torch.Tensor`: A scaled input sample. r)rzrrrrrrscale_model_inputs z#SASolverScheduler.scale_model_inputoriginal_samplesrpcCs|jj|jd|_|jj|jd}||j}||d}|}t|jt|jkr:|d}t|jt|jks+d||d}|}t|jt|jkr_|d}t|jt|jksP||||}|S)a2 Add noise to the original samples according to the noise magnitude at each timestep (this is the forward diffusion process). Args: original_samples (`torch.Tensor`): The original samples to which noise will be added. noise (`torch.Tensor`): The noise to add to the samples. timesteps (`torch.IntTensor`): The timesteps indicating the noise level for each sample. Returns: `torch.Tensor`: The noisy samples. )rr%rrZr )rfryrr&flattenrrr)rzr&rrprfsqrt_alpha_prodsqrt_one_minus_alpha_prod noisy_samplesrrr add_noises    zSASolverScheduler.add_noisecCs|jjSN)r_r>r}rrr__len__9szSASolverScheduler.__len__)r)NN)rrr,)NT)/__name__ __module__ __qualname____doc__r _compatiblesrrrintstrrmndarraylistrboolr{propertyr~rrr+rrrrrrrrrrrrrrr rrrr tupler$r% IntTensorr+r-rrrrr3WsXC      U   R,%'# 4 [,m    % f *r3)rr)rtypingrrrrmr+configuration_utilsrrutilsrrutils.torch_utilsr scheduling_utilsr r r scipy.statsrr3rrr2r3rrrrs,  4