o ۷i @sddlZddlmZddlZddlZddlmZmZddl m Z m Z ddl m Z mZmZe r3ddlZ  dd ed ed ed dejfddZGdddeeZdS)N)Literal) ConfigMixinregister_to_config) deprecateis_scipy_available)KarrasDiffusionSchedulersSchedulerMixinSchedulerOutput+?cosinenum_diffusion_timestepsmax_betaalpha_transform_type)r explaplacereturncCs|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. r cSs t|ddtjddS)NgMb?gT㥛 ?r)mathcospitrd/home/ubuntu/vllm_env/lib/python3.10/site-packages/diffusers/schedulers/scheduling_deis_multistep.py alpha_bar_fn<s z)betas_for_alpha_bar..alpha_bar_fnrc SsPdtdd|tddtd|d}t|}t|d|S)Ngr?rgư>)rcopysignlogfabsrsqrt)rlmbsnrrrrrAs4 rcSst|dS)Ng()rrrrrrrHsz"Unsupported alpha_transform_type: rdtype) ValueErrorrangeappendmintorchtensorfloat32)rrrrbetasit1t2rrrbetas_for_alpha_bar"s     "r0c0@s*eZdZdZddeDZdZe            dmdede de de dde j e e Bd Bdede ddede d e d!e dd"e dd#ed$ed%ed&ed'ed(e d)e d*d+ed,ed-e dd.d f.d/d0Zed.efd1d2Zed.efd3d4Zdnd5ed.d fd6d7Zdod8ed9eejBd:e d Bfd;d<Zd=ejd.ejfd>d?Zd@e j dAe j d.e j fdBdCZd@ejd.eejejffdDdEZdFejd8ed.ejfdGdHZdFejd8ed.ejfdIdJZ KdpdFejd8edLe dMe d.ejf dNdOZd dPdQejd=ejd.ejfdRdSZ d dPdQejd=ejd.ejfdTdUZ!d dPdVe ejd=ejd.ejfdWdXZ"d dPdVe ejd=ejd.ejfdYdZZ# dqd[eejBd\ejd Bd.efd]d^Z$d[eejBd.d fd_d`Z% drdQejd[eejBd=ejdaed.e&eBf dbdcZ'd=ejd.ejfdddeZ(dfejdgejdhej)d.ejfdidjZ*d.efdkdlZ+d S)sDEISMultistepScheduleru  `DEISMultistepScheduler` is a fast high order solver for diffusion ordinary differential equations (ODEs). 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 (`"linear"`, `"scaled_linear"`, or `"squaredcos_cap_v2"`, 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` or `list[float]`, *optional*): Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. solver_order (`int`, defaults to `2`): The DEIS order which can be `1` or `2` or `3`. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for unconditional sampling. prediction_type (`"epsilon"`, `"sample"`, `"v_prediction"`, or `"flow_prediction"`, defaults to `"epsilon"`): Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process), `sample` (directly predicts the noisy sample`), `v_prediction` (see section 2.4 of [Imagen Video](https://huggingface.co/papers/2210.02303) paper), or `flow_prediction`. 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`. algorithm_type (`"deis"`, defaults to `"deis"`): The algorithm type for the solver. solver_type (`"logrho"`, defaults to `"logrho"`): Solver type for DEIS. lower_order_final (`bool`, defaults to `True`): Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. 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. use_flow_sigmas (`bool`, *optional*, defaults to `False`): Whether to use flow sigmas for step sizes in the noise schedule during the sampling process. flow_shift (`float`, *optional*, defaults to `1.0`): The flow shift parameter for flow-based models. timestep_spacing (`"linspace"`, `"leading"`, or `"trailing"`, 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. use_dynamic_shifting (`bool`, defaults to `False`): Whether to use dynamic shifting for the noise schedule. time_shift_type (`"exponential"`, defaults to `"exponential"`): The type of time shifting to apply. cCsg|]}|jqSr)name).0errr sz!DEISMultistepScheduler.r-C6?{Gz?linearNrepsilonFףp= ??deislogrhoTlinspacer exponentialnum_train_timesteps beta_startbeta_end beta_schedule)r9 scaled_linearsquaredcos_cap_v2 trained_betas solver_orderprediction_type)r:sample v_predictionflow_prediction thresholdingdynamic_thresholding_ratiosample_max_valuealgorithm_type solver_typelower_order_finaluse_karras_sigmasuse_exponential_sigmasuse_beta_sigmasuse_flow_sigmas flow_shifttimestep_spacing)r?leadingtrailing steps_offsetuse_dynamic_shiftingtime_shift_typercCs|jjr ts tdt|jj|jj|jjgdkrtd|dur,tj |tj d|_ n:|dkr)midpointheunbh1bh2r>)rQz solver type cpu)%configrUr ImportErrorsumrTrSr%r)r*r+r,r?r0NotImplementedError __class__alphascumprodalphas_cumprodr alpha_tsigma_trlambda_tsigmasinit_noise_sigmarnum_inference_stepsnpcopy from_numpy timesteps model_outputslower_order_nums _step_index _begin_indexto)selfrArBrCrDrGrHrIrMrNrOrPrQrRrSrTrUrVrWrXr[r\r]rxrrr__init__sj   &  zDEISMultistepScheduler.__init__cC|jS)zg The index counter for current timestep. It will increase 1 after each scheduler step. )r{r~rrr step_indexz!DEISMultistepScheduler.step_indexcCr)zq The index for the first timestep. It should be set from pipeline with `set_begin_index` method. r|rrrr begin_indexrz"DEISMultistepScheduler.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)r~rrrrset_begin_indexs z&DEISMultistepScheduler.set_begin_indexrtdevicemuc s|durjjrjjdksJt|j_jjdkrz8DEISMultistepScheduler.set_timesteps..crrrrrrrr5Crcrrrrrrrr5Hrr<rr$rf)*rgr\r]rurrWrXr?rAroundrvastypeint64aranger[r%arrayrnrrSflip_convert_to_karras concatenater+rT_convert_to_exponentialrU_convert_to_betarVinterplenr)rwrrr}rxrtrHryrzr{r|) r~rtrrrx step_ratiorrrl sigma_lastrrr set_timestepssx    6  $         2   z$DEISMultistepScheduler.set_timestepsrJcCs|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. rr^)r(max)r$shaper)r+float64floatreshaperuprodabsquantilergrNclamprO unsqueezer})r~rJr$ batch_sizechannelsremaining_dims abs_samplesrrr_threshold_samplees    z(DEISMultistepScheduler._threshold_samplerrc 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)rr) rurmaximumnewaxiscumsumargmaxcliprr) r~rr log_sigmadistslow_idxhigh_idxlowhighwrrrrrs, z"DEISMultistepScheduler._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. rrr)rgrV)r~rrorprrr_sigma_to_alpha_sigma_ts z.DEISMultistepScheduler._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_maxrerg@r)hasattrrgrritemrur?) r~rrtrrrhoramp min_inv_rho max_inv_rhorrrrrrs      z)DEISMultistepScheduler._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. rNrrer) rrgrrrrurr?rr)r~rrtrrrrrrrrs     z.DEISMultistepScheduler._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. rNrrercsg|] }|qSrr)r3ppf)rrrrr5;sz;DEISMultistepScheduler._convert_to_beta..csg|] }tjj|qSr)scipystatsrr)r3timestep)rrrrr5=sr)rrgrrrrurr?)r~rrtrrrrr)rrrrrrs       z'DEISMultistepScheduler._convert_to_betarJ model_outputc Os2t|dkr |dn|dd}|dur#t|dkr|d}ntd|dur-tddd|j|j}||\}}|jjd krI||||} n5|jjd krR|} n,|jjd kra||||} n|jjd krt|j|j}|||} n td |jjd|jj r| | } |jj dkr||| |St d)a Convert the model output to the corresponding type the DEIS algorithm needs. Args: model_output (`torch.Tensor`): The direct output from the 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. Returns: `torch.Tensor`: The converted model output. rrNr/missing `sample` as a required keyword argumentrx1.0.0Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`r:rJrKrLzprediction_type given as zi must be one of `epsilon`, `sample`, `v_prediction`, or `flow_prediction` for the DEISMultistepScheduler.r='only support log-rho multistep deis now) rpopr%rrrrrrgrIrMrrPrj) r~rrJargskwargsrrrorpx0_predrrrconvert_model_outputEs<           z+DEISMultistepScheduler.convert_model_outputcOs2t|dkr |dn|dd}t|dkr|dn|dd}|dur3t|dkr/|d}ntd|dur=tdd d |durGtdd d |j|jd|j|j}}||\} }||\} }t| t|} t| t|} | | } |j j d kr| | ||t | d |}|St d)au One step for the first-order DEIS (equivalent to DDIM). Args: model_output (`torch.Tensor`): The direct output from the learned diffusion model. timestep (`int`): The current discrete timestep in the diffusion chain. 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. Returns: `torch.Tensor`: The sample tensor at the previous timestep. rrNr prev_timesteprrrxrrPassing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`r=r<r) rrr%rrrrrr)rrgrPrrj)r~rrJrrrrrpsigma_sroalpha_srqlambda_shx_trrrdeis_first_order_updates<      "z.DEISMultistepScheduler.deis_first_order_updatemodel_output_listcOst|dkr |dn|dd}t|dkr|dn|dd}|dur3t|dkr/|d}ntd|dur=tddd |durGtddd |j|jd|j|j|j|jd}}} ||\} }||\} }|| \} } |d |d } }|| || | | }}}|jjd krdd}||||||||}||||||||}| || || ||}|St d)a One step for the second-order multistep DEIS. Args: model_output_list (`list[torch.Tensor]`): The direct outputs from learned diffusion model at current and latter timesteps. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. Returns: `torch.Tensor`: The sample tensor at the previous timestep. r timestep_listNrrrrrPassing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`rrer=cSs2|t| t|dt|t|S)Nrrur)rbcrrrind_fns2zIDEISMultistepScheduler.multistep_deis_second_order_update..ind_fnr rrr%rrrrrrgrPrj)r~rrJrrrrrpsigma_s0sigma_s1roalpha_s0alpha_s1m0m1rho_trho_s0rho_s1rcoef1coef2rrrr"multistep_deis_second_order_updatesJ       z9DEISMultistepScheduler.multistep_deis_second_order_updatecOst|dkr |dn|dd}t|dkr|dn|dd}|dur3t|dkr/|d}ntd|dur=tddd |durGtddd |j|jd|j|j|j|jd|j|jdf\}}} } ||\} }||\} }|| \} } || \}} |d |d |d }}}|| || | | | |f\}}}}|jjdkrdd}||||||||||}||||||||||}||||||||||}| || ||||||}|St d)a One step for the third-order multistep DEIS. Args: model_output_list (`list[torch.Tensor]`): The direct outputs from learned diffusion model at current and latter timesteps. sample (`torch.Tensor`): A current instance of a sample created by diffusion process. Returns: `torch.Tensor`: The sample tensor at the previous timestep. rrNrrrrrrrrerr=cSs|t|t|t|dt|t|t|t|ddt|d}t|t|t|t|}||S)Nrrr)rrrd numerator denominatorrrrrHs   (zHDEISMultistepScheduler.multistep_deis_third_order_update..ind_fnrr)r~rrJrrrrrprrsigma_s2rorralpha_s2rrm2rrrrho_s2rrrcoef3rrrr!multistep_deis_third_order_updatesR        $z8DEISMultistepScheduler.multistep_deis_third_order_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)rxnonzerorr)r~rrindex_candidatesrrrrindex_for_timestep`s    z)DEISMultistepScheduler.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)Tensorr}rxrrr{r|)r~rrrr_init_step_indexs  z'DEISMultistepScheduler._init_step_index return_dictc Cs`|jdur td|jdur|||jt|jdko'|jjo't|jdk}|jt|jdko<|jjo.rre)rrr}rr$typer)is_floating_pointrxr+rrrflattenrrr) r~rrrxrr step_indicesrrorp noisy_samplesrrr add_noises$     z DEISMultistepScheduler.add_noisecCs|jjSN)rgrArrrr__len__szDEISMultistepScheduler.__len__)r6r7r8r9Nrr:Fr;r<r=r>TFFFFr<r?rFr@)r)NN)rrr)T),__name__ __module__ __qualname____doc__r _compatiblesorderrintrrrundarraylistboolrpropertyrrrstrr)rrrrrtuplerrrrrrrrrrr r r  IntTensorrrrrrrr1VsB=      _" U, %'# 4 @ A M \ % @ 1r1)r r )rtypingrnumpyrur)configuration_utilsrrutilsrrscheduling_utilsr r r scipy.statsrr rrr0r1rrrrs*  4