Add multimodal flow matching example

README.md

@@ -69,7 +69,7 @@ pip install -e .
 
 #### I want to train a Flow Matching model, where can I find the training code?
 
-We provide [training examples](examples). Under this folder, you can find synthetic data for [continuous](examples/2d_flow_matching.ipynb), [discrete](examples/2d_discrete_flow_matching.ipynb), and [Riemannian](examples/2d_riemannian_flow_matching_flat_torus.ipynb) Flow Matching. We also provide full training [examples](examples/image) (continuous and discrete) on CIFAR10 and face-blurred ImageNet, and a scalable discrete Flow Matching example for [text modeling](examples/text).
+We provide [training examples](examples). Under this folder, you can find synthetic data for [continuous](examples/2d_flow_matching.ipynb), [discrete](examples/2d_discrete_flow_matching.ipynb), [multimodal](examples/2d_multimodal_flow_matching.ipynb), and [Riemannian](examples/2d_riemannian_flow_matching_flat_torus.ipynb) Flow Matching. We also provide full training [examples](examples/image) (continuous and discrete) on CIFAR10 and face-blurred ImageNet, and a scalable discrete Flow Matching example for [text modeling](examples/text).
 
 #### Do you release pre-trained models?
 

docs/Makefile

@@ -13,6 +13,7 @@ ROOT_DIR:=$(shell dirname $(realpath $(firstword $(MAKEFILE_LIST))))
 links:
 	mkdir -p source/notebooks && ln -sfn $(ROOT_DIR)/../examples/standalone_flow_matching.ipynb source/notebooks/standalone_flow_matching.ipynb
 	mkdir -p source/notebooks && ln -sfn $(ROOT_DIR)/../examples/2d_discrete_flow_matching.ipynb source/notebooks/2d_discrete_flow_matching.ipynb
+	mkdir -p source/notebooks && ln -sfn $(ROOT_DIR)/../examples/2d_multimodal_flow_matching.ipynb source/notebooks/2d_multimodal_flow_matching.ipynb
 	mkdir -p source/notebooks && ln -sfn $(ROOT_DIR)/../examples/2d_riemannian_flow_matching_flat_torus.ipynb source/notebooks/2d_riemannian_flow_matching_flat_torus.ipynb
 	mkdir -p source/notebooks && ln -sfn $(ROOT_DIR)/../examples/2d_riemannian_flow_matching_sphere.ipynb source/notebooks/2d_riemannian_flow_matching_sphere.ipynb
 	ln -sfn $(ROOT_DIR)/../assets/teaser.png  source/_images/teaser.png

docs/source/dummy.rst

@@ -5,5 +5,6 @@
 
    notebooks/standalone_flow_matching
    notebooks/2d_discrete_flow_matching
+   notebooks/2d_multimodal_flow_matching
    notebooks/2d_riemannian_flow_matching_flat_torus
    notebooks/2d_riemannian_flow_matching_sphere

docs/source/flow_matching.solver.rst

@@ -14,5 +14,6 @@ Solvers
    Solver
    ODESolver
    MixtureDiscreteEulerSolver
+   MultimodalSolver
    RiemannianODESolver
 

docs/source/flow_matching.utils.multimodal.rst

@@ -0,0 +1,18 @@
+``flow_matching.utils.multimodal``
+=============================
+
+.. currentmodule:: flow_matching.utils.multimodal
+
+
+Flow
+--------------------------------
+
+Generic multimodal flow class
+
+.. autosummary::
+   :toctree: generated
+   :nosignatures:
+   :template: classtemplate.rst
+
+   Flow
+

docs/source/modules.rst

@@ -10,3 +10,4 @@ API Reference
    flow_matching.solver
    flow_matching.utils.model_wrapper
    flow_matching.utils.manifolds
+   flow_matching.utils.multimodal

docs/source/notebooks.rst

@@ -29,4 +29,10 @@ Notebooks
    :image: _static/riemannian_torus.png
    :link: notebooks/2d_riemannian_flow_matching_flat_torus.html
 
+.. customcarditem::
+   :header: Multimodal Flow Matching
+   :card_description: Train and sample from a 2D Multimodal Flow Matching model.
+   :image: _static/multimodal.png
+   :link: notebooks/2d_multimodal_flow_matching.html
+
 .. customcardend::