update to the more efficient frac-connections

Author: lucidrains

README.md

@@ -265,3 +265,14 @@ $ pip install -U diffusers transformers accelerate scipy ftfy safetensors
     url     = {https://api.semanticscholar.org/CorpusID:272987528}
 }
 ```
+
+```bibtex
+@article{Zhu2025FracConnectionsFE,
+    title   = {Frac-Connections: Fractional Extension of Hyper-Connections},
+    author  = {Defa Zhu and Hongzhi Huang and Jundong Zhou and Zihao Huang and Yutao Zeng and Banggu Wu and Qiyang Min and Xun Zhou},
+    journal = {ArXiv},
+    year    = {2025},
+    volume  = {abs/2503.14125},
+    url     = {https://api.semanticscholar.org/CorpusID:277104144}
+}
+```

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "transfusion-pytorch"
-version = "0.10.5"
+version = "0.11.0"
 description = "Transfusion in Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

transfusion_pytorch/transfusion.py

@@ -1007,7 +1007,8 @@ def __init__(
         attn_laser = False,
         unet_skips = True,
         use_flex_attn = False,
-        num_residual_streams = 4
+        num_residual_streams = 1,
+        num_residual_fracs = 4
     ):
         super().__init__()
         self.use_flex_attn = use_flex_attn
@@ -1023,13 +1024,9 @@ def __init__(
 
         # hyper connections
 
-        assert num_residual_streams > 0
-        is_hyper_connection = num_residual_streams > 1
-        self.num_residual_streams = num_residual_streams
-
         counter = count()
 
-        init_residual_fn, self.expand_stream, self.reduce_stream = HyperConnections.get_init_and_expand_reduce_stream_functions(num_residual_streams, disable = num_residual_streams == 1)
+        init_residual_fn, self.expand_stream, self.reduce_stream = HyperConnections.get_init_and_expand_reduce_stream_functions(num_residual_streams, num_fracs = num_residual_fracs)
 
         # layers