o j]@sZddlZddlZGdddZdidddddifddZGd d d Zd d Zd dZdS)Nc@sLeZdZdddddejfddZddZd d Zd d Zd dZ ddZ dS)NoiseScheduleVPdiscreteNg?g4@cCsV|dvr td|||_|dkr[|dur$dtd|jdd}n |dus*Jdt|}t||_d |_t d d |jddd d j |d |_ | d j |d |_ dSd |_||_||_d|_d|_t|jd |jtjdd |jtj|j|_tt|jd |jtjd|_||_|dkrd|_dSd |_dS)aLCreate a wrapper class for the forward SDE (VP type). *** Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t. We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images. *** The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ). We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper). Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have: log_alpha_t = self.marginal_log_mean_coeff(t) sigma_t = self.marginal_std(t) lambda_t = self.marginal_lambda(t) Moreover, as lambda(t) is an invertible function, we also support its inverse function: t = self.inverse_lambda(lambda_t) =============================================================== We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]). 1. For discrete-time DPMs: For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by: t_i = (i + 1) / N e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1. We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3. Args: betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details) alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details) Note that we always have alphas_cumprod = cumprod(1 - betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`. **Important**: Please pay special attention for the args for `alphas_cumprod`: The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ). Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have alpha_{t_n} = \sqrt{\hat{alpha_n}}, and log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}). 2. For continuous-time DPMs: We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise schedule are the default settings in DDPM and improved-DDPM: Args: beta_min: A `float` number. The smallest beta for the linear schedule. beta_max: A `float` number. The largest beta for the linear schedule. cosine_s: A `float` number. The hyperparameter in the cosine schedule. cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule. T: A `float` number. The ending time of the forward process. =============================================================== Args: schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs, 'linear' or 'cosine' for continuous-time DPMs. Returns: A wrapper object of the forward SDE (VP type). =============================================================== Example: # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1): >>> ns = NoiseScheduleVP('discrete', betas=betas) # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1): >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod) # For continuous-time DPMs (VPSDE), linear schedule: >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.) )rlinearcosinezZUnsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'rN?rdim?g)r)dtypeigMb?g8@@rgO@a?) ValueErrorformatscheduletorchlogcumsumlentotal_NTlinspacereshapetot_arraylog_alpha_arraybeta_0beta_1cosine_scosine_beta_maxmathatanpi cosine_t_maxcoscosine_log_alpha_0)selfrbetasalphas_cumprodcontinuous_beta_0continuous_beta_1r log_alphasr,U/home/kuhnn/.local/lib/python3.10/site-packages/TTS/tts/layers/tortoise/dpm_solver.py__init__sPY  , (  zNoiseScheduleVP.__init__csjdkrt|dj|jj|jdSjdkr3d|djjd|jSjdkrGfd d }||j }|Sd S) zT Compute log(alpha_t) of a given continuous-time label t in [0, T]. rr rr rgпrrcs*tt|jdjtjdS)Nr r )rrr$rr r")sr&r,r- log_alpha_fn*z=NoiseScheduleVP.marginal_log_mean_coeff..log_alpha_fnN) rinterpolate_fnrrrdevicerrrr%)r&tr3 log_alpha_tr,r2r-marginal_log_mean_coeffs    &  z'NoiseScheduleVP.marginal_log_mean_coeffcCst||S)zO Compute alpha_t of a given continuous-time label t in [0, T]. )rexpr9r&r7r,r,r-marginal_alphaszNoiseScheduleVP.marginal_alphac Cstdtd||S)zO Compute sigma_t of a given continuous-time label t in [0, T]. r r )rsqrtr:r9r;r,r,r- marginal_stdszNoiseScheduleVP.marginal_stdcCs.||}dtdtd|}||S)zn Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T]. rr r )r9rrr:)r&r7log_mean_coefflog_stdr,r,r-marginal_lambdas zNoiseScheduleVP.marginal_lambdacs jdkr2djjtd|td|}jd|}|t|jjjSjdkrjdttd|jd|}t | dt j |jd gt j |jd g}| d Sdtd|td|}fd d }||}|S) z` Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t. rr grr0rgr/r)r cs0tt|jddjtjjS)Nr r )rarccosr:r%rr r")r8r2r,r-t_fnsz,NoiseScheduleVP.inverse_lambda..t_fn)rrrr logaddexpzerosrr=r6r5rfliprr)r&lambtmpDelta log_alphar7rDr,r2r-inverse_lambdas ,  "    zNoiseScheduleVP.inverse_lambda) __name__ __module__ __qualname__rfloat32r.r9r<r>rArLr,r,r,r-rs  rnoiseuncondr c sj fddd  fdd fdd f dd } d vs-Jd vs3J| S) a!Create a wrapper function for the noise prediction model. DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to firstly wrap the model function to a noise prediction model that accepts the continuous time as the input. We support four types of the diffusion model by setting `model_type`: 1. "noise": noise prediction model. (Trained by predicting noise). 2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0). 3. "v": velocity prediction model. (Trained by predicting the velocity). The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2]. [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models." arXiv preprint arXiv:2202.00512 (2022). [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models." arXiv preprint arXiv:2210.02303 (2022). 4. "score": marginal score function. (Trained by denoising score matching). Note that the score function and the noise prediction model follows a simple relationship: ``` noise(x_t, t) = -sigma_t * score(x_t, t) ``` We support three types of guided sampling by DPMs by setting `guidance_type`: 1. "uncond": unconditional sampling by DPMs. The input `model` has the following format: `` model(x, t_input, **model_kwargs) -> noise | x_start | v | score `` 2. "classifier": classifier guidance sampling [3] by DPMs and another classifier. The input `model` has the following format: `` model(x, t_input, **model_kwargs) -> noise | x_start | v | score `` The input `classifier_fn` has the following format: `` classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond) `` [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis," in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794. 3. "classifier-free": classifier-free guidance sampling by conditional DPMs. The input `model` has the following format: `` model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score `` And if cond == `unconditional_condition`, the model output is the unconditional DPM output. [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance." arXiv preprint arXiv:2207.12598 (2022). The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999) or continuous-time labels (i.e. epsilon to T). We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise: `` def model_fn(x, t_continuous) -> noise: t_input = get_model_input_time(t_continuous) return noise_pred(model, x, t_input, **model_kwargs) `` where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver. =============================================================== Args: model: A diffusion model with the corresponding format described above. noise_schedule: A noise schedule object, such as NoiseScheduleVP. model_type: A `str`. The parameterization type of the diffusion model. "noise" or "x_start" or "v" or "score". model_kwargs: A `dict`. A dict for the other inputs of the model function. guidance_type: A `str`. The type of the guidance for sampling. "uncond" or "classifier" or "classifier-free". condition: A pytorch tensor. The condition for the guided sampling. Only used for "classifier" or "classifier-free" guidance type. unconditional_condition: A pytorch tensor. The condition for the unconditional sampling. Only used for "classifier-free" guidance type. guidance_scale: A `float`. The scale for the guided sampling. classifier_fn: A classifier function. Only used for the classifier guidance. classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function. Returns: A noise prediction model that accepts the noised data and the continuous time as the inputs. cs jdkr|djdS|S)a Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time. For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N]. For continuous-time DPMs, we just use `t_continuous`. rr g@@)rr) t_continuous)noise_scheduler,r-get_model_input_time:s z+model_wrapper..get_model_input_timeNcs|}|dur||fi}n |||fi}dkr"|Sdkr9||}}||||SdkrP||}}||||Sdkr^|}| |SdS)NrQx_startvscore)r<r>)xrScondt_inputoutputalpha_tsigma_t)rUmodel model_kwargs model_typerTr,r- noise_pred_fnEs   z$model_wrapper..noise_pred_fncsdt$|d}||fi}tj||dWdS1s+wYdS)z] Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t). TrN)r enable_graddetachrequires_grad_autogradgradsum)rYr[x_inlog_prob) classifier_fnclassifier_kwargs conditionr,r- cond_grad_fnWs $z#model_wrapper..cond_grad_fnc sdkr ||Sdkr.dusJ|}||}|}||}|||Sdkrldks:durA||dSt|gd}t|gd}tg}|||dd\} }| || SdS)zS The noise predicition model function that is used for DPM-Solver. rR classifierNclassifier-freer )rZr0)r>rcatchunk) rYrSr[ cond_gradr^rQrit_inc_in noise_uncond) rkrnrmrUguidance_scale guidance_typerbrTunconditional_conditionr,r-model_fn`s$     zmodel_wrapper..model_fn)rQrVrWrX)rRrorpNr,) r_rTrar`rxrmryrwrkrlrzr,) rkrlrnrmrUrwrxr_r`rarbrTryr- model_wrappers e   r|c@seZdZ     d9ddZddZd d Zd d Zd dZddZddZ ddZ d:ddZ    d;ddZ      dd$d%Zd=d&d'Z ( ) ( * + d?d,d-Zd@d.d/Z 0   1 2 3 4   ) ( dAd5d6Z 0   1 2 3 4   ) ( dAd7d8ZdS)B DPM_Solver dpmsolver++Nr ףp= ?csTfdd|_||_|dvsJ||_|dkr|j|_n||_||_||_||_dS)aN Construct a DPM-Solver. We support both DPM-Solver (`algorithm_type="dpmsolver"`) and DPM-Solver++ (`algorithm_type="dpmsolver++"`). We also support the "dynamic thresholding" method in Imagen[1]. For pixel-space diffusion models, you can set both `algorithm_type="dpmsolver++"` and `correcting_x0_fn="dynamic_thresholding"` to use the dynamic thresholding. The "dynamic thresholding" can greatly improve the sample quality for pixel-space DPMs with large guidance scales. Note that the thresholding method is **unsuitable** for latent-space DPMs (such as stable-diffusion). To support advanced algorithms in image-to-image applications, we also support corrector functions for both x0 and xt. Args: model_fn: A noise prediction model function which accepts the continuous-time input (t in [epsilon, T]): `` def model_fn(x, t_continuous): return noise `` The shape of `x` is `(batch_size, **shape)`, and the shape of `t_continuous` is `(batch_size,)`. noise_schedule: A noise schedule object, such as NoiseScheduleVP. algorithm_type: A `str`. Either "dpmsolver" or "dpmsolver++". correcting_x0_fn: A `str` or a function with the following format: ``` def correcting_x0_fn(x0, t): x0_new = ... return x0_new ``` This function is to correct the outputs of the data prediction model at each sampling step. e.g., ``` x0_pred = data_pred_model(xt, t) if correcting_x0_fn is not None: x0_pred = correcting_x0_fn(x0_pred, t) xt_1 = update(x0_pred, xt, t) ``` If `correcting_x0_fn="dynamic_thresholding"`, we use the dynamic thresholding proposed in Imagen[1]. correcting_xt_fn: A function with the following format: ``` def correcting_xt_fn(xt, t, step): x_new = ... return x_new ``` This function is to correct the intermediate samples xt at each sampling step. e.g., ``` xt = ... xt = correcting_xt_fn(xt, t, step) ``` thresholding_max_val: A `float`. The max value for thresholding. Valid only when use `dpmsolver++` and `correcting_x0_fn="dynamic_thresholding"`. dynamic_thresholding_ratio: A `float`. The ratio for dynamic thresholding (see Imagen[1] for details). Valid only when use `dpmsolver++` and `correcting_x0_fn="dynamic_thresholding"`. [1] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487, 2022b. cs|||jdS)Nr)expandshape)rYr7rzr,r-sz%DPM_Solver.__init__..) dpmsolverr~dynamic_thresholdingN)r_rTalgorithm_typedynamic_thresholding_fncorrecting_x0_fncorrecting_xt_fndynamic_thresholding_ratiothresholding_max_val)r&rzrTrrrrrr,rr-r.}sB   zDPM_Solver.__init__cCsr|}|j}tjt||jddf|dd}tt||j t | |j |}t || ||}|S)z2 The dynamic thresholding method. rr rr)r rrquantileabsrr expand_dimsmaximumr ones_likerr6clamp)r&x0r7dimspr1r,r,r-rs&z"DPM_Solver.dynamic_thresholding_fncC |||S)z4 Return the noise prediction model. )r_r&rYr7r,r,r-noise_prediction_fn zDPM_Solver.noise_prediction_fncCsP|||}|j||j|}}||||}|jdur&|||}|S)zD Return the data prediction model (with corrector). N)rrTr<r>r)r&rYr7rQr]r^rr,r,r-data_prediction_fns   zDPM_Solver.data_prediction_fncCs"|jdkr |||S|||S)z_ Convert the model to the noise prediction model or the data prediction model. r~)rrrrr,r,r-rzs   zDPM_Solver.model_fnc Cs|dkr6|jt||}|jt||}t|||d|}|j|S|dkrFt|||d|S|dkred} t|d| |d| |d | |} | St d |)a7Compute the intermediate time steps for sampling. Args: skip_type: A `str`. The type for the spacing of the time steps. We support three types: - 'logSNR': uniform logSNR for the time steps. - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.) - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.) t_T: A `float`. The starting time of the sampling (default is T). t_0: A `float`. The ending time of the sampling (default is epsilon). N: A `int`. The total number of the spacing of the time steps. device: A torch device. Returns: A pytorch tensor of the time steps, with the shape (N + 1,). logSNRr time_uniformtime_quadraticr0r zSUnsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic') rTrArtensorrrcpuitemrLpowrr) r& skip_typet_Tt_0Nr6lambda_Tlambda_0 logSNR_stepst_orderr7r,r,r-get_time_stepss( .zDPM_Solver.get_time_stepsc Cs4|dkr8|dd}|ddkrdg|dddg}nQ|ddkr-dg|ddg}n@dg|ddg}n5|dkr]|ddkrL|d}dg|}n!|dd}dg|ddg}n|dkrid}dg|}ntd|dkr~||||||} | |fS||||||ttdg|d|} | |fS)a Get the order of each step for sampling by the singlestep DPM-Solver. We combine both DPM-Solver-1,2,3 to use all the function evaluations, which is named as "DPM-Solver-fast". Given a fixed number of function evaluations by `steps`, the sampling procedure by DPM-Solver-fast is: - If order == 1: We take `steps` of DPM-Solver-1 (i.e. DDIM). - If order == 2: - Denote K = (steps // 2). We take K or (K + 1) intermediate time steps for sampling. - If steps % 2 == 0, we use K steps of DPM-Solver-2. - If steps % 2 == 1, we use K steps of DPM-Solver-2 and 1 step of DPM-Solver-1. - If order == 3: - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling. - If steps % 3 == 0, we use (K - 2) steps of DPM-Solver-3, and 1 step of DPM-Solver-2 and 1 step of DPM-Solver-1. - If steps % 3 == 1, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-1. - If steps % 3 == 2, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-2. ============================================ Args: order: A `int`. The max order for the solver (2 or 3). steps: A `int`. The total number of function evaluations (NFE). skip_type: A `str`. The type for the spacing of the time steps. We support three types: - 'logSNR': uniform logSNR for the time steps. - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.) - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.) t_T: A `float`. The starting time of the sampling (default is T). t_0: A `float`. The ending time of the sampling (default is epsilon). device: A torch device. Returns: orders: A list of the solver order of each step. rrr0z"'order' must be '1' or '2' or '3'.r)rrrrrr) r&stepsorderrrrr6Korderstimesteps_outerr,r,r-.get_orders_and_timesteps_for_singlestep_solversv       z9DPM_Solver.get_orders_and_timesteps_for_singlestep_solvercCr)z Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization. )r)r&rYr1r,r,r-denoise_to_zero_fnkrzDPM_Solver.denoise_to_zero_fnFcCs|j}|}||||}} | |} ||||} } ||||} }t| }|jdkr\t| }|durF| ||}|| ||||}|rZ|d|ifS|St| }|durk| ||}t| | ||||}|r|d|ifS|S)a DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`. Args: x: A pytorch tensor. The initial value at time `s`. s: A pytorch tensor. The starting time, with the shape (1,). t: A pytorch tensor. The ending time, with the shape (1,). model_s: A pytorch tensor. The model function evaluated at time `s`. If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it. return_intermediate: A `bool`. If true, also return the model value at time `s`. Returns: x_t: A pytorch tensor. The approximated solution at time `t`. r~Nmodel_s) rTr rAr9r>rr:rexpm1rz)r&rYr1r7rreturn_intermediatensrlambda_slambda_th log_alpha_sr8sigma_sr^r]phi_1x_tr,r,r-dpm_solver_first_updateqs,        z"DPM_Solver.dpm_solver_first_updaterrcCs|dvr td||durd}|j}||||} } | | } | || } || } |||| ||}}}|||| ||}}}t|t|}}|j dkrt | | }t | }|durx| ||}||||||}| || }|dkr||||||d|||||}n|dkr||||||d|||| d||}npt || }t | }|dur| ||}t||||||}| || }|dkrt||||||d|||||}n$|dkr8t||||||d|||| d||}|rB|||d fS|S) a Singlestep solver DPM-Solver-2 from time `s` to time `t`. Args: x: A pytorch tensor. The initial value at time `s`. s: A pytorch tensor. The starting time, with the shape (1,). t: A pytorch tensor. The ending time, with the shape (1,). r1: A `float`. The hyperparameter of the second-order solver. model_s: A pytorch tensor. The model function evaluated at time `s`. If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it. return_intermediate: A `bool`. If true, also return the model value at time `s` and `s1` (the intermediate time). solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers. The type slightly impacts the performance. We recommend to use 'dpmsolver' type. Returns: x_t: A pytorch tensor. The approximated solution at time `t`. rtaylor<'solver_type' must be either 'dpmsolver' or 'taylor', got {}Nrr~rrr )rmodel_s1 rrrTrArLr9r>rr:rrrz)r&rYr1r7r1rr solver_typerrrr lambda_s1s1r log_alpha_s1r8rsigma_s1r^alpha_s1r]phi_11rx_s1rrr,r,r-#singlestep_dpm_solver_second_updates~                   z.DPM_Solver.singlestep_dpm_solver_second_updateUUUUUU?UUUUUU?c +Csr| dvr td| |durd}|durd}|j} | || |} } | | } | || }| || }| |}| |}| || || || |f\}}}}| || || || |f\}}}}t|t|t|}}}|j dkrSt | | }t | | }t | }t | | || d} || d}!|!| d}"|dur| ||}|dur||||||}#| |#|}|||||||||| ||}$| |$|}%| d kr||||||d|||!|%|}&n| d krRd|||}'d||%|}(||'||(||})d |(|'||}*||||||||!|)||"|*}&nt || }t || }t | }t || || d} || d}!|!| d}"|dur| ||}|durt||||||}#| |#|}t|||||||||| ||}$| |$|}%| d krt||||||d|||!|%|}&nF| d kr,d|||}'d||%|}(||'||(||})d |(|'||}*t||||||||!|)||"|*}&|r7|&|||%d fS|&S) a Singlestep solver DPM-Solver-3 from time `s` to time `t`. Args: x: A pytorch tensor. The initial value at time `s`. s: A pytorch tensor. The starting time, with the shape (1,). t: A pytorch tensor. The ending time, with the shape (1,). r1: A `float`. The hyperparameter of the third-order solver. r2: A `float`. The hyperparameter of the third-order solver. model_s: A pytorch tensor. The model function evaluated at time `s`. If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it. model_s1: A pytorch tensor. The model function evaluated at time `s1` (the intermediate time given by `r1`). If `model_s1` is None, we evaluate the model at `s1`; otherwise we directly use it. return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times). solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers. The type slightly impacts the performance. We recommend to use 'dpmsolver' type. Returns: x_t: A pytorch tensor. The approximated solution at time `t`. rrNrrr~r rrrr )rrmodel_s2r)+r&rYr1r7rr2rrrrrrrrr lambda_s2rs2rr log_alpha_s2r8rrsigma_s2r^ralpha_s2r]rphi_12rphi_22phi_2phi_3rx_s2rrD1_0D1_1D1D2r,r,r-"singlestep_dpm_solver_third_updates                                      z-DPM_Solver.singlestep_dpm_solver_third_updatecCs|dvr td||j}|d|d}}|d|d} } || || ||} } } || ||}}|| ||}}t|}| | }| | }||}d|||}|jdkrt | }|dkr||||||d|||}|S|d kr|||||||||d|}|St |}|dkrt||||||d|||}|S|d krt|||||||||d|}|S) a Multistep solver DPM-Solver-2 from time `t_prev_list[-1]` to time `t`. Args: x: A pytorch tensor. The initial value at time `s`. model_prev_list: A list of pytorch tensor. The previous computed model values. t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (1,) t: A pytorch tensor. The ending time, with the shape (1,). solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers. The type slightly impacts the performance. We recommend to use 'dpmsolver' type. Returns: x_t: A pytorch tensor. The approximated solution at time `t`. rrr r r~rrr) rrrTrAr9r>rr:rr)r&rYmodel_prev_list t_prev_listr7rr model_prev_1 model_prev_0t_prev_1t_prev_0 lambda_prev_1 lambda_prev_0rlog_alpha_prev_0r8 sigma_prev_0r^r]h_0rr0rrrr,r,r-"multistep_dpm_solver_second_updates\    (      z-DPM_Solver.multistep_dpm_solver_second_updatec#Cs|j}|\}}} |\} } } || || || ||f\} }}}|| ||}}|| ||}}t|}|| }||}||}||||}}d|| |}d|||}||||||}d||||}|jdkrt| }||d} | |d}!|||||| || |||!|}"|"St|}||d} | |d}!t|||||| || |||!|}"|"S)a Multistep solver DPM-Solver-3 from time `t_prev_list[-1]` to time `t`. Args: x: A pytorch tensor. The initial value at time `s`. model_prev_list: A list of pytorch tensor. The previous computed model values. t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (1,) t: A pytorch tensor. The ending time, with the shape (1,). solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers. The type slightly impacts the performance. We recommend to use 'dpmsolver' type. Returns: x_t: A pytorch tensor. The approximated solution at time `t`. r r~r)rTrAr9r>rr:rr)#r&rYrrr7rr model_prev_2rrt_prev_2rr lambda_prev_2rrrrr8rr^r]h_1rrrrrrrrrrrrr,r,r-!multistep_dpm_solver_third_updatesX                 z,DPM_Solver.multistep_dpm_solver_third_updatec Csf|dkr |j||||dS|dkr|j||||||dS|dkr,|j|||||||dStd|)a Singlestep DPM-Solver with the order `order` from time `s` to time `t`. Args: x: A pytorch tensor. The initial value at time `s`. s: A pytorch tensor. The starting time, with the shape (1,). t: A pytorch tensor. The ending time, with the shape (1,). order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3. return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times). solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers. The type slightly impacts the performance. We recommend to use 'dpmsolver' type. r1: A `float`. The hyperparameter of the second-order or third-order solver. r2: A `float`. The hyperparameter of the third-order solver. Returns: x_t: A pytorch tensor. The approximated solution at time `t`. rrr0)rrrr)rrrr(Solver order must be 1 or 2 or 3, got {})rrrrr) r&rYr1r7rrrrrr,r,r-singlestep_dpm_solver_updates, z'DPM_Solver.singlestep_dpm_solver_updatecCsh|dkr|j||d||ddS|dkr|j|||||dS|dkr-|j|||||dStd|)a0 Multistep DPM-Solver with the order `order` from time `t_prev_list[-1]` to time `t`. Args: x: A pytorch tensor. The initial value at time `s`. model_prev_list: A list of pytorch tensor. The previous computed model values. t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (1,) t: A pytorch tensor. The ending time, with the shape (1,). order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3. solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers. The type slightly impacts the performance. We recommend to use 'dpmsolver' type. Returns: x_t: A pytorch tensor. The approximated solution at time `t`. rr )rr0rrr)rrrrr)r&rYrrr7rrr,r,r-multistep_dpm_solver_update(sz&DPM_Solver.multistep_dpm_solver_update皙?q??h㈵>c  sj} |td|} | | } | |t| |}|t| |}|}d}|dkrBdfdd}fdd}n!|d kr\d \fd d}fd d}ntd |t| | | kr| | |}||| |\}}||| |fi|}t t||||t t|t|}dd}|||| }t |dkr|}|} |}| | } t ||t|d||| }||7}t| | | ksntd||S)u The adaptive step size solver based on singlestep DPM-Solver. Args: x: A pytorch tensor. The initial value at time `t_T`. order: A `int`. The (higher) order of the solver. We only support order == 2 or 3. t_T: A `float`. The starting time of the sampling (default is T). t_0: A `float`. The ending time of the sampling (default is epsilon). h_init: A `float`. The initial step size (for logSNR). atol: A `float`. The absolute tolerance of the solver. For image data, the default setting is 0.0078, followed [1]. rtol: A `float`. The relative tolerance of the solver. The default setting is 0.05. theta: A `float`. The safety hyperparameter for adapting the step size. The default setting is 0.9, followed [1]. t_err: A `float`. The tolerance for the time. We solve the diffusion ODE until the absolute error between the current time and `t_0` is less than `t_err`. The default setting is 1e-5. solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers. The type slightly impacts the performance. We recommend to use 'dpmsolver' type. Returns: x_0: A pytorch tensor. The approximated solution at time `t_0`. [1] A. Jolicoeur-Martineau, K. Li, R. Piché-Taillefer, T. Kachman, and I. Mitliagkas, "Gotta go fast when generating data with score-based models," arXiv preprint arXiv:2105.14080, 2021. rBrr0rcsj|||ddS)NTr)rrYr1r7r2r,r- lower_updatelsz4DPM_Solver.dpm_solver_adaptive..lower_updatecsj|||fd|S)N)rrrrYr1r7kwargsrr&rr,r- higher_updateosz5DPM_Solver.dpm_solver_adaptive..higher_updater)rrcsj|||ddS)NT)rrrrrrr,r-rus csj|||fd|S)N)rrr)rrrrr&rr,r-rzsz;For adaptive step size solver, order must be 2 or 3, got {}cSs*tt||jddfjdddS)Nrr T)r keepdim)rr=squarerrmean)rWr,r,r-norm_fnr4z/DPM_Solver.dpm_solver_adaptive..norm_fnr gzadaptive solver nfe)rTronesrrArrrrrrLmaxallmin float_powerfloatprint)r&rYrrrh_initatolrtolthetat_errrrr1rrrx_prevnferrr7x_lowerlower_noise_kwargsx_higherdeltarEr,rr-dpm_solver_adaptive@sN"    zDPM_Solver.dpm_solver_adaptivecCs|j||j|}}|dur!tj|jdg|jR|jd}|dg|jR}t|| |t|| |}|jddkrI| dS|S)a# Compute the noised input xt = alpha_t * x + sigma_t * noise. Args: x: A `torch.Tensor` with shape `(batch_size, *shape)`. t: A `torch.Tensor` with shape `(t_size,)`. Returns: xt with shape `(t_size, batch_size, *shape)`. Nrr6r r) rTr<r>rrandnrr6rrr squeeze)r&rYr7rQr]r^xtr,r,r- add_noises  $ zDPM_Solver.add_noiser0r multistepTcCsh|dur d|jjn|}|dur|jjn|}|dkr|dks"Jd|j||||||||| | | | | d S)z Inverse the sample `x` from time `t_start` to `t_end` by DPM-Solver. For discrete-time DPMs, we use `t_start=1/N`, where `N` is the total time steps during training. Nr rTime range needs to be greater than 0. For discrete-time DPMs, it needs to be in [1 / N, 1], where N is the length of betas array) rt_startt_endrrmethodlower_order_finaldenoise_to_zerorrrr)rTrrsample)r&rYrr$r%rrr&r'r(rrrrrrr,r,r-inverses(zDPM_Solver.inversec" CsX|dur d|jjn|}|dur|jjn|}|dkr|dks"Jd| r,|dvs,Jd|jdur9|dvs9Jd|j}g}t|dkrV|j||||| | | d }n|d kr=||ksaJ|j|||||d }|j dd |ksvJd}||}|g}| ||g}|jdur||||}| r| |t d |D]2}||}|j |||||| d }|jdur||||}| r| || || | ||qt ||d D]c}||}|r|dkrt||d |}n|}|j |||||| d }|jdur||||}| r| |t |d D]}||d ||<||d ||<q||d<||kr;| |||d<qn|dvr|dkrU|j||||||d\}}n|dkrm||}|g|}|j|||||d }t|D]o\}}||||d }}|j|||||d }|j|}|d|d}|d krdn |d |d|} |dkrdn |d|d|}!|j||||| | |!d}|jdur||||}| r| |qqntd|| rtd||}|||}|jdur ||||d }| r| |Wdn 1swY| r*||fS|S)a6 Compute the sample at time `t_end` by DPM-Solver, given the initial `x` at time `t_start`. ===================================================== We support the following algorithms for both noise prediction model and data prediction model: - 'singlestep': Singlestep DPM-Solver (i.e. "DPM-Solver-fast" in the paper), which combines different orders of singlestep DPM-Solver. We combine all the singlestep solvers with order <= `order` to use up all the function evaluations (steps). The total number of function evaluations (NFE) == `steps`. Given a fixed NFE == `steps`, the sampling procedure is: - If `order` == 1: - Denote K = steps. We use K steps of DPM-Solver-1 (i.e. DDIM). - If `order` == 2: - Denote K = (steps // 2) + (steps % 2). We take K intermediate time steps for sampling. - If steps % 2 == 0, we use K steps of singlestep DPM-Solver-2. - If steps % 2 == 1, we use (K - 1) steps of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1. - If `order` == 3: - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling. - If steps % 3 == 0, we use (K - 2) steps of singlestep DPM-Solver-3, and 1 step of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1. - If steps % 3 == 1, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of DPM-Solver-1. - If steps % 3 == 2, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of singlestep DPM-Solver-2. - 'multistep': Multistep DPM-Solver with the order of `order`. The total number of function evaluations (NFE) == `steps`. We initialize the first `order` values by lower order multistep solvers. Given a fixed NFE == `steps`, the sampling procedure is: Denote K = steps. - If `order` == 1: - We use K steps of DPM-Solver-1 (i.e. DDIM). - If `order` == 2: - We firstly use 1 step of DPM-Solver-1, then use (K - 1) step of multistep DPM-Solver-2. - If `order` == 3: - We firstly use 1 step of DPM-Solver-1, then 1 step of multistep DPM-Solver-2, then (K - 2) step of multistep DPM-Solver-3. - 'singlestep_fixed': Fixed order singlestep DPM-Solver (i.e. DPM-Solver-1 or singlestep DPM-Solver-2 or singlestep DPM-Solver-3). We use singlestep DPM-Solver-`order` for `order`=1 or 2 or 3, with total [`steps` // `order`] * `order` NFE. - 'adaptive': Adaptive step size DPM-Solver (i.e. "DPM-Solver-12" and "DPM-Solver-23" in the paper). We ignore `steps` and use adaptive step size DPM-Solver with a higher order of `order`. You can adjust the absolute tolerance `atol` and the relative tolerance `rtol` to balance the computatation costs (NFE) and the sample quality. - If `order` == 2, we use DPM-Solver-12 which combines DPM-Solver-1 and singlestep DPM-Solver-2. - If `order` == 3, we use DPM-Solver-23 which combines singlestep DPM-Solver-2 and singlestep DPM-Solver-3. ===================================================== Some advices for choosing the algorithm: - For **unconditional sampling** or **guided sampling with small guidance scale** by DPMs: Use singlestep DPM-Solver or DPM-Solver++ ("DPM-Solver-fast" in the paper) with `order = 3`. e.g., DPM-Solver: >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver") >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3, skip_type='time_uniform', method='singlestep') e.g., DPM-Solver++: >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++") >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3, skip_type='time_uniform', method='singlestep') - For **guided sampling with large guidance scale** by DPMs: Use multistep DPM-Solver with `algorithm_type="dpmsolver++"` and `order = 2`. e.g. >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++") >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=2, skip_type='time_uniform', method='multistep') We support three types of `skip_type`: - 'logSNR': uniform logSNR for the time steps. **Recommended for low-resolutional images** - 'time_uniform': uniform time for the time steps. **Recommended for high-resolutional images**. - 'time_quadratic': quadratic time for the time steps. ===================================================== Args: x: A pytorch tensor. The initial value at time `t_start` e.g. if `t_start` == T, then `x` is a sample from the standard normal distribution. steps: A `int`. The total number of function evaluations (NFE). t_start: A `float`. The starting time of the sampling. If `T` is None, we use self.noise_schedule.T (default is 1.0). t_end: A `float`. The ending time of the sampling. If `t_end` is None, we use 1. / self.noise_schedule.total_N. e.g. if total_N == 1000, we have `t_end` == 1e-3. For discrete-time DPMs: - We recommend `t_end` == 1. / self.noise_schedule.total_N. For continuous-time DPMs: - We recommend `t_end` == 1e-3 when `steps` <= 15; and `t_end` == 1e-4 when `steps` > 15. order: A `int`. The order of DPM-Solver. skip_type: A `str`. The type for the spacing of the time steps. 'time_uniform' or 'logSNR' or 'time_quadratic'. method: A `str`. The method for sampling. 'singlestep' or 'multistep' or 'singlestep_fixed' or 'adaptive'. denoise_to_zero: A `bool`. Whether to denoise to time 0 at the final step. Default is `False`. If `denoise_to_zero` is `True`, the total NFE is (`steps` + 1). This trick is firstly proposed by DDPM (https://arxiv.org/abs/2006.11239) and score_sde (https://arxiv.org/abs/2011.13456). Such trick can improve the FID for diffusion models sampling by diffusion SDEs for low-resolutional images (such as CIFAR-10). However, we observed that such trick does not matter for high-resolutional images. As it needs an additional NFE, we do not recommend it for high-resolutional images. lower_order_final: A `bool`. Whether to use lower order solvers at the final steps. Only valid for `method=multistep` and `steps < 15`. We empirically find that this trick is a key to stabilizing the sampling by DPM-Solver with very few steps (especially for steps <= 10). So we recommend to set it to be `True`. solver_type: A `str`. The taylor expansion type for the solver. `dpmsolver` or `taylor`. We recommend `dpmsolver`. atol: A `float`. The absolute tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'. rtol: A `float`. The relative tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'. return_intermediate: A `bool`. Whether to save the xt at each step. When set to `True`, method returns a tuple (x0, intermediates); when set to False, method returns only x0. Returns: x_end: A pytorch tensor. The approximated solution at time `t_end`. Nr rr#)r" singlestepsinglestep_fixedz:Cannot use adaptive solver when saving intermediate valuesz *       r5cCs|dd|dS)z Expand the tensor `v` to the dim `dims`. Args: `v`: a PyTorch tensor with shape [N]. `dim`: a `int`. Returns: a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`. ).r{rr,)rWrr,r,r-rs r)r rrr|r}r5rr,r,r,r-s2S (n .