adding layers and math test

ml-gde · Oct 7, 2024 · a1c72e2 · a1c72e2
1 parent 44431ec
commit a1c72e2
Show file tree

Hide file tree

Showing 5 changed files with 545 additions and 240 deletions.
diff --git a/jflux/cli.py b/jflux/cli.py
@@ -10,13 +10,7 @@
 from jax.typing import DTypeLike
 
 from jflux.sampling import denoise, get_noise, get_schedule, prepare, unpack
-from jflux.util import (
-    configs,
-    load_ae,
-    load_clip,
-    load_flow_model,
-    load_t5,
-)
+from jflux.util import configs, load_ae, load_clip, load_flow_model, load_t5
 
 
 @dataclass

diff --git a/jflux/model.py b/jflux/model.py
@@ -7,12 +7,7 @@
 from jax.typing import DTypeLike
 
 from jflux.modules import DoubleStreamBlock, MLPEmbedder, SingleStreamBlock
-from jflux.modules.layers import (
-    AdaLayerNorm,
-    Embed,
-    Identity,
-    timestep_embedding,
-)
+from jflux.modules.layers import AdaLayerNorm, Embed, Identity, timestep_embedding
 
 
 @dataclass