Use stochastic interpolants class (atong01#101)

* use SI class within tutorials and cifar training

* add credit to SI and Rectified flow Flow

Author: kilianFatras

examples/2D_tutorials/Flow_matching_tutorial.ipynb

@@ -5,7 +5,7 @@
    "id": "fb2b2856",
    "metadata": {},
    "source": [
-    "# Flow Matching tutorial: making ODE generative models as good as Diffusion models"
+    "# Flow Matching tutorial: training ODE generative models like Diffusion models"
    ]
   },
   {
@@ -15,11 +15,13 @@
    "source": [
     "## Introduction\n",
     "\n",
-    "Flow Matching, a recently introduced generative model, leverages an ordinary differential equation (ODE) to mold a basic density into the desired data distribution. In contrast, diffusion is based on a stochastic differential equation (SDE). This notebook illustrates the training of Flow Matching methods, highlighting key components. We introduce two models: Independent Conditional Flow Matching (I-CFM) and Optimal Transport Conditional Flow Matching (OT-CFM).\n",
+    "**Flow Matching was introduced in three different ICLR 2023 papers and has drawn a lot of attention in the machine learning community recently. We would like to highlight all of them here: Flow Matching [(Lipman et al.)](https://arxiv.org/abs/2210.02747), Stochastic Interpolants [(Albergo et al.)](https://arxiv.org/abs/2209.15571) and Rectified Flow [(Liu et al.)](https://arxiv.org/abs/2209.03003).**\n",
+    "\n",
+    "Flow Matching, a recently introduced generative model, leverages an ordinary differential equation (ODE) to mold a base density into the desired data distribution. In contrast, diffusion is based on a stochastic differential equation (SDE). This notebook illustrates the training of Flow Matching methods, highlighting key components. In this notebook, we present two Flow Matching models built upon the original formulation: Independent Conditional Flow Matching (I-CFM) and Optimal Transport Conditional Flow Matching (OT-CFM).\n",
     "\n",
     "In our notation, $\\alpha$ represents the noise distribution, typically a Gaussian, while $\\beta$ denotes the distribution corresponding to real data.\n",
     "\n",
-    "Note from the authors: this is a beta turotial! Do not hesitate to suggest improvements through the opened issue https://github.com/atong01/conditional-flow-matching/issues/88"
+    "Note from the authors: this is a beta tutorial! Do not hesitate to suggest improvements through the opened issue https://github.com/atong01/conditional-flow-matching/issues/88"
    ]
   },
   {

examples/2D_tutorials/tutorial_training_8_gaussians_to_moons.ipynb

@@ -47,7 +47,7 @@
 
 
 # Load the model
-PATH = f"{FLAGS.input_dir}/{FLAGS.model}/cifar10_weights_step_{FLAGS.step}.pt"
+PATH = f"{FLAGS.input_dir}/{FLAGS.model}/{FLAGS.model}_cifar10_weights_step_{FLAGS.step}.pt"
 print("path: ", PATH)
 checkpoint = torch.load(PATH)
 state_dict = checkpoint["ema_model"]

examples/images/cifar10/compute_fid.py

@@ -17,6 +17,7 @@
     ConditionalFlowMatcher,
     ExactOptimalTransportConditionalFlowMatcher,
     TargetConditionalFlowMatcher,
+    VariancePreservingConditionalFlowMatcher,
 )
 from torchcfm.models.unet.unet import UNetModelWrapper
 
@@ -128,9 +129,11 @@ def train(argv):
         FM = ConditionalFlowMatcher(sigma=sigma)
     elif FLAGS.model == "fm":
         FM = TargetConditionalFlowMatcher(sigma=sigma)
+    elif FLAGS.model == "si":
+        FM = VariancePreservingConditionalFlowMatcher(sigma=sigma)
     else:
         raise NotImplementedError(
-            f"Unknown model {FLAGS.model}, must be one of ['otcfm', 'icfm', 'fm']"
+            f"Unknown model {FLAGS.model}, must be one of ['otcfm', 'icfm', 'fm', 'si']"
         )
 
     savedir = FLAGS.output_dir + FLAGS.model + "/"
@@ -162,7 +165,7 @@ def train(argv):
                         "optim": optim.state_dict(),
                         "step": step,
                     },
-                    savedir + f"cifar10_weights_step_{step}.pt",
+                    savedir + f"{FLAGS.model}_cifar10_weights_step_{step}.pt",
                 )