Adding FLUX porting code (#11)

Co-authored-by: Saurav Maheshkar <[email protected]>

.github/CODEOWNERS

@@ -1 +1 @@
-* @SauravMaheshkar
+* @SauravMaheshkar @ariG23498

jflux/cli.py

@@ -6,13 +6,12 @@
 
 import jax
 import jax.numpy as jnp
-from flax import nnx
+from einops import rearrange
 from fire import Fire
+from flax import nnx
 from jax.typing import DTypeLike
-
 from PIL import Image
 
-from einops import rearrange
 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
 

jflux/modules/conditioner.py

@@ -1,7 +1,7 @@
 # Note: This is a torch module not a Jax module
-from torch import nn
-from chex import Array
 import jax.numpy as jnp
+from chex import Array
+from torch import nn
 from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer
 
 

jflux/modules/layers.py

@@ -214,6 +214,7 @@ def __init__(
         self.img_norm1 = nnx.LayerNorm(
             num_features=hidden_size,
             use_scale=False,
+            use_bias=False,
             epsilon=1e-6,
             rngs=rngs,
             param_dtype=param_dtype,
@@ -229,6 +230,7 @@ def __init__(
         self.img_norm2 = nnx.LayerNorm(
             num_features=hidden_size,
             use_scale=False,
+            use_bias=False,
             epsilon=1e-6,
             rngs=rngs,
             param_dtype=param_dtype,
@@ -257,6 +259,7 @@ def __init__(
         self.txt_norm1 = nnx.LayerNorm(
             num_features=hidden_size,
             use_scale=False,
+            use_bias=False,
             epsilon=1e-6,
             rngs=rngs,
             param_dtype=param_dtype,
@@ -272,6 +275,7 @@ def __init__(
         self.txt_norm2 = nnx.LayerNorm(
             num_features=hidden_size,
             use_scale=False,
+            use_bias=False,
             epsilon=1e-6,
             rngs=rngs,
             param_dtype=param_dtype,
@@ -382,6 +386,7 @@ def __init__(
         self.pre_norm = nnx.LayerNorm(
             num_features=hidden_size,
             use_scale=False,
+            use_bias=False,
             epsilon=1e-6,
             rngs=rngs,
             param_dtype=param_dtype,
@@ -419,6 +424,7 @@ def __init__(
         self.norm_final = nnx.LayerNorm(
             num_features=hidden_size,
             use_scale=False,
+            use_bias=False,
             epsilon=1e-6,
             rngs=rngs,
             param_dtype=param_dtype,