tiny fixes to normalization

Author: amepas

src/maxdiffusion/configs/base_flux2klein.yml

@@ -33,6 +33,7 @@ t5xxl_model_name_or_path: 'ariG23498/t5-v1-1-xxl-flax'
 
 # Flux params
 flux_name: "flux2klein"
+scale_shift_order: "scale_shift"
 use_latents: False
 max_sequence_length: 512
 time_shift: True

src/maxdiffusion/configs/base_flux2klein_9B.yml

@@ -33,6 +33,7 @@ t5xxl_model_name_or_path: 'ariG23498/t5-v1-1-xxl-flax'
 
 # Flux params
 flux_name: "flux2klein_9B"
+scale_shift_order: "scale_shift"
 use_latents: False
 max_sequence_length: 512
 time_shift: True

src/maxdiffusion/generate_flux2klein.py

@@ -130,6 +130,10 @@ def load_or_generate_latents(config):
     """
     Loads saved latents if use_latents is True, otherwise generates random latents.
     """
+    if isinstance(config, dict):
+        from types import SimpleNamespace
+        config = SimpleNamespace(**config)
+        
     batch_size = config.batch_size
     height = config.height
     width = config.width
@@ -186,10 +190,21 @@ def get_qwen3_models(repo_id="black-forest-labs/FLUX.2-klein-4B"):
     """
     global _tokenizer, _text_encoder
     if _tokenizer is None:
+        import os
         import torch
         from transformers import Qwen2TokenizerFast, Qwen3ForCausalLM
 
-        print(f"Loading Qwen3 models from cached repo: {repo_id}...")
+        # Resolve absolute local path from HF cache if it exists to bypass buggy from_pretrained resolving
+        hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
+        cache_dir = os.path.join(hf_home, "hub", f"models--{repo_id.replace('/', '--')}", "snapshots")
+        if os.path.exists(cache_dir):
+            snapshots = os.listdir(cache_dir)
+            if snapshots:
+                snapshot_dir = os.path.join(cache_dir, snapshots[0])
+                print(f"Detected local cache directory: {snapshot_dir}")
+                repo_id = snapshot_dir
+
+        print(f"Loading Qwen3 models from repo path: {repo_id}...")
         try:
             _tokenizer = Qwen2TokenizerFast.from_pretrained(repo_id, local_files_only=True)
         except Exception:
@@ -657,6 +672,7 @@ def main(argv):
         dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32,
         weights_dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32,
         attention_kernel=config.attention,
+        scale_shift_order=getattr(config, "scale_shift_order", "shift_scale"),
     )
 
     # 5b. Instantiate JAX FlaxAutoencoderKL VAE

src/maxdiffusion/generate_flux2klein_9B.py

@@ -137,6 +137,10 @@ def load_or_generate_latents(config):
     """
     Loads saved latents if use_latents is True, otherwise generates random latents.
     """
+    if isinstance(config, dict):
+        from types import SimpleNamespace
+        config = SimpleNamespace(**config)
+        
     batch_size = config.batch_size
     height = config.height
     width = config.width
@@ -193,10 +197,21 @@ def get_qwen3_models(repo_id="black-forest-labs/FLUX.2-klein-4B"):
     """
     global _tokenizer, _text_encoder
     if _tokenizer is None:
+        import os
         import torch
         from transformers import Qwen2TokenizerFast, Qwen3ForCausalLM
 
-        print(f"Loading Qwen3 models from cached repo: {repo_id}...")
+        # Resolve absolute local path from HF cache if it exists to bypass buggy from_pretrained resolving
+        hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
+        cache_dir = os.path.join(hf_home, "hub", f"models--{repo_id.replace('/', '--')}", "snapshots")
+        if os.path.exists(cache_dir):
+            snapshots = os.listdir(cache_dir)
+            if snapshots:
+                snapshot_dir = os.path.join(cache_dir, snapshots[0])
+                print(f"Detected local cache directory: {snapshot_dir}")
+                repo_id = snapshot_dir
+
+        print(f"Loading Qwen3 models from repo path: {repo_id}...")
         try:
             _tokenizer = Qwen2TokenizerFast.from_pretrained(repo_id, local_files_only=True)
         except Exception:
@@ -702,6 +717,7 @@ def main(argv):
         dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32,
         weights_dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32,
         attention_kernel=config.attention,
