Salesforce/FOFPred

1from dataclasses import dataclass2from typing import List, Optional, Tuple, Union3 4import numpy as np5import torch6from diffusers.configuration_utils import ConfigMixin, register_to_config7from diffusers.loaders.lora_base import (  # noqa8    LoraBaseMixin,9    _fetch_state_dict,10)11from diffusers.schedulers.scheduling_utils import SchedulerMixin12from diffusers.utils import BaseOutput13 14 15@dataclass16class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput):17    """18    Output class for the scheduler's `step` function output.19 20    Args:21        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):22            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the23            denoising loop.24    """25 26    prev_sample: torch.FloatTensor27 28 29class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):30    """31    Euler scheduler.32 33    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic34    methods the library implements for all schedulers such as loading and saving.35 36    Args:37        num_train_timesteps (`int`, defaults to 1000):38            The number of diffusion steps to train the model.39        timestep_spacing (`str`, defaults to `"linspace"`):40            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and41            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.42        shift (`float`, defaults to 1.0):43            The shift value for the timestep schedule.44    """45 46    _compatibles = []47    order = 148 49    @register_to_config50    def __init__(51        self, num_train_timesteps: int = 1000, dynamic_time_shift: bool = True52    ):53        timesteps = torch.linspace(0, 1, num_train_timesteps + 1, dtype=torch.float32)[54            :-155        ]56 57        self.timesteps = timesteps58 59        self._step_index = None60        self._begin_index = None61 62    @property63    def step_index(self):64        """65        The index counter for current timestep. It will increase 1 after each scheduler step.66        """67        return self._step_index68 69    @property70    def begin_index(self):71        """72        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.73        """74        return self._begin_index75 76    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index77    def set_begin_index(self, begin_index: int = 0):78        """79        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.80 81        Args:82            begin_index (`int`):83                The begin index for the scheduler.84        """85        self._begin_index = begin_index86 87    def index_for_timestep(self, timestep, schedule_timesteps=None):88        if schedule_timesteps is None:89            schedule_timesteps = self._timesteps90 91        indices = (schedule_timesteps == timestep).nonzero()92 93        # The sigma index that is taken for the **very** first `step`94        # is always the second index (or the last index if there is only 1)95        # This way we can ensure we don't accidentally skip a sigma in96        # case we start in the middle of the denoising schedule (e.g. for image-to-image)97        pos = 1 if len(indices) > 1 else 098 99        return indices[pos].item()100 101    # def time_shift(self, mu: float, sigma: float, t: torch.Tensor):102    #     return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)103 104    def set_timesteps(105        self,106        num_inference_steps: int = None,107        device: Union[str, torch.device] = None,108        timesteps: Optional[List[float]] = None,109        num_tokens: Optional[int] = None,110    ):111        """112        Sets the discrete timesteps used for the diffusion chain (to be run before inference).113 114        Args:115            num_inference_steps (`int`):116                The number of diffusion steps used when generating samples with a pre-trained model.117            device (`str` or `torch.device`, *optional*):118                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.119        """120 121        if timesteps is None:122            self.num_inference_steps = num_inference_steps123            timesteps = np.linspace(0, 1, num_inference_steps + 1, dtype=np.float32)[124                :-1125            ]126            if self.config.dynamic_time_shift and num_tokens is not None:127                m = (128                    np.sqrt(num_tokens) / 40129                )  # when input resolution is 320 * 320, m = 1, when input resolution is 1024 * 1024, m = 3.2130                timesteps = timesteps / (m - m * timesteps + timesteps)131 132        timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32, device=device)133        _timesteps = torch.cat([timesteps, torch.ones(1, device=timesteps.device)])134 135        self.timesteps = timesteps136        self._timesteps = _timesteps137        self._step_index = None138        self._begin_index = None139 140    def _init_step_index(self, timestep):141        if self.begin_index is None:142            if isinstance(timestep, torch.Tensor):143                timestep = timestep.to(self.timesteps.device)144            self._step_index = self.index_for_timestep(timestep)145        else:146            self._step_index = self._begin_index147 148    def step(149        self,150        model_output: torch.FloatTensor,151        timestep: Union[float, torch.FloatTensor],152        sample: torch.FloatTensor,153        generator: Optional[torch.Generator] = None,154        return_dict: bool = True,155    ) -> Union[FlowMatchEulerDiscreteSchedulerOutput, Tuple]:156        """157        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion158        process from the learned model outputs (most often the predicted noise).159 160        Args:161            model_output (`torch.FloatTensor`):162                The direct output from learned diffusion model.163            timestep (`float`):164                The current discrete timestep in the diffusion chain.165            sample (`torch.FloatTensor`):166                A current instance of a sample created by the diffusion process.167            s_churn (`float`):168            s_tmin  (`float`):169            s_tmax  (`float`):170            s_noise (`float`, defaults to 1.0):171                Scaling factor for noise added to the sample.172            generator (`torch.Generator`, *optional*):173                A random number generator.174            return_dict (`bool`):175                Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or176                tuple.177 178        Returns:179            [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:180                If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is181                returned, otherwise a tuple is returned where the first element is the sample tensor.182        """183 184        if (185            isinstance(timestep, int)186            or isinstance(timestep, torch.IntTensor)187            or isinstance(timestep, torch.LongTensor)188        ):189            raise ValueError(190                (191                    "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"192                    " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"193                    " one of the `scheduler.timesteps` as a timestep."194                ),195            )196 197        if self.step_index is None:198            self._init_step_index(timestep)199        # Upcast to avoid precision issues when computing prev_sample200        sample = sample.to(torch.float32)201        t = self._timesteps[self.step_index]202        t_next = self._timesteps[self.step_index + 1]203 204        prev_sample = sample + (t_next - t) * model_output205 206        # Cast sample back to model compatible dtype207        prev_sample = prev_sample.to(model_output.dtype)208 209        # upon completion increase step index by one210        self._step_index += 1211 212        if not return_dict:213            return (prev_sample,)214 215        return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)216 217    def __len__(self):218        return self.config.num_train_timesteps219