o . i8@sdddlmZddlmZddlZddlZddlmZddl m Z GdddeZ Gd d d e Z dS) ) annotations)AnyN)StrEnum) Schedulerc@seZdZdZdZdZdS)PNDMPredictionTypea Set of valid prediction type names for the PNDM scheduler's `prediction_type` argument. epsilon: predicting the noise of the diffusion process v_prediction: velocity prediction, see section 2.4 https://imagen.research.google/video/paper.pdf epsilon v_predictionN)__name__ __module__ __qualname____doc__EPSILON V_PREDICTIONrr`/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/schedulers/pndm.pyr,srcsfeZdZdZddddejdfd(fdd Zd)d*ddZd+ddZd,d d!Z d-d#d$Z d.d&d'Z Z S)/ PNDMScheduleraS Pseudo numerical methods for diffusion models (PNDM) proposes using more advanced ODE integration techniques, namely Runge-Kutta method and a linear multi-step method. Based on: Liu et al., "Pseudo Numerical Methods for Diffusion Models on Manifolds" https://arxiv.org/abs/2202.09778 Args: num_train_timesteps: number of diffusion steps used to train the model. schedule: member of NoiseSchedules, name of noise schedule function in component store skip_prk_steps: allows the scheduler to skip the Runge-Kutta steps that are defined in the original paper as being required before plms step. set_alpha_to_one: each diffusion step uses the value of alphas product at that step and at the previous one. For the final step there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`, otherwise it uses the value of alpha at step 0. prediction_type: member of DDPMPredictionType steps_offset: an offset added to the inference steps. You can use a combination of `offset=1` and `set_alpha_to_one=False`, to make the last step use step 0 for the previous alpha product, as done in stable diffusion. schedule_args: arguments to pass to the schedule function i linear_betaFrnum_train_timestepsintschedulestrskip_prk_stepsboolset_alpha_to_oneprediction_type steps_offsetreturnNonec stj||fi||tjvrtd||_|r tdn|j d|_ d|_ d|_ ||_ ||_t|_d|_t|_g|_||dS)NzAArgument `prediction_type` must be a member of PNDMPredictionTypeg?r)super__init__r __members__values ValueErrorrtorchtensoralphas_cumprodfinal_alpha_cumprodinit_noise_sigma pndm_orderrrTensorcur_model_outputcounter cur_sampleets set_timesteps)selfrrrrrr schedule_args __class__rrr!Ps   zPNDMScheduler.__init__Nnum_inference_stepsdevicestr | torch.device | NonecCs`||jkrtd|d|jd|jd||_|j|j}td||tj|_|j|j 7_|j rHt g|_ |jddd|_ nDt |j|j ddtt d|j|dg|j}|dddd dddd|_ |jdd ddd|_ t|j |j gtj}t|||_t|j|_g|_d|_dS) a/ Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference. Args: num_inference_steps: number of diffusion steps used when generating samples with a pre-trained model. device: target device to put the data. z`num_inference_steps`: z3 cannot be larger than `self.num_train_timesteps`: zG as the unet model trained with this scheduler can only handle maximal z timesteps.rNr)rr$r5nparangeroundastypeint64Z _timestepsrrarray prk_timestepsZplms_timestepsr*repeattilecopy concatenater% from_numpyto timestepslenr/r-)r1r5r6 step_ratiorArHrrrr0ws:    *   zPNDMScheduler.set_timesteps model_output torch.Tensortimestepsampletuple[torch.Tensor, Any]cCs>|jt|jkr|js|j|||ddfS|j|||ddfS)an Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion process from the learned model outputs (most often the predicted noise). This function calls `step_prk()` or `step_plms()` depending on the internal variable `counter`. Args: model_output: direct output from learned diffusion model. timestep: current discrete timestep in the diffusion chain. sample: current instance of sample being created by diffusion process. Returns: pred_prev_sample: Predicted previous sample )rKrMrNN)r-rIrArstep_prk step_plms)r1rKrMrNrrrstepszPNDMScheduler.stepcCs(|jdur td|jdrdn|j|jd}||}|j|jdd}|jddkrs