o . i,@shddlmZddlZddlZddlmZddlmZGdddeZ Gdd d eZ Gd d d eZ dS) ) annotationsN)StrEnum) Schedulerc@s eZdZdZdZdZdZdZdS)DDPMVarianceTypez Valid names for DDPM Scheduler's `variance_type` argument. Options to clip the variance used when adding noise to the denoised sample. Z fixed_smallZ fixed_largelearned learned_rangeN)__name__ __module__ __qualname____doc__ FIXED_SMALL FIXED_LARGELEARNED LEARNED_RANGErr`/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/schedulers/ddpm.pyr*s rc@seZdZdZdZdZdZdS)DDPMPredictionTypea4 Set of valid prediction type names for the DDPM scheduler's `prediction_type` argument. epsilon: predicting the noise of the diffusion process sample: directly predicting the noisy sample v_prediction: velocity prediction, see section 2.4 https://imagen.research.google/video/paper.pdf epsilonsample v_predictionN)r r r r EPSILONSAMPLE V_PREDICTIONrrrrr6s rcsfeZdZdZddejdejddfd-fdd Zd.d/ddZ d0d d!Z d.d1d$d%Z d.d2d+d,Z Z S)3 DDPMScheduleraZ Denoising diffusion probabilistic models (DDPMs) explores the connections between denoising score matching and Langevin dynamics sampling. Based on: Ho et al., "Denoising Diffusion Probabilistic Models" https://arxiv.org/abs/2006.11239 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 variance_type: member of DDPMVarianceType clip_sample: option to clip predicted sample between -1 and 1 for numerical stability. prediction_type: member of DDPMPredictionType clip_sample_min: minimum clipping value when clip_sample equals True clip_sample_max: maximum clipping value when clip_sample equals True schedule_args: arguments to pass to the schedule function i linear_betaTgg?num_train_timestepsintschedulestr variance_type clip_sampleboolprediction_typeclip_sample_minfloatclip_sample_maxreturnNonec  sbtj||fi||tjvrtd|tjvr!td||_||g|_||_ ||_ dS)Nz?Argument `variance_type` must be a member of `DDPMVarianceType`zCArgument `prediction_type` must be a member of `DDPMPredictionType`) super__init__r __members__values ValueErrorrr!clip_sample_valuesr r#) selfrrr r!r#r$r& schedule_args __class__rrr*Us   zDDPMScheduler.__init__Nnum_inference_stepsdevicestr | torch.device | NonecCsx||jkrtd|d|jd|jd||_|j|j}td||dddtj}t | ||_ 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.rN) rr-r3nparangeroundastypeint64torch from_numpyto timesteps)r/r3r4 step_ratior?rrr set_timestepsms  &zDDPMScheduler.set_timestepstimestepx_0 torch.Tensorx_tc Csv|j|}|j|}|dkr|j|dn|j}||j|d|}|d|d|}||||} | S)z Compute the mean of the posterior at timestep t. Args: timestep: current timestep. x0: the noise-free input. x_t: the input noised to timestep t. Returns: Returns the mean rr)alphasalphas_cumprodonesqrtbetas) r/rBrCrEalpha_t alpha_prod_talpha_prod_t_prevZx_0_coefficientZx_t_coefficientmeanrrr _get_means  zDDPMScheduler._get_meanpredicted_variancetorch.Tensor | Nonec Cs|j|}|dkr|j|dn|j}d|d||j|}|jtjkr/tj|dd}|S|jtjkr<|j|}|S|jtj krH|durH|S|jtj kri|duri|}|j|}|dd}||d||}|S)z Compute the variance of the posterior at timestep t. Args: timestep: current timestep. predicted_variance: variance predicted by the model. Returns: Returns the variance rrg#B ;)minN) rGrHrJr rr r<clamprrr) r/rBrPrLrMvarianceZmin_logZmax_logfracrrr _get_variances"      zDDPMScheduler._get_variance model_outputr generatortorch.Generator | None!tuple[torch.Tensor, torch.Tensor]cCs|jd|jddkr|jdvrtj||jddd\}}nd}|j|}|dkr1|j|dn|j}d|}d|} |jtjkrO||d||d} n|jtj krX|} n|jtj krj|d||d|} |j rzt | |j d|j d} |d|j||} |j|d| |} | | | |} td}|dkrtj||j|j||jd}|j||d d|}| |} | | fS) a7 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). 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. generator: random number generator. Returns: pred_prev_sample: Predicted previous sample rrS)rr)dimNrg?)dtypelayoutrYr4)rP)shaper r<splitrGrHr#rrrrr!rTr.rJrFtensorrandnsizer]r^r4rW)r/rXrBrrYrPrLrM beta_prod_tbeta_prod_t_prevpred_original_samplepred_original_sample_coeffcurrent_sample_coeffpred_prev_samplerUnoiserrrsteps@"     zDDPMScheduler.step)rrrrr rr!r"r#rr$r%r&r%r'r()N)r3rr4r5r'r()rBrrCrDrErDr'rD)rBrrPrQr'rD) rXrDrBrrrDrYrZr'r[)r r r r rr rrr*rArOrWrk __classcell__rrr1rrDs   "r) __future__rnumpyr7r< monai.utilsr schedulerrrrrrrrrs