Onboard FLUX.2-klein pipeline and models into MaxDiffusion (JAX+TPU)

Author: amepas

src/maxdiffusion/configs/base_flux2klein.yml

@@ -161,9 +161,9 @@ dcn_fsdp_parallelism: -1
 dcn_context_parallelism: 1
 dcn_tensor_parallelism: 1
 ici_data_parallelism: 1
-ici_fsdp_parallelism: 1
+ici_fsdp_parallelism: -1
 ici_context_parallelism: 1
-ici_tensor_parallelism: 4
+ici_tensor_parallelism: 1
 
 allow_split_physical_axes: False
 

src/maxdiffusion/generate_flux2klein.py

@@ -15,6 +15,7 @@
 import gc
 import os
 import time
+import sys
 from typing import List
 
 from absl import app
@@ -30,7 +31,7 @@
 from maxdiffusion.max_utils import create_device_mesh
 from maxdiffusion.train_utils import transformer_engine_context
 
-from maxdiffusion.models.flux.transformers.transformer_flux2klein_flax import Flux2KleinTransformer2DModel
+from maxdiffusion.models.flux.transformers.transformer_flux_flax import Flux2KleinTransformer2DModel
 from maxdiffusion.models.vae_flax import FlaxAutoencoderKL
 from maxdiffusion.models.qwen3_flax import FlaxQwen3Config, FlaxQwen3Model, load_and_convert_qwen3_weights
 from maxdiffusion.schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler