make sure when doing velocity consistency that base model only receives times from 0. to 1. - delta_time

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "transfusion-pytorch"
-version = "0.1.11"
+version = "0.1.12"
 description = "Transfusion in Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

transfusion_pytorch/transfusion.py

@@ -1647,6 +1647,13 @@ def forward(
 
         times_per_token = einsum(is_modalities.float(), times, 'b t m n, b m -> b t n')
 
+        # if needs velocity matching, make sure times are in the range of 0 - (1. - <velocity consistency delta time>)
+
+        if need_velocity_matching:
+            times_per_token = times_per_token * (1. - velocity_consistency_delta_time)
+
+        # noise only if returning loss
+
         if return_loss:
             noised_modality_tokens = []
             flows = []