o . i,8@sLddlmZddlZddlZddlmZddlmZeZ GdddeZ dS)) annotationsN)DDPMPredictionType) SchedulercsfeZdZdZdddddejddfd.fdd Zd/d0ddZd1d!d"Z # d2d3d*d+Z d4d,d-Z Z S)5 DDIMSchedulera Denoising diffusion implicit models is a scheduler that extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with non-Markovian guidance. Based on: Song et al. "Denoising Diffusion Implicit Models" https://arxiv.org/abs/2010.02502 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 clip_sample: option to clip predicted sample between -1 and 1 for numerical stability. 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. steps_offset: an offset added to the inference steps. You can use a combination of `steps_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. 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 iZ linear_betaTrg?num_train_timestepsintschedulestr clip_sampleboolset_alpha_to_one steps_offsetprediction_typeclip_sample_minfloatclip_sample_maxreturnNonec  stj||fi| |tjvrtd||_|r tdn|j d|_ d|_ t t d|jdddt j|_||_||g|_||_|||jdS)NzAArgument `prediction_type` must be a member of DDIMPredictionTyperr)super__init__DDIMPredictionType __members__values ValueErrorrtorchtensoralphas_cumprodfinal_alpha_cumprodZinit_noise_sigma from_numpynparangerastypeint64 timestepsr clip_sample_valuesr set_timesteps) selfrr r rrrrrZ schedule_args __class__`/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/schedulers/ddim.pyrBs ( zDDIMScheduler.__init__Nnum_inference_stepsdevicestr | torch.device | NonecCs||jkrtd|d|jd|jd||_|j|j}|j|kr/td|jd|dtd||d d d tj }t | ||_ |j |j7_ d S) 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.z`steps_offset`: zR cannot be greater than or equal to `num_train_timesteps // num_inference_steps : z@` as this will cause timesteps to exceed the max train timestep.rNr)rrr.rr"r#roundcopyr$r%rr!tor&)r)r.r/ step_ratior&r,r,r-r(hs&    *zDDIMScheduler.set_timestepstimestep prev_timestep torch.TensorcCsJ|j|}|dkr|j|n|j}d|}d|}||d||}|S)Nrr)rr )r)r5r6 alpha_prod_talpha_prod_t_prev beta_prod_tbeta_prod_t_prevvariancer,r,r- _get_variances zDDIMScheduler._get_variance model_outputsampleeta generatortorch.Generator | None!tuple[torch.Tensor, torch.Tensor]cCs||j|j}|j|}|dkr|j|n|j}d|} |} |} |jtjkr6|| d||d} |} n3|jtjkrK|} ||d| | d} n|jtjkri|d|| d|} |d|| d|} |j ryt | |j d|j d} | ||} || d} d|| dd| }|d| |}|dkrt t |r|jnd}t j|j|j||d}| ||d||} || }|| fS)a 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. eta: weight of noise for added noise in diffusion step. generator: random number generator. Returns: pred_prev_sample: Predicted previous sample pred_original_sample: Predicted original sample rr?cpu)dtyperBr/)rr.rr rrEPSILONSAMPLE V_PREDICTIONr rclampr'r=r/ is_tensorrandnshaperH)r)r?r5r@rArBr6r8r9r:pred_original_sample pred_epsilonr<Z std_dev_tpred_sample_directionZpred_prev_sampler/noiser,r,r-steps:#      zDDIMScheduler.stepc Cs||j|j}|j|}|dkr|j|n|j}d|}|}|} |jtjkr6||d||d}|} n3|jtjkrK|}||d||d} n|jtjkri|d||d|}|d||d|} |j ryt ||j d|j d}d|d| } |d|| } | |fS)a? Predict the sample at the next 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. Returns: pred_prev_sample: Predicted previous sample pred_original_sample: Predicted original sample rrrE) rr.rr rrrIrJrKr rrLr') r)r?r5r@r6r8r9r:rPrQrRZpred_post_sampler,r,r- reversed_steps,    zDDIMScheduler.reversed_step)rr r r r r rr rr rr rrrrrr)N)r.r r/r0rr)r5r r6r rr7)r>N) r?r7r5r r@r7rArrBrCrrD)r?r7r5r r@r7rrD) __name__ __module__ __qualname____doc__rrIrr(r=rTrU __classcell__r,r,r*r-r+s" &  Vr) __future__rnumpyr"rddpmr schedulerrrrr,r,r,r-s