diffusers icon indicating copy to clipboard operation
diffusers copied to clipboard

Flow Matching Euler Scheduler added

Open leffff opened this issue 1 year ago • 11 comments

In PR I have implemented a scheduler for Flow Matching with Euler method.

The popularity for Flow Matching models grows rapidly. Even novel models like SD3 (https://arxiv.org/pdf/2403.03206) use flow matching instead of Diffusion algorithm.
The FlowMatchingEulerScheduler features a full range of methods. However, unlike regular diffusion, in Flow Matching, timesteps are continous, but not discrete, therefore x_t = t * x_1 + (1 - t) * x_0. To infer Flow-Mathcing models one should perform Euler Method x_t = h * f(x_{t-1}, t) + x_{t-1}.

My implementation and usage examples could be seen here: https://github.com/leffff/euler-scheduler

leffff avatar Jun 09 '24 21:06 leffff

how to use this with StableCascadeCombinedPipeline?

s9anus98a avatar Jun 10 '24 21:06 s9anus98a

Well It is not useable with StableCascadeCombinedPipeline as is. This this scheduler should be used to train new models. With this scheduler diffusion models tend to converge faster, due to a straight trajectory between noise and data.

leffff avatar Jun 10 '24 23:06 leffff

Well It is not useable with StableCascadeCombinedPipeline as is. This this scheduler should be used to train new models. With this scheduler diffusion models tend to converge faster, due to a straight trajectory between noise and data.

Is it possible to use this scheduler for finetune of ready-made models that were trained with a regular ddpm scheduler, or is it only possible to train from scratch?

drimeF0 avatar Jun 16 '24 13:06 drimeF0

Of course! With this scheduler you can "re tune" the model form DDPM prediction to Flow Matching velocity prediction

leffff avatar Jun 19 '24 13:06 leffff

Hi, leffff, I'm confused about this part:

    def add_noise(original_samples: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
        """
        Add noise to the given sample

        Args:
            original_samples (`torch.Tensor`):
                The original sample that is to be noised
            noise (`torch.Tensor`):
                The noise that is used to noise the image
            timestep (`torch.Tensor`):
                Timestep used to create linear interpolation `x_t = t * x_1 + (1 - t) * x_0`.
                Where x_1 is a target distribution, x_0 is a source distribution and t (timestep) ∈ [0, 1]
        """

        t = get_time_coefficients(timestep, original_samples.ndim)

        noised_sample = t * original_samples + (1 - t) * noise

        return noised_sample

In SD3 paper, $z_t$ is caculated as follows: image and velocity should be $\epsilon-x_0$, but in your code $v=x_0-\epsilon$

    def get_velocity(original_samples: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
        """
        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
        process from the learned model outputs (most often the predicted noise).

        Args:
            original_samples (`torch.Tensor`):
                The original sample that is to be noised
            noise (`torch.Tensor`):
                The noise that is used to noise the image

        Returns:
            `torch.Tensor`
        """

        return original_samples - noise

does it matter?

xxheyu avatar Jun 21 '24 08:06 xxheyu

The thing is its the same $x_t = t * x_1 + (1 - t) * x_0$ with $target = x_1 - x_0$ and $x_t = (1 - t) * x_1 + t * x_0$ with $target = x_0 - x_1$ For proof https://github.com/leffff/InstructFlow look at the generations in the readme file)

leffff avatar Jun 21 '24 09:06 leffff

The thing is its the same xt=t∗x1+(1−t)∗x0 with target=x1−x0 and xt=(1−t)∗x1+t∗x0 with target=x0−x1 For proof https://github.com/leffff/InstructFlow look at the generations in the readme file)

$x_t=t∗x_1+(1−t)∗x_0$ and $x_t=(1−t)∗x_1+t∗x_0$ are equivalent since $t$ is uniformly sampled?

xxheyu avatar Jun 21 '24 10:06 xxheyu

If you respectively change the target it should work well

leffff avatar Jun 21 '24 11:06 leffff

Apologies from our end that we haven't been able to get back to you sooner. I just requested for a review from our core maintainer @yiyixuxu. Expect feedback from Yiyi soon.

Additionally, out of curiosity, does this work off the shelf with SD3?

sayakpaul avatar Jun 24 '24 08:06 sayakpaul

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

we added a flow-match Euler scheduler for sd3 (with a shift value that's specific to sd3 but that can be set to 0)

yiyixuxu avatar Jun 29 '24 19:06 yiyixuxu

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

github-actions[bot] avatar Sep 14 '24 15:09 github-actions[bot]

I'm closing this now since we already have a flow matching scheduler :)

yiyixuxu avatar Dec 03 '24 04:12 yiyixuxu