o . i5@sddlmZddlmZddlZddlZddlmZddl m Z ddl m Z ddl mZGd d d e Z dddZGdddeZdS)) annotations)UnionN)LogisticNormal)StrEnum)DDPMPredictionType) Schedulerc@seZdZdZejZdS)RFlowPredictionTypez Set of valid prediction type names for the RFlow scheduler's `prediction_type` argument. v_prediction: velocity prediction, see section 2.4 https://imagen.research.google/video/paper.pdf N)__name__ __module__ __qualname____doc__r V_PREDICTIONrrj/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/schedulers/rectified_flow.pyr +s r ?c CsD||}||d|}||}||d|d|}||}|S)aQ Applies a transformation to the timestep based on image resolution scaling. Args: t (torch.Tensor): The original timestep(s). input_img_size_numel (torch.Tensor): The input image's size (H * W * D). base_img_size_numel (int): reference H*W*D size, usually smaller than input_img_size_numel. scale (float): Scaling factor for the transformation. num_train_timesteps (int): Total number of training timesteps. spatial_dim (int): Number of spatial dimensions in the image. Returns: torch.Tensor: Transformed timestep(s). rrr) tinput_img_size_numelbase_img_size_numelscalenum_train_timesteps spatial_dimZ ratio_spacerationew_trrrtimestep_transform5s rc@s`eZdZdZ         d5d6ddZd7d d!Z " "d8d9d)d*Zd+d,Z "d:d;d3d4Zd"S)<RFlowSchedulera A rectified flow scheduler for guiding the diffusion process in a generative model. Supports uniform and logit-normal sampling methods, timestep transformation for different resolutions, and noise addition during diffusion. Args: num_train_timesteps (int): Total number of training timesteps. use_discrete_timesteps (bool): Whether to use discrete timesteps. sample_method (str): Training time step sampling method ('uniform' or 'logit-normal'). loc (float): Location parameter for logit-normal distribution, used only if sample_method='logit-normal'. scale (float): Scale parameter for logit-normal distribution, used only if sample_method='logit-normal'. use_timestep_transform (bool): Whether to apply timestep transformation. If true, there will be more inference timesteps at early(noisy) stages for larger image volumes. transform_scale (float): Scaling factor for timestep transformation, used only if use_timestep_transform=True. steps_offset (int): Offset added to computed timesteps, used only if use_timestep_transform=True. base_img_size_numel (int): Reference image volume size for scaling, used only if use_timestep_transform=True. spatial_dim (int): 2 or 3, incidcating 2D or 3D images, used only if use_timestep_transform=True. Example: .. code-block:: python # define a scheduler noise_scheduler = RFlowScheduler( num_train_timesteps = 1000, use_discrete_timesteps = True, sample_method = 'logit-normal', use_timestep_transform = True, base_img_size_numel = 32 * 32 * 32, spatial_dim = 3 ) # during training inputs = torch.ones(2,4,64,64,32) noise = torch.randn_like(inputs) timesteps = noise_scheduler.sample_timesteps(inputs) noisy_inputs = noise_scheduler.add_noise(original_samples=inputs, noise=noise, timesteps=timesteps) predicted_velocity = diffusion_unet( x=noisy_inputs, timesteps=timesteps ) loss = loss_l1(predicted_velocity, (inputs - noise)) # during inference noisy_inputs = torch.randn(2,4,64,64,32) input_img_size_numel = torch.prod(torch.tensor(noisy_inputs.shape[-3:]) noise_scheduler.set_timesteps( num_inference_steps=30, input_img_size_numel=input_img_size_numel) ) all_next_timesteps = torch.cat( (noise_scheduler.timesteps[1:], torch.tensor([0], dtype=noise_scheduler.timesteps.dtype)) ) for t, next_t in tqdm( zip(noise_scheduler.timesteps, all_next_timesteps), total=min(len(noise_scheduler.timesteps), len(all_next_timesteps)), ): predicted_velocity = diffusion_unet( x=noisy_inputs, timesteps=timesteps ) noisy_inputs, _ = noise_scheduler.step(predicted_velocity, t, noisy_inputs, next_t) final_output = noisy_inputs rTuniformrFrrrrintuse_discrete_timestepsbool sample_methodstrlocfloatruse_timestep_transformtransform_scale steps_offsetrrc stj_|_|_| _| _|dvrtd|d|_|dkr8t t |gt |g_ fdd_ |_|_|_dS)N)r logit-normalzsample_method = z:, which has to be chosen from ['uniform', 'logit-normal'].r+cs(j|jdfdddf|jS)Nr) distributionsampleshapetodevice)xselfrrs(z)RFlowScheduler.__init__..)r rprediction_typerr"rr ValueErrorr$rtorchtensorr,sample_tr(r)r*) r3rr"r$r&rr(r)r*rrrr2r__init__s   zRFlowScheduler.__init__original_samples torch.Tensornoise timestepsreturncCs||j}d|}|jdkr!|djdg|jddR}n|jdkr7|djdg|jddR}ntd|j||d||}|S) aV Add noise to the original samples. Args: original_samples: original samples noise: noise to add to samples timesteps: timesteps tensor with shape of (N,), indicating the timestep to be computed for each sample. Returns: noisy_samples: sample with added noise r).NNNNN).NNNz9noise tensor has to be 4D or 5D tensor, yet got shape of )r'rndimexpandr.r6)r3r;r=r>Z timepoints noisy_samplesrrr add_noises  " "zRFlowScheduler.add_noiseNnum_inference_stepsr0str | torch.device | Noner int | NoneNonecs|jks |dkrtd|djdjd|_fddtjD}jr2dd|D}jr?fd d|D}t|tj }jrQ|tj }t | |_jj7_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. input_img_size_numel: int, H*W*D of the image, used with self.use_timestep_transform is True. rz`num_inference_steps`: zM should be at least 1, and cannot be larger than `self.num_train_timesteps`: zG as the unet model trained with this scheduler can only handle maximal z timesteps.cs g|] }d|jjqS)r)rGr).0ir2rr sz0RFlowScheduler.set_timesteps..cSsg|]}tt|qSr)r!roundrKrrrrrMsc s$g|]}t|jjjdqS)rrrr)rrrrrOrr3rrrMsN)rr6rGranger"r(nparrayastypefloat16int64r7 from_numpyr/r>r*)r3rGr0rr>Z timesteps_nprrQr set_timestepss.    zRFlowScheduler.set_timestepscCs|jdkrtj|jdf|jd|j}n |jdkr"|||j}|jr)|}|j rIt t |jdd}t |||j |jt|jdd}|S)z Randomly samples training timesteps using the chosen sampling method. Args: x_start (torch.Tensor): The input tensor for sampling. Returns: torch.Tensor: Sampled timesteps. rr)r0r+NrP)r$r7randr.r0rr9r"longr(prodr8rrlen)r3Zx_startrrrrrsample_timestepss   zRFlowScheduler.sample_timesteps model_outputtimestepr- next_timestepUnion[int, None]!tuple[torch.Tensor, torch.Tensor]c Cst|dr t|jtstd|}|dur#t|}t|||j}n|jdkr/dt|jnd}|||}||||j}||fS)a Predicts the next sample in the diffusion process. Args: model_output (torch.Tensor): Output from the trained diffusion model. timestep (int): Current timestep in the diffusion chain. sample (torch.Tensor): Current sample in the process. next_timestep (Union[int, None]): Optional next timestep. Returns: tuple[torch.Tensor, torch.Tensor]: Predicted sample at the next step and additional info. rGznum_inference_steps is missing or not an integer in the class.Please run self.set_timesteps(num_inference_steps,device,input_img_size_numel) to set it.Nrrr )hasattr isinstancerGr!AttributeErrorr'r) r3r`rar-rbZv_preddtpred_post_samplepred_original_samplerrrsteps zRFlowScheduler.step) rTrr rFrrrr)rr!r"r#r$r%r&r'rr'r(r#r)r'r*r!rr!rr!)r;r<r=r<r>r<r?r<)NN)rGr!r0rHrrIr?rJ)N) r`r<rar!r-r<rbrcr?rd) r r r r r:rFrYr_rkrrrrrPs(C $ .r)rrrr) __future__rtypingrnumpyrSr7Ztorch.distributionsr monai.utilsrddpmr schedulerrr rrrrrrs