refactoring code + new more streamlined test scripts

Author: amepas

src/maxdiffusion/configs/base_flux2klein.yml

@@ -270,3 +270,9 @@ width: 1024
 batch_size: 4
 interactive: False
 
+# 4B Architecture Dimensions
+depth: 20 # num_single_layers
+num_double_layers: 5
+hidden_size: 3072
+num_attention_heads: 24
+

src/maxdiffusion/configs/base_flux_dev.yml

@@ -33,6 +33,7 @@ t5xxl_model_name_or_path: 'ariG23498/t5-v1-1-xxl-flax'
 
 # Flux params
 flux_name: "flux-dev"
+scale_shift_order: "shift_scale"
 max_sequence_length: 512
 time_shift: True
 base_shift: 0.5

src/maxdiffusion/generate_flux2klein.py

@@ -462,12 +462,15 @@ def get_w(name: str, transpose: bool = True) -> np.ndarray:
 
         else:
             print(f"WARNING: JAX parameter '{path_str}' did not match any PyTorch weights!")
-            converted_flat[k] = v
+            converted_flat[k] = np.zeros(v.shape, dtype=np.float32) if hasattr(v, 'shape') and not isinstance(v, np.ndarray) else v
 
     # Clean up PyTorch memory immediately
     del torch_weights
     import gc
     gc.collect()
 
-    print("Qwen3 weight conversion completed successfully!")
-    return flax.traverse_util.unflatten_dict(converted_flat)
+    res = flax.traverse_util.unflatten_dict(converted_flat)
+    return jax.tree_util.tree_map(
+        lambda leaf: jnp.zeros(leaf.shape, dtype=leaf.dtype) if isinstance(leaf, jax.ShapeDtypeStruct) else leaf,
+        res
+    )

src/maxdiffusion/models/flux/util.py

@@ -19,3 +19,6 @@
 from .flux_pipeline import (
     FluxPipeline,
 )
+from .flux2klein_pipeline import (
+    FlaxFlux2KleinPipeline,
+)