+        scale_shift_order=getattr(config, "scale_shift_order", "shift_scale"),
     )
 
     # 6b. Instantiate JAX FlaxAutoencoderKL VAE

src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py

@@ -453,6 +453,7 @@ class FluxTransformer2DModel(nn.Module, FlaxModelMixin, ConfigMixin):
   flash_min_seq_length: int = 4096
   flash_block_sizes: BlockSizes = None
   mesh: jax.sharding.Mesh = None
+  scale_shift_order: str = "shift_scale"
   dtype: jnp.dtype = jnp.float32
   weights_dtype: jnp.dtype = jnp.float32
   precision: jax.lax.Precision = None
@@ -581,6 +582,7 @@ def setup(self):
         dtype=self.dtype,
         weights_dtype=self.weights_dtype,
         precision=self.precision,
+        scale_shift_order=self.scale_shift_order,
     )
 
     self.proj_out = nn.Dense(

src/maxdiffusion/models/normalization_flax.py

@@ -29,6 +29,7 @@ class AdaLayerNormContinuous(nn.Module):
   dtype: jnp.dtype = jnp.float32
   weights_dtype: jnp.dtype = jnp.float32
   precision: jax.lax.Precision = None
+  scale_shift_order: str = "shift_scale"
 
   @nn.compact
   def __call__(self, x, conditioning_embedding):
@@ -42,7 +43,14 @@ def __call__(self, x, conditioning_embedding):
         param_dtype=self.weights_dtype,
         precision=self.precision,
     )(nn.silu(conditioning_embedding))
-    scale, shift = jnp.split(emb, 2, axis=1)
+
+    if self.scale_shift_order == "scale_shift":
+      scale, shift = jnp.split(emb, 2, axis=1)
+    elif self.scale_shift_order == "shift_scale":
+      shift, scale = jnp.split(emb, 2, axis=1)
+    else:
+      raise ValueError(f"Unsupported scale_shift_order: {self.scale_shift_order}")
+
     scale = nn.with_logical_constraint(scale, ("activation_batch", "activation_embed"))
     shift = nn.with_logical_constraint(shift, ("activation_batch", "activation_embed"))
     x = nn.LayerNorm(epsilon=self.eps, use_bias=self.elementwise_affine, use_scale=self.elementwise_affine)(x)

src/maxdiffusion/tests/flux2klein/generate_flux2klein_9b_parity_test.py

@@ -233,6 +233,7 @@ def test_e2e_parity_9b(self):
             theta=2000,
             mesh=mesh,
             dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32,
+            scale_shift_order=getattr(config, "scale_shift_order", "shift_scale"),
         )
 
         vae = FlaxAutoencoderKL(

src/maxdiffusion/tests/flux2klein/generate_flux2klein_e2e_test.py

@@ -135,6 +135,7 @@ def test_end_to_end_parity_and_offloading(self):
             theta=2000,
             mesh=mesh,
             dtype=jnp.float32,
+            scale_shift_order=getattr(config, "scale_shift_order", "shift_scale"),
         )
 
         # VAE Config

src/maxdiffusion/tests/flux2klein/generate_flux2klein_test.py

@@ -134,6 +134,8 @@ def test_context_embedder_projection(self):
                 "run_name=flux_test",
                 "output_dir=/tmp/",
                 "jax_cache_dir=/tmp/cache_dir",
+                "ici_data_parallelism=1",
+                "ici_fsdp_parallelism=1",
             ], unittest=True)
         config = pyconfig.config
 
@@ -372,6 +374,8 @@ def test_attention_blocks_parity(self):
                 "run_name=flux_test",
                 "output_dir=/tmp/",
                 "jax_cache_dir=/tmp/cache_dir",
+                "ici_data_parallelism=1",
+                "ici_fsdp_parallelism=1",
             ], unittest=True)
         config = pyconfig.config
 
@@ -626,6 +630,8 @@ def test_full_transformer_and_multistep_parity(self):
                 "run_name=flux_test",
                 "output_dir=/tmp/",
                 "jax_cache_dir=/tmp/cache_dir",
+                "ici_data_parallelism=1",
+                "ici_fsdp_parallelism=1",
             ], unittest=True)
         config = pyconfig.config