# Copyright (c) MONAI Consortium # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # ========================================================================= # Adapted from https://github.com/hpcaitech/Open-Sora/blob/main/opensora/schedulers/rf/rectified_flow.py # which has the following license: # https://github.com/hpcaitech/Open-Sora/blob/main/LICENSE # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ========================================================================= from __future__ import annotations from typing import Union import numpy as np import torch from torch.distributions import LogisticNormal from monai.utils import StrEnum from .ddpm import DDPMPredictionType from .scheduler import Scheduler class RFlowPredictionType(StrEnum): """ 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 """ V_PREDICTION = DDPMPredictionType.V_PREDICTION def timestep_transform( t, input_img_size_numel, base_img_size_numel=32 * 32 * 32, scale=1.0, num_train_timesteps=1000, spatial_dim=3 ): """ 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). """ t = t / num_train_timesteps ratio_space = (input_img_size_numel / base_img_size_numel) ** (1.0 / spatial_dim) ratio = ratio_space * scale new_t = ratio * t / (1 + (ratio - 1) * t) new_t = new_t * num_train_timesteps return new_t class RFlowScheduler(Scheduler): """ 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 """ def __init__( self, num_train_timesteps: int = 1000, use_discrete_timesteps: bool = True, sample_method: str = "uniform", loc: float = 0.0, scale: float = 1.0, use_timestep_transform: bool = False, transform_scale: float = 1.0, steps_offset: int = 0, base_img_size_numel: int = 32 * 32 * 32, spatial_dim: int = 3, ): # rectified flow only accepts velocity prediction self.prediction_type = RFlowPredictionType.V_PREDICTION self.num_train_timesteps = num_train_timesteps self.use_discrete_timesteps = use_discrete_timesteps self.base_img_size_numel = base_img_size_numel self.spatial_dim = spatial_dim # sample method if sample_method not in ["uniform", "logit-normal"]: raise ValueError( f"sample_method = {sample_method}, which has to be chosen from ['uniform', 'logit-normal']." ) self.sample_method = sample_method if sample_method == "logit-normal": self.distribution = LogisticNormal(torch.tensor([loc]), torch.tensor([scale])) self.sample_t = lambda x: self.distribution.sample((x.shape[0],))[:, 0].to(x.device) # timestep transform self.use_timestep_transform = use_timestep_transform self.transform_scale = transform_scale self.steps_offset = steps_offset def add_noise(self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor: """ 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 """ timepoints: torch.Tensor = timesteps.float() / self.num_train_timesteps timepoints = 1 - timepoints # [1,1/1000] # expand timepoint to noise shape if noise.ndim == 5: timepoints = timepoints[..., None, None, None, None].expand(-1, *noise.shape[1:]) elif noise.ndim == 4: timepoints = timepoints[..., None, None, None].expand(-1, *noise.shape[1:]) else: raise ValueError(f"noise tensor has to be 4D or 5D tensor, yet got shape of {noise.shape}") noisy_samples: torch.Tensor = timepoints * original_samples + (1 - timepoints) * noise return noisy_samples def set_timesteps( self, num_inference_steps: int, device: str | torch.device | None = None, input_img_size_numel: int | None = None, ) -> None: """ 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. """ if num_inference_steps > self.num_train_timesteps or num_inference_steps < 1: raise ValueError( f"`num_inference_steps`: {num_inference_steps} should be at least 1, " "and cannot be larger than `self.num_train_timesteps`:" f" {self.num_train_timesteps} as the unet model trained with this scheduler can only handle" f" maximal {self.num_train_timesteps} timesteps." ) self.num_inference_steps = num_inference_steps # prepare timesteps timesteps = [ (1.0 - i / self.num_inference_steps) * self.num_train_timesteps for i in range(self.num_inference_steps) ] if self.use_discrete_timesteps: timesteps = [int(round(t)) for t in timesteps] if self.use_timestep_transform: timesteps = [ timestep_transform( t, input_img_size_numel=input_img_size_numel, base_img_size_numel=self.base_img_size_numel, num_train_timesteps=self.num_train_timesteps, spatial_dim=self.spatial_dim, ) for t in timesteps ] timesteps_np = np.array(timesteps).astype(np.float16) if self.use_discrete_timesteps: timesteps_np = timesteps_np.astype(np.int64) self.timesteps = torch.from_numpy(timesteps_np).to(device) self.timesteps += self.steps_offset def sample_timesteps(self, x_start): """ Randomly samples training timesteps using the chosen sampling method. Args: x_start (torch.Tensor): The input tensor for sampling. Returns: torch.Tensor: Sampled timesteps. """ if self.sample_method == "uniform": t = torch.rand((x_start.shape[0],), device=x_start.device) * self.num_train_timesteps elif self.sample_method == "logit-normal": t = self.sample_t(x_start) * self.num_train_timesteps if self.use_discrete_timesteps: t = t.long() if self.use_timestep_transform: input_img_size_numel = torch.prod(torch.tensor(x_start.shape[2:])) t = timestep_transform( t, input_img_size_numel=input_img_size_numel, base_img_size_numel=self.base_img_size_numel, num_train_timesteps=self.num_train_timesteps, spatial_dim=len(x_start.shape) - 2, ) return t def step( self, model_output: torch.Tensor, timestep: int, sample: torch.Tensor, next_timestep: Union[int, None] = None ) -> tuple[torch.Tensor, torch.Tensor]: """ 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. """ # Ensure num_inference_steps exists and is a valid integer if not hasattr(self, "num_inference_steps") or not isinstance(self.num_inference_steps, int): raise AttributeError( "num_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." ) v_pred = model_output if next_timestep is not None: next_timestep = int(next_timestep) dt: float = ( float(timestep - next_timestep) / self.num_train_timesteps ) # Now next_timestep is guaranteed to be int else: dt = ( 1.0 / float(self.num_inference_steps) if self.num_inference_steps > 0 else 0.0 ) # Avoid division by zero pred_post_sample = sample + v_pred * dt pred_original_sample = sample + v_pred * timestep / self.num_train_timesteps return pred_post_sample, pred_original_sample