diff --git a/README.md b/README.md index bcc73e6..89fb968 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,100 @@ -# comfy-diff2flow -Unofficial Implementation of diff2flow method for ComfyUI. +# ComfyUI-Diff2Flow + +[![arXiv](https://img.shields.io/badge/arXiv-2506.02221-b31b1b.svg)](https://arxiv.org/abs/2506.02221) + +This is an unofficial implementation of the **Diff2Flow** method for ComfyUI, based on the paper: + +**[Diff2Flow: Training Flow Matching Models via Diffusion Model Alignment (arXiv:2506.02221)](https://arxiv.org/abs/2506.02221)**. + +This repository provides a custom KSampler node, `Diff2FlowODESampler`, that applies the Diff2Flow principles to standard diffusion models at inference time. + +## What is Diff2Flow? + +**Diff2Flow** is a framework designed to bridge the gap between two powerful generative modeling paradigms: + +1. **Diffusion Models (DMs):** Like Stable Diffusion, these are known for their high-quality results but often require many sampling steps and can suffer from issues like the "non-zero terminal SNR" problem (difficulty producing pure black or white). +2. **Flow Matching (FM) Models:** A newer paradigm known for faster inference, improved performance , and simpler, straighter sampling paths. + +The challenge is that these two models are not directly compatible. They use different timestep definitions , paths from noise to data (interpolants) , and training objectives. + +The key idea of the paper is to **efficiently transfer the knowledge** from a pre-trained diffusion model into a flow matching model. It does this by "warping" the diffusion process into a flow matching one by: + +* **Rescaling Timesteps:** Systematically mapping the diffusion model's discrete 0-1000 timesteps to the flow matching 0-1 continuous time. +* **Aligning Interpolants:** Mathematically aligning the path from noise-to-data between the two model types. +* **Deriving Velocity:** Calculating the **velocity field** (which FMs use) directly from the diffusion model's standard prediction (e.g., `epsilon` or `v-prediction`). + +This allows a model to be sampled using a flow-matching objective, which can lead to faster, more efficient inference. + +## How This Node Works + +This implementation provides a single node, `Diff2Flow ODE KSampler`. + +Instead of requiring a special model finetuned with the Diff2Flow objective, this node **patches your existing diffusion model** (e.g., SD1.5, SD2.1, SDXL) at runtime. + +When you sample, the `enable_diff2flow` function: +1. **Patches** your loaded model by adding the `Diff2Flow` mathematical conversions based on the model's original `beta_schedule`. +2. The `Diff2FlowODESampler` then intercepts the sampling call. At each step, it: + * Rescales the FM timestep to the equivalent DM timestep. + * Aligns the FM latent to the equivalent DM latent. + * Feeds these aligned inputs into the original UNet. + * Takes the UNet's output (`epsilon` or `v-prediction`) and converts it into a **velocity field** of Flow-Matching. +3. This derived velocity is then used by a true **Ordinary Differential Equation (ODE) solver** from the `torchdiffeq` library to perform the sampling step. + +## 🚀 Key Features + +* **Works with Standard Checkpoints:** Use your existing `v-prediction` (like SD 2.1) and `epsilon-prediction` (like SDXL) models. +* **True ODE Solvers:** Integrates with `torchdiffeq` to provide a wide range of solvers, including adaptive-step-size solvers (`dopri5`, `bosh3`, etc.) and fixed-step-size solvers (`euler`, `rk4`, `midpoint`, etc.). +* **Single Node:** A simple, drop-in KSampler alternative. + +## 💾 Installation + +1. **Clone the Repository:** + Navigate to your `ComfyUI/custom_nodes/` directory and clone this repo: + ```bash + git clone https://github.com/Koratahiu/comfy-diff2flow + ``` + +2. **Install Dependencies:** + This node has a **critical dependency** on `torchdiffeq`. + + Activate your ComfyUI virtual environment (if you have one) and run: + ```bash + pip install torchdiffeq + ``` + +3. **Restart ComfyUI:** + Completely shut down and restart ComfyUI. + +## 💡 Usage + +1. After installation, find the **`Diff2Flow ODE KSampler`** node in the `sampling/custom_sampling` menu. +2. Connect your `MODEL`, `positive`, `negative`, and `latent_image` just like a standard KSampler. +3. **Select an ODE `solver`**. + * `euler` is the simplest and fastest, equivalent to the standard Euler sampler. + * Adaptive solvers like `dopri5` or `bosh3` may offer different quality/speed trade-offs. You can tune them with the `log_relative_tolerance` and `log_absolute_tolerance` parameters. +4. Adjust `steps`, `cfg`, and `seed`. +5. Generate! + +--- + +> ### ⚠️ Important Note: Inference vs. Finetuning +> +> The original paper achieves its most impressive results (e.g., high-quality 2-step generation ) by **finetuning** a model *using* the Diff2Flow objective. +> +> This node applies the Diff2Flow math for **inference only**. Your results will vary! As you need to finetune the model on diff2flow for replicating the results of a fully finetuned diff2flow model in the paper. +> +> That said, the paper also shows that simply converting a diffusion model to its flow-matching counterpart for the *same* task can lead to performance improvements. +> +> You can finetune the model in diff2flow, either using [the offical code](https://github.com/CompVis/diff2flow) or using this PR of OneTrainer (which is up to now is yet to be merged but it's pretty much stable and tested): +https://github.com/Nerogar/OneTrainer/pull/1052 + +## Known Issues + +- ODE samplers don't support masks well. + +## Acknowledgements + +* Main code adopted from: + https://github.com/CompVis/diff2flow +* All credit for the methodology goes to the paper's authors: Johannes Schusterbauer, Ming Gui, Frank Fundel, and Björn Ommer! +* Also part of the credit for ODE solvers goes to this node that inspired us: https://github.com/redhottensors/ComfyUI-ODE