able to ues the consistency EMA as the main EMA being sampled from the Trainer. remove wip as things are working

Author: lucidrains

README.md

@@ -1,6 +1,6 @@
 <img src="./rf.png" width="400px"></img>
 
-## Rectified Flow - Pytorch (wip)
+## Rectified Flow - Pytorch
 
 Implementation of <a href="https://www.cs.utexas.edu/~lqiang/rectflow/html/intro.html">rectified flow</a> and some of its followup research / improvements in Pytorch
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rectified-flow-pytorch"
-version = "0.0.22"
+version = "0.0.23"
 description = "Rectified Flow in Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

rectified_flow_pytorch/rectified_flow.py

@@ -824,19 +824,25 @@ def __init__(
         num_samples: int = 16,
         adam_kwargs: dict = dict(),
         accelerate_kwargs: dict = dict(),
-        ema_kwargs: dict = dict()
+        ema_kwargs: dict = dict(),
+        use_consistency_ema = False  # whether to just use the EMA from the velocity consistency from the consistency FM paper
     ):
         super().__init__()
         self.accelerator = Accelerator(**accelerate_kwargs)
 
         self.model = rectified_flow
 
         if self.is_main:
-            self.ema_model = EMA(
-                self.model,
-                forward_method_names = ('sample',),
-                **ema_kwargs
-            )
+            if use_consistency_ema:
+                assert self.model.use_consistency, 'model must be using the consistency EMA for it to be reused as the main EMA model during sampling'
+
+                self.ema_model = self.model.ema_model
+            else:
+                self.ema_model = EMA(
+                    self.model,
+                    forward_method_names = ('sample',),
+                    **ema_kwargs
+                )
 
             self.ema_model.to(self.accelerator.device)