diff --git a/src/maxdiffusion/configs/base_flux2klein.yml b/src/maxdiffusion/configs/base_flux2klein.yml
index f1a7e0fa4..3c5d8cc04 100644
--- a/src/maxdiffusion/configs/base_flux2klein.yml
+++ b/src/maxdiffusion/configs/base_flux2klein.yml
@@ -40,6 +40,9 @@ max_sequence_length: 512
 time_shift: True
 base_shift: 0.5
 max_shift: 1.15
+image_paths: []
+use_base2_exp: True
+use_kv: False
 
 
 unet_checkpoint: ''
diff --git a/src/maxdiffusion/configs/base_flux2klein_9B.yml b/src/maxdiffusion/configs/base_flux2klein_9B.yml
index a3a0afeac..e3d20dda1 100644
--- a/src/maxdiffusion/configs/base_flux2klein_9B.yml
+++ b/src/maxdiffusion/configs/base_flux2klein_9B.yml
@@ -40,6 +40,9 @@ max_sequence_length: 512
 time_shift: True
 base_shift: 0.5
 max_shift: 1.15
+image_paths: []
+use_base2_exp: True
+use_kv: False
 
 
 unet_checkpoint: ''
diff --git a/src/maxdiffusion/generate_flux2klein.py b/src/maxdiffusion/generate_flux2klein.py
index b1427f937..1230de568 100644
--- a/src/maxdiffusion/generate_flux2klein.py
+++ b/src/maxdiffusion/generate_flux2klein.py
@@ -20,6 +20,7 @@
 import sys
 from typing import List
 
+from PIL import Image, UnidentifiedImageError
 from absl import app
 import jax
 import jax.numpy as jnp
@@ -35,7 +36,10 @@
 from maxdiffusion.max_utils import create_device_mesh
 from maxdiffusion.train_utils import transformer_engine_context
 
-from maxdiffusion.models.vae_flax import FlaxAutoencoderKL
+from maxdiffusion.models.flux.vae.autoencoder_kl_flux2_nnx import (
+    NNXAutoencoderKLFlux2,
+    load_and_convert_flux2klein_nnx_vae_weights,
+)
 from maxdiffusion.models.qwen3_flax import FlaxQwen3Config, FlaxQwen3Model
 from maxdiffusion.models.qwen3_utils import load_and_convert_qwen3_weights
 from maxdiffusion.schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler
@@ -144,11 +148,6 @@ def main(argv):
 
   pyconfig.initialize(default_args)
 
-  # Import modules after jax.distributed.initialize() has run via pyconfig.initialize()
-  from maxdiffusion.models.flux.util import (
-      load_and_convert_vae_weights,
-  )
-
   config = pyconfig.config
   os.makedirs(config.output_dir, exist_ok=True)
 
@@ -220,6 +219,17 @@ def main(argv):
   repo_id = getattr(config, "pretrained_model_name_or_path", None)
   if not repo_id:
     raise ValueError("pretrained_model_name_or_path must be specified in configuration YAML or CLI.")
+
+  use_kv = getattr(config, "use_kv", False)
+  if use_kv:
+    if repo_id in ("black-forest-labs/FLUX.2-klein-4B", "black-forest-labs/FLUX.2-klein-4b"):
+      max_logging.log("⚠️ Warning: KV cache not supported for 4B model, ignoring use_kv=True.")
+      pyconfig._config.keys["use_kv"] = False
+    elif repo_id in ("black-forest-labs/FLUX.2-klein-9B", "black-forest-labs/FLUX.2-klein-9b"):
+      repo_id = "black-forest-labs/FLUX.2-klein-9b-kv"
+      pyconfig._config.keys["pretrained_model_name_or_path"] = repo_id
+      max_logging.log(f"ℹ️ use_kv=True: switched pretrained_model_name_or_path to KV model variant: {repo_id}")
+
   max_logging.log(f"Target model detected: {repo_id}")
 
   if os.path.exists(repo_id):
@@ -240,9 +250,10 @@ def main(argv):
   safetensors_path = os.path.join(snapshot_dir, "transformer")
   vae_safetensors_path = os.path.join(snapshot_dir, "vae", "diffusion_pytorch_model.safetensors")
   text_encoder_path = os.path.join(snapshot_dir, "text_encoder")
+  tokenizer_path = os.path.join(snapshot_dir, "tokenizer")
 
   # 4. Load Qwen3 Config & Setup model layout
-  from transformers import AutoConfig
+  from transformers import AutoConfig, Qwen2TokenizerFast
   from maxdiffusion.max_utils import get_flash_block_sizes
   from flax import nnx
   from maxdiffusion.models.flux.transformers.transformer_flux_flax import NNXFlux2KleinTransformer2DModel
@@ -267,10 +278,10 @@ def main(argv):
       intermediate_size=pt_config.intermediate_size,
       num_hidden_layers=pt_config.num_hidden_layers,
       num_attention_heads=pt_config.num_attention_heads,
-      num_key_value_heads=pt_config.num_key_value_heads,
-      max_position_embeddings=pt_config.max_position_embeddings,
-      rms_norm_eps=pt_config.rms_norm_eps,
-      rope_theta=pt_config.rope_theta,
+      num_key_value_heads=getattr(pt_config, "num_key_value_heads", pt_config.num_attention_heads),
+      max_position_embeddings=getattr(pt_config, "max_position_embeddings", 32768),
+      rms_norm_eps=getattr(pt_config, "rms_norm_eps", 1e-6),
+      rope_theta=getattr(pt_config, "rope_theta", getattr(pt_config, "rope_base", 1000000.0)),
       dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32,
       attention_kernel=getattr(config, "text_encoder_attention", "flash"),
       flash_block_sizes=te_bs,
@@ -299,6 +310,14 @@ def main(argv):
   num_attention_heads = getattr(config, "num_attention_heads", None) or transformer_pt_cfg.get("num_attention_heads", 24)
 
   # 5. Instantiate JAX NNXFlux2KleinTransformer2DModel
+  guidance_embeds = getattr(config, "guidance_embeds", None)
+  if guidance_embeds is None:
+    guidance_embeds = transformer_pt_cfg.get("guidance_embeds", False)
+
+  pooled_projection_dim = getattr(config, "pooled_projection_dim", None)
+  if pooled_projection_dim is None:
+    pooled_projection_dim = transformer_pt_cfg.get("pooled_projection_dim", None)
+
   transformer = NNXFlux2KleinTransformer2DModel(
       rngs=nnx.Rngs(0),
       in_channels=128,
@@ -307,8 +326,8 @@ def main(argv):
       attention_head_dim=128,
       num_attention_heads=num_attention_heads,
       joint_attention_dim=3 * pt_config.hidden_size,
-      pooled_projection_dim=768,
-      guidance_embeds=True,
+      pooled_projection_dim=pooled_projection_dim,
+      guidance_embeds=guidance_embeds,
       axes_dim=(32, 32, 32, 32),
       theta=2000.0,
       mlp_ratio=3.0,
@@ -321,56 +340,42 @@ def main(argv):
       scale_shift_order=getattr(config, "scale_shift_order", "scale_shift"),
       ulysses_shards=getattr(config, "ulysses_shards", -1),
       ulysses_attention_chunks=getattr(config, "ulysses_attention_chunks", 1),
+      use_base2_exp=getattr(config, "use_base2_exp", True),
   )
 
-  # 6. Instantiate JAX VAE
-  vae = FlaxAutoencoderKL(
-      in_channels=3,
-      out_channels=3,
-      down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"),
-      up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"),
-      block_out_channels=(128, 256, 512, 512),
-      layers_per_block=2,
-      act_fn="silu",
-      latent_channels=32,
-      norm_num_groups=32,
-      sample_size=512,
-      use_quant_conv=True,
-      use_post_quant_conv=True,
-      dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32,
+  # 6. Instantiate JAX NNX VAE
+  weight_dtype = jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32
+  vae = NNXAutoencoderKLFlux2(
+      dtype=weight_dtype,
+      param_dtype=weight_dtype,
   )
 
   # 7. Evaluate shapes & extract mesh shardings
   max_logging.log("Evaluating model shapes and shardings...")
   seq_len_txt = config.max_sequence_length
-  dummy_img = jnp.zeros((config.batch_size, 3, 512, 512))
   dummy_ids = jnp.zeros((config.batch_size, seq_len_txt), dtype=jnp.int32)
   dummy_mask = jnp.zeros((config.batch_size, seq_len_txt), dtype=jnp.int32)
 
   key = jax.random.PRNGKey(0)
-  vae_key, qwen_key = jax.random.split(key, 2)
+  qwen_key = jax.random.split(key, 1)[0]
 
   abstract_state = nnx.state(transformer, nnx.Param)
-
-  def vae_init_fn():
-    return vae.init(vae_key, dummy_img)
+  abstract_vae_state = nnx.state(vae, nnx.Param)
 
   def qwen3_init_fn():
     return qwen3_model.init(qwen_key, dummy_ids, dummy_mask)
 
   with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
     logical_transformer_specs = nnx.get_partition_spec(abstract_state)
-    abstract_vae_vars = jax.eval_shape(vae_init_fn)
+    logical_vae_specs = nnx.get_partition_spec(abstract_vae_state)
     abstract_qwen3_vars = jax.eval_shape(qwen3_init_fn)
-
-    logical_vae_specs = nn.get_partition_spec(abstract_vae_vars)
     logical_qwen3_specs = nn.get_partition_spec(abstract_qwen3_vars)
 
     transformer_mesh_shardings = nn.logical_to_mesh_sharding(logical_transformer_specs, mesh, config.logical_axis_rules)
     vae_mesh_shardings = nn.logical_to_mesh_sharding(logical_vae_specs, mesh, config.logical_axis_rules)
     qwen3_mesh_shardings = nn.logical_to_mesh_sharding(logical_qwen3_specs, mesh, config.logical_axis_rules)
 
-  vae_shardings = flax.core.freeze(vae_mesh_shardings["params"])
+  vae_shardings = vae_mesh_shardings
   qwen3_shardings = flax.core.freeze(qwen3_mesh_shardings["params"])
   transformer_shardings = transformer_mesh_shardings
 
@@ -386,33 +391,26 @@ def unbox_fn(x):
         return x.unbox() if isinstance(x, flax_spmd.LogicallyPartitioned) else x
 
       t_sub0 = time.time()
-      vae_params = jax.tree_util.tree_map(
-          unbox_fn, abstract_vae_vars["params"], is_leaf=lambda k: isinstance(k, flax_spmd.LogicallyPartitioned)
-      )
-      vae_params = flax.core.unfreeze(vae_params)
-
       qwen3_params = jax.tree_util.tree_map(
           unbox_fn, abstract_qwen3_vars["params"], is_leaf=lambda k: isinstance(k, flax_spmd.LogicallyPartitioned)
       )
       qwen3_params = flax.core.unfreeze(qwen3_params)
-
       max_logging.log(f" -> [SUB-TIMING 1/3] PyTree unboxing template setup: {time.time() - t_sub0:.2f}s")
-      t_sub1 = time.time()
-
-      weight_dtype = jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32
 
+      t_sub1 = time.time()
       params = load_and_convert_flux_klein_nnx_weights(
           safetensors_path, abstract_state, num_double_layers, depth, dtype=weight_dtype
       )
-      vae_params, vae_bn_mean, vae_bn_std = load_and_convert_vae_weights(
-          vae_safetensors_path, vae_params, dtype=weight_dtype
-      )
+      nnx.update(transformer, params)
+      params = nnx.state(transformer, nnx.Param)
+
+      vae_bn_mean, vae_bn_std = load_and_convert_flux2klein_nnx_vae_weights(vae_safetensors_path, vae, dtype=weight_dtype)
+      vae_params = nnx.state(vae, nnx.Param)
       qwen3_params = load_and_convert_qwen3_weights(text_encoder_path, qwen3_params, qwen3_config)
       max_logging.log(
           f" -> [SUB-TIMING 2/3] Safetensors loading & key mapping (in target dtype): {time.time() - t_sub1:.4f}s"
       )
 
-      vae_params = flax.core.freeze(vae_params)
       qwen3_params = flax.core.freeze(qwen3_params)
 
       max_logging.log("\n" + "=" * 80)
@@ -447,13 +445,18 @@ def unbox_fn(x):
       time_shift_type="exponential",
   )
 
+  try:
+    tokenizer = Qwen2TokenizerFast.from_pretrained(tokenizer_path, local_files_only=True)
+  except Exception:
+    tokenizer = Qwen2TokenizerFast.from_pretrained(snapshot_dir, subfolder="tokenizer", local_files_only=True)
+
   # 10. Instantiate and invoke FlaxFlux2KleinPipeline
   max_logging.log("Instantiating JAX FlaxFlux2KleinPipeline...")
   pipeline = FlaxFlux2KleinPipeline(
       transformer=transformer,
       vae=vae,
       text_encoder=qwen3_model,
-      tokenizer=None,
+      tokenizer=tokenizer,
       scheduler=scheduler,
       config=config,
       mesh=mesh,
@@ -464,6 +467,34 @@ def unbox_fn(x):
     raise ValueError("Prompt must be specified in the configuration YAML or passed via CLI prompt='...'")
   active_prompts = partition_prompts(prompt_str, config.batch_size)
 
+  # Parse reference image paths for multi-image editing if provided
+  images = None
+  image_paths = getattr(config, "image_paths", None)
+  if image_paths is not None:
+    if isinstance(image_paths, str) and image_paths.strip():
+      import ast
+
+      try:
+        image_paths = ast.literal_eval(image_paths)
+      except Exception:
+        image_paths = [p.strip() for p in image_paths.split(",") if p.strip()]
+    if isinstance(image_paths, (list, tuple)) and len(image_paths) > 0:
+      max_logging.log(f" -> Loading {len(image_paths)} reference image(s) for multi-image editing...")
+      images = []
+      for p in image_paths:
+        try:
+          if not os.path.exists(p):
+            raise FileNotFoundError(f"Reference image file not found: {p}")
+          with Image.open(p) as img_raw:
+            img = img_raw.convert("RGB")
+            images.append(img)
+        except (UnidentifiedImageError, OSError, FileNotFoundError) as e:
+          max_logging.log(f"❌ Error loading reference image '{p}': {e}")
+          raise ValueError(f"Failed to load reference image '{p}': {e}") from e
+        except Exception as e:
+          max_logging.log(f"❌ Unexpected error loading reference image '{p}': {e}")
+          raise ValueError(f"Failed to load reference image '{p}': {e}") from e
+
   if getattr(config, "interactive", False):
     max_logging.log("\n" + "=" * 80)
     max_logging.log("   BATCHED INTERACTIVE GENERATION MODE ENABLED 🎮")
@@ -501,6 +532,7 @@ def unbox_fn(x):
           width=config.width,
           num_inference_steps=config.num_inference_steps,
           batch_size=config.batch_size,
+          images=images,
           use_latents=False,
           output_dir=config.output_dir,
           output_name=output_file,
@@ -528,6 +560,7 @@ def unbox_fn(x):
         batch_size=config.batch_size,
         height=config.height,
         width=config.width,
+        images=images,
     )
 
     max_logging.log("\n" + "=" * 80)
@@ -547,6 +580,7 @@ def unbox_fn(x):
         width=config.width,
         num_inference_steps=config.num_inference_steps,
         batch_size=config.batch_size,
+        images=images,
         use_latents=use_latents_flag,
         latents=latents_to_use,
         output_dir=config.output_dir,
@@ -554,7 +588,8 @@ def unbox_fn(x):
         warmup=True,
     )
     warmup_time = (
-        warmup_trace.get("prompt_encoding", 0.0)
+        warmup_trace.get("vae_encode", 0.0)
+        + warmup_trace.get("prompt_encoding", 0.0)
         + warmup_trace.get("denoise_loop", 0.0)
         + warmup_trace.get("vae_decode", 0.0)
     )
@@ -589,6 +624,7 @@ def unbox_fn(x):
               width=config.width,
               num_inference_steps=config.num_inference_steps,
               batch_size=config.batch_size,
+              images=images,
               use_latents=use_latents_flag,
               latents=latents_to_use,
               output_dir=config.output_dir,
@@ -609,6 +645,7 @@ def unbox_fn(x):
             width=config.width,
             num_inference_steps=config.num_inference_steps,
             batch_size=config.batch_size,
+            images=images,
             use_latents=use_latents_flag,
             latents=latents_to_use,
             output_dir=config.output_dir,
@@ -617,16 +654,22 @@ def unbox_fn(x):
 
       tot_time_i = trace_i.get(
           "e2e_pipeline_total",
-          trace_i.get("prompt_encoding", 0.0) + trace_i.get("denoise_loop", 0.0) + trace_i.get("vae_decode", 0.0),
+          trace_i.get("vae_encode", 0.0)
+          + trace_i.get("prompt_encoding", 0.0)
+          + trace_i.get("denoise_loop", 0.0)
+          + trace_i.get("vae_decode", 0.0),
       )
       main_traces.append(trace_i)
       main_times.append(tot_time_i)
       if num_reps > 1:
+        vae_enc_str = f" | VAE_Enc={trace_i.get('vae_encode', 0.0):.4f}s" if trace_i.get("vae_encode", 0.0) > 0 else ""
         max_logging.log(
-            f"   -> Rep {rep+1}/{num_reps} Completed: Total={tot_time_i:.4f}s | Qwen3={trace_i.get('qwen3_encoding', 0.0):.4f}s | Denoise={trace_i.get('denoise_loop', 0.0):.4f}s | VAE={trace_i.get('vae_decode', 0.0):.4f}s"
+            f"   -> Rep {rep+1}/{num_reps} Completed: Total={tot_time_i:.4f}s{vae_enc_str} | Qwen3={trace_i.get('qwen3_encoding', 0.0):.4f}s | Denoise={trace_i.get('denoise_loop', 0.0):.4f}s | VAE_Dec={trace_i.get('vae_decode', 0.0):.4f}s"
         )
 
     avg_main_time = sum(main_times) / num_reps
+    avg_vae_encode = sum(tr.get("vae_encode", 0.0) for tr in main_traces) / num_reps
+    avg_vae_to_qwen3 = sum(tr.get("vae_encode_to_qwen3", 0.0) for tr in main_traces) / num_reps
     avg_start_to_qwen3 = sum(tr.get("start_to_qwen3", 0.0) for tr in main_traces) / num_reps
     avg_prompt_enc = sum(tr.get("qwen3_encoding", tr.get("prompt_encoding", 0.0)) for tr in main_traces) / num_reps
     avg_qwen3_to_denoise = sum(tr.get("qwen3_to_denoise", 0.0) for tr in main_traces) / num_reps
@@ -643,19 +686,38 @@ def unbox_fn(x):
     max_logging.log(f"1) Model Loading & Placement Time:              {load_time:.4f} seconds ⏱️")
     max_logging.log(f"2) Concurrent AOT XLA Compilation Time:         {aot_time:.4f} seconds ⚡")
     max_logging.log(f"3) Warmup Pass Execution Time:                   {warmup_time:.4f} seconds ⏱️")
+    if warmup_trace.get("vae_encode", 0.0) > 0:
+      max_logging.log(f"   - VAE Encoding:    {warmup_trace.get('vae_encode', 0.0):.4f}s")
     max_logging.log(f"   - Qwen3 Encoding:  {warmup_trace.get('prompt_encoding', 0.0):.4f}s")
     max_logging.log(f"   - Flux Denoising:  {warmup_trace.get('denoise_loop', 0.0):.4f}s")
     max_logging.log(f"   - VAE Decoding:    {warmup_trace.get('vae_decode', 0.0):.4f}s")
     max_logging.log(f"👉 TOTAL COLD-START TIME (Loading + AOT + Warmup): {total_cold_start:.4f} seconds 🎯")
     rep_label = f" (Average across {num_reps} reps)" if num_reps > 1 else ""
     max_logging.log(f"4) Main Warmed-Up Pass (Pure Inference Latency){rep_label}: {avg_main_time:.4f} seconds ⏱️")
-    max_logging.log(f"   - 1. Start -> Qwen3:          {avg_start_to_qwen3*1000:.2f} ms ({avg_start_to_qwen3:.4f}s)")
-    max_logging.log(f"   - 2. Qwen3 Encoding:         {avg_prompt_enc*1000:.2f} ms ({avg_prompt_enc:.4f}s)")
-    max_logging.log(f"   - 3. Qwen3 -> Denoising:      {avg_qwen3_to_denoise*1000:.2f} ms ({avg_qwen3_to_denoise:.4f}s)")
-    max_logging.log(f"   - 4. Flux Denoising Loop:    {avg_denoise*1000:.2f} ms ({avg_denoise:.4f}s)")
-    max_logging.log(f"   - 5. Denoising -> VAE:       {avg_denoise_to_vae*1000:.2f} ms ({avg_denoise_to_vae:.4f}s)")
-    max_logging.log(f"   - 6. VAE Decoding:           {avg_vae_decode*1000:.2f} ms ({avg_vae_decode:.4f}s)")
-    max_logging.log(f"   - 7. Image Saving:           {avg_image_saving*1000:.2f} ms ({avg_image_saving:.4f}s)")
+    step_num = 1
+    if avg_vae_encode > 0:
+      max_logging.log(f"   - {step_num}. VAE Image Encoding:       {avg_vae_encode*1000:.2f} ms ({avg_vae_encode:.4f}s)")
+      step_num += 1
+      max_logging.log(f"   - {step_num}. VAE -> Qwen3:             {avg_vae_to_qwen3*1000:.2f} ms ({avg_vae_to_qwen3:.4f}s)")
+      step_num += 1
+    else:
+      max_logging.log(
+          f"   - {step_num}. Start -> Qwen3:          {avg_start_to_qwen3*1000:.2f} ms ({avg_start_to_qwen3:.4f}s)"
+      )
+      step_num += 1
+    max_logging.log(f"   - {step_num}. Qwen3 Encoding:         {avg_prompt_enc*1000:.2f} ms ({avg_prompt_enc:.4f}s)")
+    step_num += 1
+    max_logging.log(
+        f"   - {step_num}. Qwen3 -> Denoising:      {avg_qwen3_to_denoise*1000:.2f} ms ({avg_qwen3_to_denoise:.4f}s)"
+    )
+    step_num += 1
+    max_logging.log(f"   - {step_num}. Flux Denoising Loop:    {avg_denoise*1000:.2f} ms ({avg_denoise:.4f}s)")
+    step_num += 1
+    max_logging.log(f"   - {step_num}. Denoising -> VAE:       {avg_denoise_to_vae*1000:.2f} ms ({avg_denoise_to_vae:.4f}s)")
+    step_num += 1
+    max_logging.log(f"   - {step_num}. VAE Decoding:           {avg_vae_decode*1000:.2f} ms ({avg_vae_decode:.4f}s)")
+    step_num += 1
+    max_logging.log(f"   - {step_num}. Image Saving:           {avg_image_saving*1000:.2f} ms ({avg_image_saving:.4f}s)")
     max_logging.log(f"   - 👉 TOTAL E2E PIPELINE:     {avg_main_time*1000:.2f} ms ({avg_main_time:.4f}s)")
     max_logging.log("=" * 80)
 
diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py
index 5a2754b6d..be4e46176 100644
--- a/src/maxdiffusion/models/attention_flax.py
+++ b/src/maxdiffusion/models/attention_flax.py
@@ -274,40 +274,44 @@ def _select_flash_block_sizes(
     dtype: jnp.dtype,
     attention_kernel: str,
 ) -> BlockSizes:
+  """Selects Flash/Splash attention block sizes, clamping only for short sequences and preserving user config."""
   query_seq_len = _flash_sequence_length(query)
   key_seq_len = _flash_sequence_length(key)
+  default_max_block = 1024 if dtype == jnp.bfloat16 else 512
+  use_tokamax = "tokamax" in attention_kernel
 
-  q_max_block_size = 1024 if dtype == jnp.bfloat16 else 512
-  if key_seq_len != query_seq_len:
-    kv_max_block_size = ((key_seq_len + 127) // 128) * 128
-  else:
-    kv_max_block_size = q_max_block_size
-
-  # Custom kernels use a lightweight carrier that omits the standard Splash
-  # backward fields. A remapped/local standard kernel still needs a complete
-  # BlockSizes object, including when cross-attention happens to have q_len ==
-  # kv_len.
-  if flash_block_sizes is not None and not hasattr(flash_block_sizes, "use_fused_bwd_kernel"):
-    flash_block_sizes = _coerce_tokamax_block_sizes(flash_block_sizes)
-
-  # Keep configured block sizes for self-attention, but let
-  # cross-attention derive safe KV-aware sizes when q_len != kv_len.
-  if flash_block_sizes and key_seq_len == query_seq_len:
-    if attention_kernel in ["tokamax_flash", "tokamax_ring"]:
-      return _coerce_tokamax_block_sizes(flash_block_sizes)
-    return flash_block_sizes
-
-  block_size_q = flash_block_sizes.block_q if flash_block_sizes else q_max_block_size
-  use_tokamax = attention_kernel in ["tokamax_flash", "tokamax_ring"]
+  if flash_block_sizes is not None:
+    user_bkv = getattr(flash_block_sizes, "block_kv", flash_block_sizes.block_q)
+    kv_max_bound = ((key_seq_len + 127) // 128) * 128
+    safe_bkv = min(user_bkv, kv_max_bound)
+
+    user_bq = flash_block_sizes.block_q
+    q_max_bound = ((query_seq_len + 127) // 128) * 128
+    safe_bq = min(user_bq, q_max_bound)
+
+    return splash_attention_kernel.BlockSizes(
+        block_q=safe_bq,
+        block_kv=safe_bkv,
+        block_kv_compute=min(getattr(flash_block_sizes, "block_kv_compute", safe_bkv), safe_bkv),
+        block_q_dkv=safe_bq,
+        block_kv_dkv=safe_bkv,
+        block_kv_dkv_compute=min(safe_bkv, query_seq_len),
+        block_q_dq=None if use_tokamax else safe_bq,
+        block_kv_dq=None if use_tokamax else min(safe_bkv, query_seq_len),
+        use_fused_bwd_kernel=True if use_tokamax else False,
+    )
+
+  block_q = min(default_max_block, query_seq_len)
+  block_kv = min(default_max_block, key_seq_len)
   return splash_attention_kernel.BlockSizes(
-      block_q=block_size_q,
-      block_kv_compute=min(kv_max_block_size, key_seq_len),
-      block_kv=min(kv_max_block_size, key_seq_len),
-      block_q_dkv=block_size_q,
-      block_kv_dkv=min(kv_max_block_size, key_seq_len),
-      block_kv_dkv_compute=min(kv_max_block_size, query_seq_len),
-      block_q_dq=None if use_tokamax else block_size_q,
-      block_kv_dq=None if use_tokamax else min(kv_max_block_size, query_seq_len),
+      block_q=block_q,
+      block_kv=block_kv,
+      block_kv_compute=block_kv,
+      block_q_dkv=block_q,
+      block_kv_dkv=block_kv,
+      block_kv_dkv_compute=min(block_kv, query_seq_len),
+      block_q_dq=None if use_tokamax else block_q,
+      block_kv_dq=None if use_tokamax else min(block_kv, query_seq_len),
       use_fused_bwd_kernel=True if use_tokamax else False,
   )
 
@@ -1963,7 +1967,8 @@ def _apply_attention(
 
   # Module-level Registry lookup
   if effective_attention_kernel in KERNEL_REGISTRY:
-    return KERNEL_REGISTRY[effective_attention_kernel](query, key, value, context)
+    with jax.named_scope(f"kernel_{effective_attention_kernel}"):
+      return KERNEL_REGISTRY[effective_attention_kernel](query, key, value, context)
 
   raise ValueError(f"Unexpected attention kernel {effective_attention_kernel=}.")
 
diff --git a/src/maxdiffusion/models/embeddings_flax.py b/src/maxdiffusion/models/embeddings_flax.py
index 61e7956ce..17ac2b07a 100644
--- a/src/maxdiffusion/models/embeddings_flax.py
+++ b/src/maxdiffusion/models/embeddings_flax.py
@@ -615,7 +615,7 @@ def __init__(
           weights_dtype=weights_dtype,
       )
 
-    if pooled_projection_dim > 0:
+    if pooled_projection_dim is not None and pooled_projection_dim > 0:
       self.pooled_embedder = NNXPixArtAlphaTextProjection(
           rngs=rngs,
           in_features=pooled_projection_dim,
@@ -643,7 +643,7 @@ def __call__(
     else:
       time_guidance_emb = timestep_emb
 
-    if pooled_projection is not None and self.pooled_projection_dim > 0:
+    if pooled_projection is not None and self.pooled_projection_dim is not None and self.pooled_projection_dim > 0:
       pooled_projections = self.pooled_embedder(pooled_projection)
       conditioning = time_guidance_emb + pooled_projections
     else:
diff --git a/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py b/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py
index 42bfca5d3..d3b411f83 100644
--- a/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py
+++ b/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py
@@ -14,7 +14,7 @@
 limitations under the License.
 """
 
-from typing import Dict, Optional, Tuple, Union
+from typing import Any, Dict, Optional, Tuple, Union
 import jax
 import math
 import jax.numpy as jnp
@@ -1333,6 +1333,7 @@ def __init__(
       qkv_bias: bool = False,
       ulysses_shards: int = -1,
       ulysses_attention_chunks: int = 1,
+      use_base2_exp: bool = False,
   ):
     self.heads = heads
     self.dim_head = dim_head
@@ -1352,6 +1353,7 @@ def __init__(
         split_head_dim=False,
         ulysses_shards=ulysses_shards,
         ulysses_attention_chunks=ulysses_attention_chunks,
+        use_base2_exp=use_base2_exp,
     )
 
     kernel_axes = ("embed", "heads")
@@ -1433,58 +1435,108 @@ def __call__(
       hidden_states: jax.Array,
       encoder_hidden_states: Optional[jax.Array] = None,
       image_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None,
-  ) -> Tuple[jax.Array, Optional[jax.Array]]:
+      kv_cache: Optional[Tuple[jax.Array, jax.Array]] = None,
+      kv_cache_mode: Optional[str] = None,
+      num_ref_tokens: int = 0,
+  ) -> Tuple[Tuple[jax.Array, Optional[jax.Array]], Optional[Tuple[jax.Array, jax.Array]]]:
     B, L = hidden_states.shape[:2]
     H, D = self.heads, self.dim_head
 
-    qkv_proj = self.i_qkv(hidden_states).reshape(B, L, 3, H, D)
-    query_proj, key_proj, value_proj = jnp.split(qkv_proj, 3, axis=2)
-    query_proj = self.query_norm(query_proj.squeeze(2))
-    key_proj = self.key_norm(key_proj.squeeze(2))
-    value_proj = value_proj.squeeze(2)
-
-    if encoder_hidden_states is not None:
-      B_enc, L_txt = encoder_hidden_states.shape[:2]
-      encoder_qkv_proj = self.e_qkv(encoder_hidden_states).reshape(B_enc, L_txt, 3, H, D)
-      enc_query_proj, enc_key_proj, enc_value_proj = jnp.split(encoder_qkv_proj, 3, axis=2)
-      enc_query_proj = self.encoder_query_norm(enc_query_proj.squeeze(2))
-      enc_key_proj = self.encoder_key_norm(enc_key_proj.squeeze(2))
-      enc_value_proj = enc_value_proj.squeeze(2)
-
-      query_proj = jnp.concatenate((enc_query_proj, query_proj), axis=1)
-      key_proj = jnp.concatenate((enc_key_proj, key_proj), axis=1)
-      value_proj = jnp.concatenate((enc_value_proj, value_proj), axis=1)
+    with jax.named_scope("qkv_projections"):
+      qkv_proj = self.i_qkv(hidden_states).reshape(B, L, 3, H, D)
+      query_proj, key_proj, value_proj = jnp.split(qkv_proj, 3, axis=2)
+      query_proj = self.query_norm(query_proj.squeeze(2))
+      key_proj = self.key_norm(key_proj.squeeze(2))
+      value_proj = value_proj.squeeze(2)
+
+      if encoder_hidden_states is not None:
+        B_enc, L_txt = encoder_hidden_states.shape[:2]
+        encoder_qkv_proj = self.e_qkv(encoder_hidden_states).reshape(B_enc, L_txt, 3, H, D)
+        enc_query_proj, enc_key_proj, enc_value_proj = jnp.split(encoder_qkv_proj, 3, axis=2)
+        enc_query_proj = self.encoder_query_norm(enc_query_proj.squeeze(2))
+        enc_key_proj = self.encoder_key_norm(enc_key_proj.squeeze(2))
+        enc_value_proj = enc_value_proj.squeeze(2)
+
+        query_proj = jnp.concatenate((enc_query_proj, query_proj), axis=1)
+        key_proj = jnp.concatenate((enc_key_proj, key_proj), axis=1)
+        value_proj = jnp.concatenate((enc_value_proj, value_proj), axis=1)
 
     if image_rotary_emb is not None:
-      if not isinstance(image_rotary_emb, (tuple, list)):
-        image_rotary_emb_reordered = rearrange(image_rotary_emb, "n d (i j) -> n d i j", i=2, j=2)
+      with jax.named_scope("rope_embeddings"):
+        if not isinstance(image_rotary_emb, (tuple, list)):
+          image_rotary_emb_reordered = rearrange(image_rotary_emb, "n d (i j) -> n d i j", i=2, j=2)
+        else:
+          image_rotary_emb_reordered = image_rotary_emb
+        query_proj = query_proj.swapaxes(1, 2)
+        key_proj = key_proj.swapaxes(1, 2)
+        query_proj, key_proj = apply_rope(query_proj, key_proj, image_rotary_emb_reordered)
+        query_proj = query_proj.swapaxes(1, 2)
+        key_proj = key_proj.swapaxes(1, 2)
+
+    layer_kv = None
+    num_txt_tokens = encoder_hidden_states.shape[1] if encoder_hidden_states is not None else 0
+
+    with jax.named_scope("joint_attention_op"):
+      if kv_cache_mode == "extract" and num_ref_tokens > 0:
+        ref_start = num_txt_tokens
+        ref_end = num_txt_tokens + num_ref_tokens
+        k_ref = key_proj[:, ref_start:ref_end, :, :]
+        v_ref = value_proj[:, ref_start:ref_end, :, :]
+        layer_kv = (k_ref, v_ref)
+
+        q_txt = query_proj[:, :ref_start]
+        q_ref = query_proj[:, ref_start:ref_end]
+        q_img = query_proj[:, ref_end:]
+
+        q_txt_img = jnp.concatenate([q_txt, q_img], axis=1).reshape(B, -1, H * D)
+        k_all = key_proj.reshape(B, -1, H * D)
+        v_all = value_proj.reshape(B, -1, H * D)
+
+        attn_txt_img = self.attention_op.apply_attention(q_txt_img, k_all, v_all)
+        attn_txt = attn_txt_img[:, :ref_start]
+        attn_img = attn_txt_img[:, ref_start:]
+
+        q_ref_flat = q_ref.reshape(B, -1, H * D)
+        k_ref_flat = k_ref.reshape(B, -1, H * D)
+        v_ref_flat = v_ref.reshape(B, -1, H * D)
+        attn_ref = self.attention_op.apply_attention(q_ref_flat, k_ref_flat, v_ref_flat)
+
+        attn_output = jnp.concatenate([attn_txt, attn_ref, attn_img], axis=1)
+
+      elif kv_cache_mode == "cached" and kv_cache is not None:
+        k_ref, v_ref = kv_cache
+        k_txt = key_proj[:, :num_txt_tokens]
+        k_img = key_proj[:, num_txt_tokens:]
+        v_txt = value_proj[:, :num_txt_tokens]
+        v_img = value_proj[:, num_txt_tokens:]
+
+        k_all = jnp.concatenate([k_txt, k_ref, k_img], axis=1).reshape(B, -1, H * D)
+        v_all = jnp.concatenate([v_txt, v_ref, v_img], axis=1).reshape(B, -1, H * D)
+        q_all = query_proj.reshape(B, -1, H * D)
+
+        attn_output = self.attention_op.apply_attention(q_all, k_all, v_all)
+
       else:
-        image_rotary_emb_reordered = image_rotary_emb
-      query_proj = query_proj.swapaxes(1, 2)
-      key_proj = key_proj.swapaxes(1, 2)
-      query_proj, key_proj = apply_rope(query_proj, key_proj, image_rotary_emb_reordered)
-      query_proj = query_proj.swapaxes(1, 2)
-      key_proj = key_proj.swapaxes(1, 2)
+        query_proj = query_proj.reshape(B, -1, H * D)
+        key_proj = key_proj.reshape(B, -1, H * D)
+        value_proj = value_proj.reshape(B, -1, H * D)
 
-    query_proj = query_proj.reshape(B, -1, H * D)
-    key_proj = key_proj.reshape(B, -1, H * D)
-    value_proj = value_proj.reshape(B, -1, H * D)
+        if encoder_hidden_states is not None:
+          query_proj = nn.with_logical_constraint(query_proj, ("activation_batch", "activation_length", "activation_heads"))
+          key_proj = nn.with_logical_constraint(key_proj, ("activation_batch", "activation_length", "activation_heads"))
+          value_proj = nn.with_logical_constraint(value_proj, ("activation_batch", "activation_length", "activation_heads"))
 
-    if encoder_hidden_states is not None:
-      query_proj = nn.with_logical_constraint(query_proj, ("activation_batch", "activation_length", "activation_heads"))
-      key_proj = nn.with_logical_constraint(key_proj, ("activation_batch", "activation_length", "activation_heads"))
-      value_proj = nn.with_logical_constraint(value_proj, ("activation_batch", "activation_length", "activation_heads"))
+        attn_output = self.attention_op.apply_attention(query_proj, key_proj, value_proj)
 
-    attn_output = self.attention_op.apply_attention(query_proj, key_proj, value_proj)
     context_attn_output = None
-
     if encoder_hidden_states is not None:
-      context_attn_output = attn_output[:, : encoder_hidden_states.shape[1]]
-      attn_output = attn_output[:, encoder_hidden_states.shape[1] :]
-      attn_output = self.i_proj(attn_output)
-      context_attn_output = self.e_proj(context_attn_output)
+      with jax.named_scope("attention_out_projections"):
+        context_attn_output = attn_output[:, : encoder_hidden_states.shape[1]]
+        attn_output = attn_output[:, encoder_hidden_states.shape[1] :]
+        attn_output = self.i_proj(attn_output)
+        context_attn_output = self.e_proj(context_attn_output)
 
-    return attn_output, context_attn_output
+    return (attn_output, context_attn_output), layer_kv
 
 
 class NNXFluxSingleAttention(nnx.Module):
@@ -1504,6 +1556,7 @@ def __init__(
       weights_dtype: jnp.dtype = jnp.float32,
       ulysses_shards: int = -1,
       ulysses_attention_chunks: int = 1,
+      use_base2_exp: bool = False,
   ):
     self.num_attention_heads = num_attention_heads
     self.attention_head_dim = attention_head_dim
@@ -1522,6 +1575,7 @@ def __init__(
         split_head_dim=False,
         ulysses_shards=ulysses_shards,
         ulysses_attention_chunks=ulysses_attention_chunks,
+        use_base2_exp=use_base2_exp,
     )
     self.query_norm = nnx.RMSNorm(
         num_features=attention_head_dim,
@@ -1560,6 +1614,7 @@ def __init__(
       qkv_bias: bool = False,
       ulysses_shards: int = -1,
       ulysses_attention_chunks: int = 1,
+      use_base2_exp: bool = False,
   ):
     self.dim = dim
     self.num_heads = num_attention_heads
@@ -1616,6 +1671,7 @@ def __init__(
         qkv_bias=qkv_bias,
         ulysses_shards=ulysses_shards,
         ulysses_attention_chunks=ulysses_attention_chunks,
+        use_base2_exp=use_base2_exp,
     )
 
     self.ff = NNXFlaxSwiGluFeedForward(
@@ -1643,46 +1699,58 @@ def __call__(
       image_rotary_emb: Tuple[jax.Array, jax.Array],
       temb_mod_img: Optional[jax.Array] = None,
       temb_mod_txt: Optional[jax.Array] = None,
-  ) -> Tuple[jax.Array, jax.Array]:
+      kv_cache: Optional[Tuple[jax.Array, jax.Array]] = None,
+      kv_cache_mode: Optional[str] = None,
+      num_ref_tokens: int = 0,
+  ) -> Tuple[jax.Array, jax.Array, Optional[Tuple[jax.Array, jax.Array]]]:
     shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = jnp.split(temb_mod_img, 6, axis=-1)
     c_shift_msa, c_scale_msa, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = jnp.split(temb_mod_txt, 6, axis=-1)
 
-    shift_msa = jnp.expand_dims(shift_msa, axis=1)
-    scale_msa = jnp.expand_dims(scale_msa, axis=1)
-    gate_msa = jnp.expand_dims(gate_msa, axis=1)
-    shift_mlp = jnp.expand_dims(shift_mlp, axis=1)
-    scale_mlp = jnp.expand_dims(scale_mlp, axis=1)
-    gate_mlp = jnp.expand_dims(gate_mlp, axis=1)
-
-    c_shift_msa = jnp.expand_dims(c_shift_msa, axis=1)
-    c_scale_msa = jnp.expand_dims(c_scale_msa, axis=1)
-    c_gate_msa = jnp.expand_dims(c_gate_msa, axis=1)
-    c_shift_mlp = jnp.expand_dims(c_shift_mlp, axis=1)
-    c_scale_mlp = jnp.expand_dims(c_scale_mlp, axis=1)
-    c_gate_mlp = jnp.expand_dims(c_gate_mlp, axis=1)
-
-    norm1_h = self.norm1(hidden_states) * (1.0 + scale_msa) + shift_msa
-    norm1_enc = self.norm1_context(encoder_hidden_states) * (1.0 + c_scale_msa) + c_shift_msa
-
-    attn_img, attn_txt = self.attn(
-        hidden_states=norm1_h,
-        encoder_hidden_states=norm1_enc,
-        image_rotary_emb=image_rotary_emb,
-    )
+    if temb_mod_img.ndim == 2:
+      shift_msa = jnp.expand_dims(shift_msa, axis=1)
+      scale_msa = jnp.expand_dims(scale_msa, axis=1)
+      gate_msa = jnp.expand_dims(gate_msa, axis=1)
+      shift_mlp = jnp.expand_dims(shift_mlp, axis=1)
+      scale_mlp = jnp.expand_dims(scale_mlp, axis=1)
+      gate_mlp = jnp.expand_dims(gate_mlp, axis=1)
+
+    if temb_mod_txt.ndim == 2:
+      c_shift_msa = jnp.expand_dims(c_shift_msa, axis=1)
+      c_scale_msa = jnp.expand_dims(c_scale_msa, axis=1)
+      c_gate_msa = jnp.expand_dims(c_gate_msa, axis=1)
+      c_shift_mlp = jnp.expand_dims(c_shift_mlp, axis=1)
+      c_scale_mlp = jnp.expand_dims(c_scale_mlp, axis=1)
+      c_gate_mlp = jnp.expand_dims(c_gate_mlp, axis=1)
 
-    hidden_states = hidden_states + gate_msa * attn_img
-    encoder_hidden_states = encoder_hidden_states + c_gate_msa * attn_txt
+    with jax.named_scope("norm1_and_modulation"):
+      norm1_h = self.norm1(hidden_states) * (1.0 + scale_msa) + shift_msa
+      norm1_enc = self.norm1_context(encoder_hidden_states) * (1.0 + c_scale_msa) + c_shift_msa
+
+    with jax.named_scope("double_attention"):
+      (attn_img, attn_txt), layer_kv = self.attn(
+          hidden_states=norm1_h,
+          encoder_hidden_states=norm1_enc,
+          image_rotary_emb=image_rotary_emb,
+          kv_cache=kv_cache,
+          kv_cache_mode=kv_cache_mode,
+          num_ref_tokens=num_ref_tokens,
+      )
 
-    norm2_h = self.norm2(hidden_states) * (1.0 + scale_mlp) + shift_mlp
-    norm2_enc = self.norm2_context(encoder_hidden_states) * (1.0 + c_scale_mlp) + c_shift_mlp
+      hidden_states = hidden_states + gate_msa * attn_img
+      encoder_hidden_states = encoder_hidden_states + c_gate_msa * attn_txt
 
-    mlp_output = self.ff(norm2_h)
-    encoder_mlp_output = self.ff_context(norm2_enc)
+    with jax.named_scope("norm2_and_modulation"):
+      norm2_h = self.norm2(hidden_states) * (1.0 + scale_mlp) + shift_mlp
+      norm2_enc = self.norm2_context(encoder_hidden_states) * (1.0 + c_scale_mlp) + c_shift_mlp
 
-    hidden_states = hidden_states + gate_mlp * mlp_output
-    encoder_hidden_states = encoder_hidden_states + c_gate_mlp * encoder_mlp_output
+    with jax.named_scope("double_mlp"):
+      mlp_output = self.ff(norm2_h)
+      encoder_mlp_output = self.ff_context(norm2_enc)
 
-    return encoder_hidden_states, hidden_states
+      hidden_states = hidden_states + gate_mlp * mlp_output
+      encoder_hidden_states = encoder_hidden_states + c_gate_mlp * encoder_mlp_output
+
+    return encoder_hidden_states, hidden_states, layer_kv
 
 
 class NNXFluxSingleTransformerBlock(nnx.Module):
@@ -1703,6 +1771,7 @@ def __init__(
       weights_dtype: jnp.dtype = jnp.float32,
       ulysses_shards: int = -1,
       ulysses_attention_chunks: int = 1,
+      use_base2_exp: bool = False,
   ):
     self.dim = dim
     self.num_attention_heads = num_attention_heads
@@ -1753,6 +1822,7 @@ def __init__(
         weights_dtype=weights_dtype,
         ulysses_shards=ulysses_shards,
         ulysses_attention_chunks=ulysses_attention_chunks,
+        use_base2_exp=use_base2_exp,
     )
 
   def __call__(
@@ -1761,50 +1831,107 @@ def __call__(
       temb: jax.Array,
       image_rotary_emb: Tuple[jax.Array, jax.Array],
       temb_mod: Optional[jax.Array] = None,
-  ) -> jax.Array:
+      kv_cache: Optional[Tuple[jax.Array, jax.Array]] = None,
+      kv_cache_mode: Optional[str] = None,
+      num_txt_tokens: int = 0,
+      num_ref_tokens: int = 0,
+  ) -> Tuple[jax.Array, Optional[Tuple[jax.Array, jax.Array]]]:
     residual = hidden_states
     shift_msa, scale_msa, gate = jnp.split(temb_mod, 3, axis=-1)
-    shift_msa = jnp.expand_dims(shift_msa, axis=1)
-    scale_msa = jnp.expand_dims(scale_msa, axis=1)
-    gate = jnp.expand_dims(gate, axis=1)
 
-    norm_hidden_states = self.norm(hidden_states)
-    norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa
+    if temb_mod.ndim == 2:
+      shift_msa = jnp.expand_dims(shift_msa, axis=1)
+      scale_msa = jnp.expand_dims(scale_msa, axis=1)
+      gate = jnp.expand_dims(gate, axis=1)
 
-    qkv, mlp = jnp.split(self.linear1(norm_hidden_states), [3 * self.dim], axis=-1)
-    qkv = nn.with_logical_constraint(qkv, ("activation_batch", "activation_length", "activation_embed"))
-    mlp = nn.with_logical_constraint(mlp, ("activation_batch", "activation_length", "activation_embed"))
+    with jax.named_scope("norm_and_modulation"):
+      norm_hidden_states = self.norm(hidden_states)
+      norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa
 
-    B, L = hidden_states.shape[:2]
-    H, D = self.num_attention_heads, qkv.shape[-1] // (self.num_attention_heads * 3)
-    qkv_proj = qkv.reshape(B, L, 3, H, D).transpose(2, 0, 3, 1, 4)
-    q, k, v = qkv_proj
+    with jax.named_scope("linear1_qkv_and_mlp"):
+      qkv, mlp = jnp.split(self.linear1(norm_hidden_states), [3 * self.dim], axis=-1)
+      qkv = nn.with_logical_constraint(qkv, ("activation_batch", "activation_length", "activation_embed"))
+      mlp = nn.with_logical_constraint(mlp, ("activation_batch", "activation_length", "activation_embed"))
 
-    q = self.attn.query_norm(q)
-    k = self.attn.key_norm(k)
+      B, L = hidden_states.shape[:2]
+      H, D = self.num_attention_heads, qkv.shape[-1] // (self.num_attention_heads * 3)
+      qkv_proj = qkv.reshape(B, L, 3, H, D).transpose(2, 0, 3, 1, 4)
+      q, k, v = qkv_proj
+
+      q = self.attn.query_norm(q)
+      k = self.attn.key_norm(k)
 
     if image_rotary_emb is not None:
-      if isinstance(image_rotary_emb, (tuple, list)):
-        image_rotary_emb_reordered = image_rotary_emb
-      else:
-        image_rotary_emb_reordered = rearrange(image_rotary_emb, "n d (i j) -> n d i j", i=2, j=2)
-      q, k = apply_rope(q, k, image_rotary_emb_reordered)
+      with jax.named_scope("rope_embeddings"):
+        if isinstance(image_rotary_emb, (tuple, list)):
+          image_rotary_emb_reordered = image_rotary_emb
+        else:
+          image_rotary_emb_reordered = rearrange(image_rotary_emb, "n d (i j) -> n d i j", i=2, j=2)
+        q, k = apply_rope(q, k, image_rotary_emb_reordered)
+
+    layer_kv = None
+    with jax.named_scope("single_attention_op"):
+      if kv_cache_mode == "extract" and num_ref_tokens > 0:
+        ref_start = num_txt_tokens
+        ref_end = num_txt_tokens + num_ref_tokens
+        k_ref = k[:, :, ref_start:ref_end, :].transpose(0, 2, 1, 3)
+        v_ref = v[:, :, ref_start:ref_end, :].transpose(0, 2, 1, 3)
+        layer_kv = (k_ref, v_ref)
+
+        q_txt_img = (
+            jnp.concatenate([q[:, :, :ref_start, :], q[:, :, ref_end:, :]], axis=2)
+            .transpose(0, 2, 1, 3)
+            .reshape(B, -1, H * D)
+        )
+        k_all = k.transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+        v_all = v.transpose(0, 2, 1, 3).reshape(B, -1, H * D)
 
-    q = q.transpose(0, 2, 1, 3).reshape(q.shape[0], q.shape[2], -1)
-    k = k.transpose(0, 2, 1, 3).reshape(k.shape[0], k.shape[2], -1)
-    v = v.transpose(0, 2, 1, 3).reshape(v.shape[0], v.shape[2], -1)
+        attn_txt_img = self.attn.attention_op.apply_attention(q_txt_img, k_all, v_all)
+        attn_txt = attn_txt_img[:, :ref_start]
+        attn_img = attn_txt_img[:, ref_start:]
 
-    attn_output = self.attn.attention_op.apply_attention(q, k, v)
+        q_ref_flat = q[:, :, ref_start:ref_end, :].transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+        k_ref_flat = k[:, :, ref_start:ref_end, :].transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+        v_ref_flat = v[:, :, ref_start:ref_end, :].transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+        attn_ref = self.attn.attention_op.apply_attention(q_ref_flat, k_ref_flat, v_ref_flat)
 
-    mlp1, mlp2 = jnp.split(mlp, 2, axis=-1)
-    mlp_activated = nnx.silu(mlp1) * mlp2
+        attn_output = jnp.concatenate([attn_txt, attn_ref, attn_img], axis=1)
 
-    attn_mlp = jnp.concatenate([attn_output, mlp_activated], axis=2)
-    attn_mlp = nn.with_logical_constraint(attn_mlp, ("activation_batch", "activation_length", "activation_embed"))
-    hidden_states = self.linear2(attn_mlp)
-    hidden_states = gate * hidden_states
-    hidden_states = residual + hidden_states
-    return hidden_states
+      elif kv_cache_mode == "cached" and kv_cache is not None:
+        k_ref, v_ref = kv_cache
+        k_ref_trans = k_ref.transpose(0, 2, 1, 3)
+        v_ref_trans = v_ref.transpose(0, 2, 1, 3)
+
+        k_all = (
+            jnp.concatenate([k[:, :, :num_txt_tokens, :], k_ref_trans, k[:, :, num_txt_tokens:, :]], axis=2)
+            .transpose(0, 2, 1, 3)
+            .reshape(B, -1, H * D)
+        )
+        v_all = (
+            jnp.concatenate([v[:, :, :num_txt_tokens, :], v_ref_trans, v[:, :, num_txt_tokens:, :]], axis=2)
+            .transpose(0, 2, 1, 3)
+            .reshape(B, -1, H * D)
+        )
+        q_all = q.transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+
+        attn_output = self.attn.attention_op.apply_attention(q_all, k_all, v_all)
+
+      else:
+        q_flat = q.transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+        k_flat = k.transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+        v_flat = v.transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+        attn_output = self.attn.attention_op.apply_attention(q_flat, k_flat, v_flat)
+
+    with jax.named_scope("swiglu_and_linear2"):
+      mlp1, mlp2 = jnp.split(mlp, 2, axis=-1)
+      mlp_activated = nnx.silu(mlp1) * mlp2
+
+      attn_mlp = jnp.concatenate([attn_output, mlp_activated], axis=2)
+      attn_mlp = nn.with_logical_constraint(attn_mlp, ("activation_batch", "activation_length", "activation_embed"))
+      hidden_states = self.linear2(attn_mlp)
+      hidden_states = gate * hidden_states
+      hidden_states = residual + hidden_states
+    return hidden_states, layer_kv
 
 
 class NNXFlux2KleinTransformer2DModel(nnx.Module):
@@ -1834,6 +1961,7 @@ def __init__(
       scale_shift_order: str = "scale_shift",
       ulysses_shards: int = -1,
       ulysses_attention_chunks: int = 1,
+      use_base2_exp: bool = False,
   ):
     self.in_channels = in_channels
     self.out_channels = in_channels
@@ -1914,6 +2042,7 @@ def __init__(
                 weights_dtype=weights_dtype,
                 ulysses_shards=ulysses_shards,
                 ulysses_attention_chunks=ulysses_attention_chunks,
+                use_base2_exp=use_base2_exp,
             )
             for _ in range(num_layers)
         ]
@@ -1935,6 +2064,7 @@ def __init__(
                 weights_dtype=weights_dtype,
                 ulysses_shards=ulysses_shards,
                 ulysses_attention_chunks=ulysses_attention_chunks,
+                use_base2_exp=use_base2_exp,
             )
             for _ in range(num_single_layers)
         ]
@@ -1967,59 +2097,177 @@ def __call__(
       txt_ids: Optional[jax.Array] = None,
       guidance: Optional[jax.Array] = None,
       return_dict: bool = True,
-  ) -> Union[jax.Array, Transformer2DModelOutput]:
-    hidden_states = self.x_embedder(hidden_states)
-    timestep = timestep * 1000.0
-    if guidance is not None:
-      guidance = guidance * 1000.0
-    temb = self.time_text_embed(timestep, guidance, pooled_projections)
-    temb = temb.astype(hidden_states.dtype)
-
-    temb_silu = nnx.silu(temb)
-    double_stream_mod_img = self.double_stream_modulation_img(temb_silu)
-    double_stream_mod_txt = self.double_stream_modulation_txt(temb_silu)
-    single_stream_mod = self.single_stream_modulation(temb_silu)
-
-    if encoder_hidden_states is not None:
-      encoder_hidden_states = self.context_embedder(encoder_hidden_states)
-
-    if txt_ids.ndim == 3:
-      txt_ids = txt_ids[0]
-    if img_ids.ndim == 3:
-      img_ids = img_ids[0]
-
-    image_rotary_emb = self.pos_embed(img_ids)
-    text_rotary_emb = self.pos_embed(txt_ids)
-    concat_rotary_emb = (
-        jnp.concatenate([text_rotary_emb[0], image_rotary_emb[0]], axis=0),
-        jnp.concatenate([text_rotary_emb[1], image_rotary_emb[1]], axis=0),
-    )
+      kv_cache: Optional[Any] = None,
+      kv_cache_mode: Optional[str] = None,
+      num_ref_tokens: int = 0,
+      ref_fixed_timestep: float = 0.0,
+  ) -> Union[jax.Array, Transformer2DModelOutput, Tuple[Any, Any]]:
+    with jax.named_scope("input_embeddings"):
+      hidden_states = self.x_embedder(hidden_states)
+      timestep_scaled = timestep * 1000.0
+      guidance_scaled = guidance * 1000.0 if guidance is not None else None
+      temb = self.time_text_embed(timestep_scaled, guidance_scaled, pooled_projections)
+      temb = temb.astype(hidden_states.dtype)
+
+      temb_silu = nnx.silu(temb)
+      double_stream_mod_img = self.double_stream_modulation_img(temb_silu)
+      double_stream_mod_txt = self.double_stream_modulation_txt(temb_silu)
+      single_stream_mod = self.single_stream_modulation(temb_silu)
 
-    for double_block in self.double_blocks:
-      encoder_hidden_states, hidden_states = double_block(
-          hidden_states=hidden_states,
-          encoder_hidden_states=encoder_hidden_states,
-          temb=temb,
-          image_rotary_emb=concat_rotary_emb,
-          temb_mod_img=double_stream_mod_img,
-          temb_mod_txt=double_stream_mod_txt,
-      )
+      if encoder_hidden_states is not None:
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
 
-    num_txt_tokens = encoder_hidden_states.shape[1]
-    hidden_states = jnp.concatenate([encoder_hidden_states, hidden_states], axis=1)
+      if txt_ids.ndim == 3:
+        txt_ids = txt_ids[0]
+      if img_ids.ndim == 3:
+        img_ids = img_ids[0]
 
-    for single_block in self.single_blocks:
-      hidden_states = single_block(
-          hidden_states=hidden_states,
-          temb=temb,
-          image_rotary_emb=concat_rotary_emb,
-          temb_mod=single_stream_mod,
+      image_rotary_emb = self.pos_embed(img_ids)
+      text_rotary_emb = self.pos_embed(txt_ids)
+      concat_rotary_emb = (
+          jnp.concatenate([text_rotary_emb[0], image_rotary_emb[0]], axis=0),
+          jnp.concatenate([text_rotary_emb[1], image_rotary_emb[1]], axis=0),
       )
 
-    hidden_states = hidden_states[:, num_txt_tokens:, ...]
-    hidden_states = self.norm_out(hidden_states, temb)
-    output = self.proj_out(hidden_states)
+    if kv_cache_mode == "extract" and num_ref_tokens > 0:
+      with jax.named_scope("extract_reference_modulation"):
+        num_img_tokens = hidden_states.shape[1] - num_ref_tokens
+        ref_timestep = jnp.full_like(timestep_scaled, ref_fixed_timestep * 1000.0)
+        ref_temb = self.time_text_embed(ref_timestep, guidance_scaled, pooled_projections).astype(hidden_states.dtype)
+        ref_temb_silu = nnx.silu(ref_temb)
+        ref_double_mod_img = self.double_stream_modulation_img(ref_temb_silu)
+        ref_single_mod = self.single_stream_modulation(ref_temb_silu)
+
+        ref_mod_expanded = jnp.repeat(jnp.expand_dims(ref_double_mod_img, 1), num_ref_tokens, axis=1)
+        img_mod_expanded = jnp.repeat(jnp.expand_dims(double_stream_mod_img, 1), num_img_tokens, axis=1)
+        double_stream_mod_img = jnp.concatenate([ref_mod_expanded, img_mod_expanded], axis=1)
+
+      double_block_caches = []
+      for idx, double_block in enumerate(self.double_blocks):
+        with jax.named_scope(f"double_block_{idx}"):
+          encoder_hidden_states, hidden_states, layer_kv = double_block(
+              hidden_states=hidden_states,
+              encoder_hidden_states=encoder_hidden_states,
+              temb=temb,
+              image_rotary_emb=concat_rotary_emb,
+              temb_mod_img=double_stream_mod_img,
+              temb_mod_txt=double_stream_mod_txt,
+              kv_cache=None,
+              kv_cache_mode="extract",
+              num_ref_tokens=num_ref_tokens,
+          )
+          double_block_caches.append(layer_kv)
+
+      num_txt_tokens = encoder_hidden_states.shape[1]
+      hidden_states = jnp.concatenate([encoder_hidden_states, hidden_states], axis=1)
+
+      txt_mod_expanded = jnp.repeat(jnp.expand_dims(single_stream_mod, 1), num_txt_tokens, axis=1)
+      ref_smod_expanded = jnp.repeat(jnp.expand_dims(ref_single_mod, 1), num_ref_tokens, axis=1)
+      img_smod_expanded = jnp.repeat(jnp.expand_dims(single_stream_mod, 1), num_img_tokens, axis=1)
+      single_stream_mod = jnp.concatenate([txt_mod_expanded, ref_smod_expanded, img_smod_expanded], axis=1)
+
+      single_block_caches = []
+      for idx, single_block in enumerate(self.single_blocks):
+        with jax.named_scope(f"single_block_{idx}"):
+          hidden_states, layer_kv = single_block(
+              hidden_states=hidden_states,
+              temb=temb,
+              image_rotary_emb=concat_rotary_emb,
+              temb_mod=single_stream_mod,
+              kv_cache=None,
+              kv_cache_mode="extract",
+              num_txt_tokens=num_txt_tokens,
+              num_ref_tokens=num_ref_tokens,
+          )
+          single_block_caches.append(layer_kv)
+
+      with jax.named_scope("output_norm_and_projection"):
+        hidden_states = hidden_states[:, num_txt_tokens + num_ref_tokens :, ...]
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+      extracted_kv_cache = {
+          "double": tuple(double_block_caches),
+          "single": tuple(single_block_caches),
+          "num_ref_tokens": num_ref_tokens,
+      }
+      if not return_dict:
+        return output, extracted_kv_cache
+      return Transformer2DModelOutput(sample=output), extracted_kv_cache
+
+    elif kv_cache_mode == "cached" and kv_cache is not None:
+      double_caches = kv_cache["double"] if isinstance(kv_cache, dict) else kv_cache[0]
+      single_caches = kv_cache["single"] if isinstance(kv_cache, dict) else kv_cache[1]
+      num_ref = kv_cache.get("num_ref_tokens", 0) if isinstance(kv_cache, dict) else 0
+
+      for idx, double_block in enumerate(self.double_blocks):
+        with jax.named_scope(f"double_block_{idx}"):
+          encoder_hidden_states, hidden_states, _ = double_block(
+              hidden_states=hidden_states,
+              encoder_hidden_states=encoder_hidden_states,
+              temb=temb,
+              image_rotary_emb=concat_rotary_emb,
+              temb_mod_img=double_stream_mod_img,
+              temb_mod_txt=double_stream_mod_txt,
+              kv_cache=double_caches[idx],
+              kv_cache_mode="cached",
+              num_ref_tokens=num_ref,
+          )
+
+      num_txt_tokens = encoder_hidden_states.shape[1]
+      hidden_states = jnp.concatenate([encoder_hidden_states, hidden_states], axis=1)
+
+      for idx, single_block in enumerate(self.single_blocks):
+        with jax.named_scope(f"single_block_{idx}"):
+          hidden_states, _ = single_block(
+              hidden_states=hidden_states,
+              temb=temb,
+              image_rotary_emb=concat_rotary_emb,
+              temb_mod=single_stream_mod,
+              kv_cache=single_caches[idx],
+              kv_cache_mode="cached",
+              num_txt_tokens=num_txt_tokens,
+              num_ref_tokens=num_ref,
+          )
+
+      with jax.named_scope("output_norm_and_projection"):
+        hidden_states = hidden_states[:, num_txt_tokens:, ...]
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+      if not return_dict:
+        return (output,)
+      return Transformer2DModelOutput(sample=output)
 
-    if not return_dict:
-      return (output,)
-    return Transformer2DModelOutput(sample=output)
+    else:
+      for idx, double_block in enumerate(self.double_blocks):
+        with jax.named_scope(f"double_block_{idx}"):
+          encoder_hidden_states, hidden_states, _ = double_block(
+              hidden_states=hidden_states,
+              encoder_hidden_states=encoder_hidden_states,
+              temb=temb,
+              image_rotary_emb=concat_rotary_emb,
+              temb_mod_img=double_stream_mod_img,
+              temb_mod_txt=double_stream_mod_txt,
+          )
+
+      num_txt_tokens = encoder_hidden_states.shape[1]
+      hidden_states = jnp.concatenate([encoder_hidden_states, hidden_states], axis=1)
+
+      for idx, single_block in enumerate(self.single_blocks):
+        with jax.named_scope(f"single_block_{idx}"):
+          hidden_states, _ = single_block(
+              hidden_states=hidden_states,
+              temb=temb,
+              image_rotary_emb=concat_rotary_emb,
+              temb_mod=single_stream_mod,
+          )
+
+      with jax.named_scope("output_norm_and_projection"):
+        hidden_states = hidden_states[:, num_txt_tokens:, ...]
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+      if not return_dict:
+        return (output,)
+      return Transformer2DModelOutput(sample=output)
diff --git a/src/maxdiffusion/models/flux/util.py b/src/maxdiffusion/models/flux/util.py
index aa43609a1..0855a5fb5 100644
--- a/src/maxdiffusion/models/flux/util.py
+++ b/src/maxdiffusion/models/flux/util.py
@@ -734,6 +734,43 @@ def set_val(var, tensor):
   return nnx.from_flat_state(flat_state)
 
 
+def patchify_latents(latents):
+  """Patchifies latents: (B, C, H, W) -> (B, C*4, H//2, W//2)."""
+  import jax.numpy as jnp
+
+  batch_size, num_channels, height, width = latents.shape
+  latents = latents.reshape((batch_size, num_channels, height // 2, 2, width // 2, 2))
+  latents = jnp.transpose(latents, (0, 1, 3, 5, 2, 4))
+  latents = latents.reshape((batch_size, num_channels * 4, height // 2, width // 2))
+  return latents
+
+
+def prepare_multi_image_ids(image_latents_list, scale=10):
+  """Generates 4D position IDs (T, H, W, L) for a sequence of reference image latents.
+
+  For the k-th image, T = scale * (k + 1).
+  image_latents_list: list of arrays with shape (1, C, H, W) or (C, H, W).
+  Returns: array of shape (1, total_tokens, 4).
+  """
+  import jax.numpy as jnp
+
+  all_ids = []
+  for idx, latent in enumerate(image_latents_list):
+    if latent.ndim == 4:
+      latent = latent[0]
+    _, h, w = latent.shape
+    t_val = scale * (idx + 1)
+    t = jnp.full((h * w, 1), t_val, dtype=jnp.int32)
+    h_grid, w_grid = jnp.meshgrid(jnp.arange(h, dtype=jnp.int32), jnp.arange(w, dtype=jnp.int32), indexing="ij")
+    h_coords = h_grid.reshape(-1, 1)
+    w_coords = w_grid.reshape(-1, 1)
+    l_coords = jnp.zeros((h * w, 1), dtype=jnp.int32)
+    coords = jnp.concatenate([t, h_coords, w_coords, l_coords], axis=-1)
+    all_ids.append(coords)
+  combined = jnp.concatenate(all_ids, axis=0)
+  return jnp.expand_dims(combined, axis=0)
+
+
 def load_and_convert_vae_weights(safetensors_path, jax_params, dtype=None, pt_state_dict=None):
   """Loads VAE weights from safetensors via zero-copy safetensors.numpy, maps them to JAX, and extracts BN stats."""
   from safetensors.numpy import load_file
@@ -756,62 +793,49 @@ def get_pytorch_weight_tensor(key, dtype_val=target_dtype):
     leaf_dtype = jnp.float32 if is_norm else dtype_val
     return jnp.array(tensor, dtype=leaf_dtype)
 
-  # Map weights
-  max_logging.log("Mapping VAE decoder weights to JAX parameters...")
-
-  # post_quant_conv
-  jax_params["post_quant_conv"]["kernel"] = jnp.array(
-      get_pytorch_weight_tensor("post_quant_conv.weight").transpose(2, 3, 1, 0)
-  )
-  jax_params["post_quant_conv"]["bias"] = jnp.array(get_pytorch_weight_tensor("post_quant_conv.bias"))
+  # 1. Map VAE Encoder Weights
+  if "encoder" in jax_params:
+    max_logging.log("Mapping VAE encoder weights to JAX parameters...")
+    enc_jax = jax_params["encoder"]
+
+    if "encoder.conv_in.weight" in pt_state_dict:
+      enc_jax["conv_in"]["kernel"] = jnp.array(get_pytorch_weight_tensor("encoder.conv_in.weight").transpose(2, 3, 1, 0))
+      enc_jax["conv_in"]["bias"] = jnp.array(get_pytorch_weight_tensor("encoder.conv_in.bias"))
+
+    for b_idx in range(4):
+      down_block_pt = f"encoder.down_blocks.{b_idx}"
+      down_block_jax = enc_jax[f"down_blocks_{b_idx}"]
+
+      for r_idx in range(2):
+        res_pt = f"{down_block_pt}.resnets.{r_idx}"
+        res_jax = down_block_jax[f"resnets_{r_idx}"]
+
+        res_jax["norm1"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm1.weight"))
+        res_jax["norm1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm1.bias"))
+        res_jax["conv1"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv1.weight").transpose(2, 3, 1, 0))
+        res_jax["conv1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv1.bias"))
+
+        res_jax["norm2"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm2.weight"))
+        res_jax["norm2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm2.bias"))
+        res_jax["conv2"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv2.weight").transpose(2, 3, 1, 0))
+        res_jax["conv2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv2.bias"))
+
+        shortcut_key = f"{res_pt}.conv_shortcut.weight"
+        if shortcut_key in pt_state_dict:
+          res_jax["conv_shortcut"]["kernel"] = jnp.array(get_pytorch_weight_tensor(shortcut_key).transpose(2, 3, 1, 0))
+          res_jax["conv_shortcut"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv_shortcut.bias"))
+
+      if b_idx < 3:
+        downsampler_pt = f"{down_block_pt}.downsamplers.0"
+        downsampler_jax = down_block_jax["downsamplers_0"]
+        downsampler_jax["conv"]["kernel"] = jnp.array(
+            get_pytorch_weight_tensor(f"{downsampler_pt}.conv.weight").transpose(2, 3, 1, 0)
+        )
+        downsampler_jax["conv"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{downsampler_pt}.conv.bias"))
 
-  # decoder.conv_in
-  jax_params["decoder"]["conv_in"]["kernel"] = jnp.array(
-      get_pytorch_weight_tensor("decoder.conv_in.weight").transpose(2, 3, 1, 0)
-  )
-  jax_params["decoder"]["conv_in"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_in.bias"))
-
-  # decoder.mid_block
-  # resnets
-  for idx in [0, 1]:
-    res_jax = jax_params["decoder"]["mid_block"][f"resnets_{idx}"]
-    res_pt_prefix = f"decoder.mid_block.resnets.{idx}"
-
-    res_jax["norm1"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm1.weight"))
-    res_jax["norm1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm1.bias"))
-    res_jax["conv1"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.conv1.weight").transpose(2, 3, 1, 0))
-    res_jax["conv1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.conv1.bias"))
-
-    res_jax["norm2"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm2.weight"))
-    res_jax["norm2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm2.bias"))
-    res_jax["conv2"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.conv2.weight").transpose(2, 3, 1, 0))
-    res_jax["conv2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.conv2.bias"))
-
-  # attentions
-  attn_pt_prefix = "decoder.mid_block.attentions.0"
-  attn_jax = jax_params["decoder"]["mid_block"]["attentions_0"]
-
-  attn_jax["group_norm"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.group_norm.weight"))
-  attn_jax["group_norm"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.group_norm.bias"))
-
-  attn_jax["query"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_q.weight").T)
-  attn_jax["query"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_q.bias"))
-  attn_jax["key"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_k.weight").T)
-  attn_jax["key"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_k.bias"))
-  attn_jax["value"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_v.weight").T)
-  attn_jax["value"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_v.bias"))
-
-  attn_jax["proj_attn"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_out.0.weight").T)
-  attn_jax["proj_attn"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_out.0.bias"))
-
-  # decoder.up_blocks
-  for b_idx in range(4):
-    up_block_jax = jax_params["decoder"][f"up_blocks_{b_idx}"]
-    up_block_pt = f"decoder.up_blocks.{b_idx}"
-
-    for r_idx in range(3):
-      res_jax = up_block_jax[f"resnets_{r_idx}"]
-      res_pt = f"{up_block_pt}.resnets.{r_idx}"
+    for r_idx in range(2):
+      res_pt = f"encoder.mid_block.resnets.{r_idx}"
+      res_jax = enc_jax["mid_block"][f"resnets_{r_idx}"]
 
       res_jax["norm1"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm1.weight"))
       res_jax["norm1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm1.bias"))
@@ -823,27 +847,112 @@ def get_pytorch_weight_tensor(key, dtype_val=target_dtype):
       res_jax["conv2"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv2.weight").transpose(2, 3, 1, 0))
       res_jax["conv2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv2.bias"))
 
-      shortcut_key = f"{res_pt}.conv_shortcut.weight"
-      if shortcut_key in pt_state_dict:
-        res_jax["conv_shortcut"]["kernel"] = jnp.array(get_pytorch_weight_tensor(shortcut_key).transpose(2, 3, 1, 0))
-        res_jax["conv_shortcut"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv_shortcut.bias"))
+    attn_enc_pt = "encoder.mid_block.attentions.0"
+    attn_enc_jax = enc_jax["mid_block"]["attentions_0"]
+    attn_enc_jax["group_norm"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.group_norm.weight"))
+    attn_enc_jax["group_norm"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.group_norm.bias"))
+    attn_enc_jax["query"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.to_q.weight").T)
+    attn_enc_jax["query"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.to_q.bias"))
+    attn_enc_jax["key"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.to_k.weight").T)
+    attn_enc_jax["key"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.to_k.bias"))
+    attn_enc_jax["value"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.to_v.weight").T)
+    attn_enc_jax["value"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.to_v.bias"))
+    attn_enc_jax["proj_attn"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.to_out.0.weight").T)
+    attn_enc_jax["proj_attn"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_enc_pt}.to_out.0.bias"))
+
+    enc_jax["conv_norm_out"]["scale"] = jnp.array(get_pytorch_weight_tensor("encoder.conv_norm_out.weight"))
+    enc_jax["conv_norm_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("encoder.conv_norm_out.bias"))
+    enc_jax["conv_out"]["kernel"] = jnp.array(get_pytorch_weight_tensor("encoder.conv_out.weight").transpose(2, 3, 1, 0))
+    enc_jax["conv_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("encoder.conv_out.bias"))
+
+  if "quant_conv" in jax_params and "quant_conv.weight" in pt_state_dict:
+    jax_params["quant_conv"]["kernel"] = jnp.array(get_pytorch_weight_tensor("quant_conv.weight").transpose(2, 3, 1, 0))
+    jax_params["quant_conv"]["bias"] = jnp.array(get_pytorch_weight_tensor("quant_conv.bias"))
+
+  # 2. Map VAE Decoder Weights
+  max_logging.log("Mapping VAE decoder weights to JAX parameters...")
 
-    if b_idx < 3:
-      upsampler_jax = up_block_jax["upsamplers_0"]
-      upsampler_pt = f"{up_block_pt}.upsamplers.0"
+  if "post_quant_conv" in jax_params and "post_quant_conv.weight" in pt_state_dict:
+    jax_params["post_quant_conv"]["kernel"] = jnp.array(
+        get_pytorch_weight_tensor("post_quant_conv.weight").transpose(2, 3, 1, 0)
+    )
+    jax_params["post_quant_conv"]["bias"] = jnp.array(get_pytorch_weight_tensor("post_quant_conv.bias"))
 
-      upsampler_jax["conv"]["kernel"] = jnp.array(
-          get_pytorch_weight_tensor(f"{upsampler_pt}.conv.weight").transpose(2, 3, 1, 0)
+  if "decoder" in jax_params:
+    dec_jax = jax_params["decoder"]
+    dec_jax["conv_in"]["kernel"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_in.weight").transpose(2, 3, 1, 0))
+    dec_jax["conv_in"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_in.bias"))
+
+    for idx in [0, 1]:
+      res_jax = dec_jax["mid_block"][f"resnets_{idx}"]
+      res_pt_prefix = f"decoder.mid_block.resnets.{idx}"
+
+      res_jax["norm1"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm1.weight"))
+      res_jax["norm1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm1.bias"))
+      res_jax["conv1"]["kernel"] = jnp.array(
+          get_pytorch_weight_tensor(f"{res_pt_prefix}.conv1.weight").transpose(2, 3, 1, 0)
       )
-      upsampler_jax["conv"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{upsampler_pt}.conv.bias"))
+      res_jax["conv1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.conv1.bias"))
 
-  # decoder.conv_norm_out & conv_out
-  jax_params["decoder"]["conv_norm_out"]["scale"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_norm_out.weight"))
-  jax_params["decoder"]["conv_norm_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_norm_out.bias"))
-  jax_params["decoder"]["conv_out"]["kernel"] = jnp.array(
-      get_pytorch_weight_tensor("decoder.conv_out.weight").transpose(2, 3, 1, 0)
-  )
-  jax_params["decoder"]["conv_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_out.bias"))
+      res_jax["norm2"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm2.weight"))
+      res_jax["norm2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm2.bias"))
+      res_jax["conv2"]["kernel"] = jnp.array(
+          get_pytorch_weight_tensor(f"{res_pt_prefix}.conv2.weight").transpose(2, 3, 1, 0)
+      )
+      res_jax["conv2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.conv2.bias"))
+
+    attn_pt_prefix = "decoder.mid_block.attentions.0"
+    attn_jax = dec_jax["mid_block"]["attentions_0"]
+
+    attn_jax["group_norm"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.group_norm.weight"))
+    attn_jax["group_norm"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.group_norm.bias"))
+
+    attn_jax["query"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_q.weight").T)
+    attn_jax["query"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_q.bias"))
+    attn_jax["key"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_k.weight").T)
+    attn_jax["key"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_k.bias"))
+    attn_jax["value"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_v.weight").T)
+    attn_jax["value"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_v.bias"))
+
+    attn_jax["proj_attn"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_out.0.weight").T)
+    attn_jax["proj_attn"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_out.0.bias"))
+
+    for b_idx in range(4):
+      up_block_jax = dec_jax[f"up_blocks_{b_idx}"]
+      up_block_pt = f"decoder.up_blocks.{b_idx}"
+
+      for r_idx in range(3):
+        res_jax = up_block_jax[f"resnets_{r_idx}"]
+        res_pt = f"{up_block_pt}.resnets.{r_idx}"
+
+        res_jax["norm1"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm1.weight"))
+        res_jax["norm1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm1.bias"))
+        res_jax["conv1"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv1.weight").transpose(2, 3, 1, 0))
+        res_jax["conv1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv1.bias"))
+
+        res_jax["norm2"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm2.weight"))
+        res_jax["norm2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm2.bias"))
+        res_jax["conv2"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv2.weight").transpose(2, 3, 1, 0))
+        res_jax["conv2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv2.bias"))
+
+        shortcut_key = f"{res_pt}.conv_shortcut.weight"
+        if shortcut_key in pt_state_dict:
+          res_jax["conv_shortcut"]["kernel"] = jnp.array(get_pytorch_weight_tensor(shortcut_key).transpose(2, 3, 1, 0))
+          res_jax["conv_shortcut"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv_shortcut.bias"))
+
+      if b_idx < 3:
+        upsampler_jax = up_block_jax["upsamplers_0"]
+        upsampler_pt = f"{up_block_pt}.upsamplers.0"
+
+        upsampler_jax["conv"]["kernel"] = jnp.array(
+            get_pytorch_weight_tensor(f"{upsampler_pt}.conv.weight").transpose(2, 3, 1, 0)
+        )
+        upsampler_jax["conv"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{upsampler_pt}.conv.bias"))
+
+    dec_jax["conv_norm_out"]["scale"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_norm_out.weight"))
+    dec_jax["conv_norm_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_norm_out.bias"))
+    dec_jax["conv_out"]["kernel"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_out.weight").transpose(2, 3, 1, 0))
+    dec_jax["conv_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_out.bias"))
 
   jax_params = jax.tree_util.tree_map(
       lambda leaf: jnp.zeros(leaf.shape, dtype=leaf.dtype) if isinstance(leaf, jax.ShapeDtypeStruct) else leaf, jax_params
diff --git a/src/maxdiffusion/models/flux/vae/__init__.py b/src/maxdiffusion/models/flux/vae/__init__.py
new file mode 100644
index 000000000..11f31009e
--- /dev/null
+++ b/src/maxdiffusion/models/flux/vae/__init__.py
@@ -0,0 +1,15 @@
+"""
+Copyright 2026 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
diff --git a/src/maxdiffusion/models/flux/vae/autoencoder_kl_flux2_nnx.py b/src/maxdiffusion/models/flux/vae/autoencoder_kl_flux2_nnx.py
new file mode 100644
index 000000000..01adabd1f
--- /dev/null
+++ b/src/maxdiffusion/models/flux/vae/autoencoder_kl_flux2_nnx.py
@@ -0,0 +1,799 @@
+"""
+Copyright 2026 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import math
+from typing import Optional, Tuple
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+
+
+class NNXUpsample2D(nnx.Module):
+  """2D Nearest-neighbor Upsample + Conv layer in NNX."""
+
+  def __init__(
+      self,
+      in_channels: int,
+      out_channels: Optional[int] = None,
+      rngs: Optional[nnx.Rngs] = None,
+      dtype: jnp.dtype = jnp.float32,
+      param_dtype: jnp.dtype = jnp.float32,
+  ):
+    out_channels = out_channels or in_channels
+    self.conv = nnx.Conv(
+        in_features=in_channels,
+        out_features=out_channels,
+        kernel_size=(3, 3),
+        strides=(1, 1),
+        padding=((1, 1), (1, 1)),
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+
+  def __call__(self, x: jax.Array) -> jax.Array:
+    batch, height, width, channels = x.shape
+    x = jnp.broadcast_to(x[:, :, None, :, None, :], (batch, height, 2, width, 2, channels))
+    x = jnp.reshape(x, (batch, height * 2, width * 2, channels))
+    return self.conv(x)
+
+
+class NNXDownsample2D(nnx.Module):
+  """2D Downsample layer with asymmetric padding in NNX."""
+
+  def __init__(
+      self,
+      in_channels: int,
+      out_channels: Optional[int] = None,
+      rngs: Optional[nnx.Rngs] = None,
+      dtype: jnp.dtype = jnp.float32,
+      param_dtype: jnp.dtype = jnp.float32,
+  ):
+    out_channels = out_channels or in_channels
+    self.conv = nnx.Conv(
+        in_features=in_channels,
+        out_features=out_channels,
+        kernel_size=(3, 3),
+        strides=(2, 2),
+        padding="VALID",
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+
+  def __call__(self, x: jax.Array) -> jax.Array:
+    pad_width = ((0, 0), (0, 1), (0, 1), (0, 0))
+    x = jnp.pad(x, pad_width)
+    return self.conv(x)
+
+
+class NNXResnetBlock2D(nnx.Module):
+  """2D ResNet Block with GroupNorm and SiLU activations in NNX."""
+
+  def __init__(
+      self,
+      in_channels: int,
+      out_channels: Optional[int] = None,
+      groups: int = 32,
+      use_conv_shortcut: Optional[bool] = None,
+      rngs: Optional[nnx.Rngs] = None,
+      dtype: jnp.dtype = jnp.float32,
+      param_dtype: jnp.dtype = jnp.float32,
+  ):
+    out_channels = out_channels or in_channels
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+
+    self.norm1 = nnx.GroupNorm(
+        num_groups=groups,
+        num_features=in_channels,
+        epsilon=1e-6,
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+    self.conv1 = nnx.Conv(
+        in_features=in_channels,
+        out_features=out_channels,
+        kernel_size=(3, 3),
+        strides=(1, 1),
+        padding=((1, 1), (1, 1)),
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+    self.norm2 = nnx.GroupNorm(
+        num_groups=groups,
+        num_features=out_channels,
+        epsilon=1e-6,
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+    self.conv2 = nnx.Conv(
+        in_features=out_channels,
+        out_features=out_channels,
+        kernel_size=(3, 3),
+        strides=(1, 1),
+        padding=((1, 1), (1, 1)),
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+
+    use_shortcut = (in_channels != out_channels) if use_conv_shortcut is None else use_conv_shortcut
+    if use_shortcut:
+      self.conv_shortcut = nnx.Conv(
+          in_features=in_channels,
+          out_features=out_channels,
+          kernel_size=(1, 1),
+          strides=(1, 1),
+          padding="VALID",
+          rngs=rngs,
+          dtype=dtype,
+          param_dtype=param_dtype,
+      )
+    else:
+      self.conv_shortcut = None
+
+  def __call__(self, x: jax.Array) -> jax.Array:
+    residual = self.conv_shortcut(x) if self.conv_shortcut is not None else x
+    h = self.norm1(x)
+    h = nnx.silu(h)
+    h = self.conv1(h)
+    h = self.norm2(h)
+    h = nnx.silu(h)
+    h = self.conv2(h)
+    return h + residual
+
+
+class NNXAttentionBlock(nnx.Module):
+  """Self-Attention block with GroupNorm in NNX."""
+
+  def __init__(
+      self,
+      channels: int,
+      groups: int = 32,
+      rngs: Optional[nnx.Rngs] = None,
+      dtype: jnp.dtype = jnp.float32,
+      param_dtype: jnp.dtype = jnp.float32,
+  ):
+    self.channels = channels
+    self.group_norm = nnx.GroupNorm(
+        num_groups=groups,
+        num_features=channels,
+        epsilon=1e-6,
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+    self.to_q = nnx.Linear(
+        in_features=channels,
+        out_features=channels,
+        use_bias=True,
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+    self.to_k = nnx.Linear(
+        in_features=channels,
+        out_features=channels,
+        use_bias=True,
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+    self.to_v = nnx.Linear(
+        in_features=channels,
+        out_features=channels,
+        use_bias=True,
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+    self.to_out = nnx.Linear(
+        in_features=channels,
+        out_features=channels,
+        use_bias=True,
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+
+  def __call__(self, x: jax.Array) -> jax.Array:
+    residual = x
+    b, h, w, c = x.shape
+    h_states = self.group_norm(x)
+    h_flat = h_states.reshape((b, h * w, c))
+
+    q = self.to_q(h_flat)
+    k = self.to_k(h_flat)
+    v = self.to_v(h_flat)
+
+    scale = 1.0 / math.sqrt(c)
+    attn_weights = jnp.einsum("bqc,bkc->bqk", q * scale, k)
+    attn_weights = jax.nn.softmax(attn_weights, axis=-1)
+
+    out = jnp.einsum("bqk,bkc->bqc", attn_weights, v)
+    out = self.to_out(out)
+    out = out.reshape((b, h, w, c))
+    return out + residual
+
+
+class NNXUNetMidBlock2D(nnx.Module):
+  """Mid-Block module in NNX with resnets and attention."""
+
+  def __init__(
+      self,
+      in_channels: int,
+      groups: int = 32,
+      rngs: Optional[nnx.Rngs] = None,
+      dtype: jnp.dtype = jnp.float32,
+      param_dtype: jnp.dtype = jnp.float32,
+  ):
+    self.resnets_0 = NNXResnetBlock2D(
+        in_channels=in_channels,
+        out_channels=in_channels,
+        groups=groups,
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+    self.attentions_0 = NNXAttentionBlock(
+        channels=in_channels,
+        groups=groups,
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+    self.resnets_1 = NNXResnetBlock2D(
+        in_channels=in_channels,
+        out_channels=in_channels,
+        groups=groups,
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+
+  def __call__(self, x: jax.Array) -> jax.Array:
+    x = self.resnets_0(x)
+    x = self.attentions_0(x)
+    x = self.resnets_1(x)
+    return x
+
+
+class NNXDownEncoderBlock2D(nnx.Module):
+  """Down-Encoder block containing ResNet layers and an optional Downsampler in NNX."""
+
+  def __init__(
+      self,
+      in_channels: int,
+      out_channels: int,
+      num_layers: int = 2,
+      groups: int = 32,
+      add_downsample: bool = True,
+      rngs: Optional[nnx.Rngs] = None,
+      dtype: jnp.dtype = jnp.float32,
+      param_dtype: jnp.dtype = jnp.float32,
+  ):
+    resnets = []
+    for i in range(num_layers):
+      in_ch = in_channels if i == 0 else out_channels
+      resnets.append(
+          NNXResnetBlock2D(
+              in_channels=in_ch,
+              out_channels=out_channels,
+              groups=groups,
+              rngs=rngs,
+              dtype=dtype,
+              param_dtype=param_dtype,
+          )
+      )
+    self.resnets = nnx.List(resnets)
+
+    if add_downsample:
+      self.downsamplers_0 = NNXDownsample2D(
+          in_channels=out_channels,
+          out_channels=out_channels,
+          rngs=rngs,
+          dtype=dtype,
+          param_dtype=param_dtype,
+      )
+    else:
+      self.downsamplers_0 = None
+
+  def __call__(self, x: jax.Array) -> jax.Array:
+    for resnet in self.resnets:
+      x = resnet(x)
+    if self.downsamplers_0 is not None:
+      x = self.downsamplers_0(x)
+    return x
+
+
+class NNXUpDecoderBlock2D(nnx.Module):
+  """Up-Decoder block containing ResNet layers and an optional Upsampler in NNX."""
+
+  def __init__(
+      self,
+      in_channels: int,
+      out_channels: int,
+      num_layers: int = 3,
+      groups: int = 32,
+      add_upsample: bool = True,
+      rngs: Optional[nnx.Rngs] = None,
+      dtype: jnp.dtype = jnp.float32,
+      param_dtype: jnp.dtype = jnp.float32,
+  ):
+    resnets = []
+    for i in range(num_layers):
+      in_ch = in_channels if i == 0 else out_channels
+      resnets.append(
+          NNXResnetBlock2D(
+              in_channels=in_ch,
+              out_channels=out_channels,
+              groups=groups,
+              rngs=rngs,
+              dtype=dtype,
+              param_dtype=param_dtype,
+          )
+      )
+    self.resnets = nnx.List(resnets)
+
+    if add_upsample:
+      self.upsamplers_0 = NNXUpsample2D(
+          in_channels=out_channels,
+          out_channels=out_channels,
+          rngs=rngs,
+          dtype=dtype,
+          param_dtype=param_dtype,
+      )
+    else:
+      self.upsamplers_0 = None
+
+  def __call__(self, x: jax.Array) -> jax.Array:
+    for resnet in self.resnets:
+      x = resnet(x)
+    if self.upsamplers_0 is not None:
+      x = self.upsamplers_0(x)
+    return x
+
+
+class NNXEncoder(nnx.Module):
+  """Complete VAE Encoder in NNX."""
+
+  def __init__(
+      self,
+      in_channels: int = 3,
+      out_channels: int = 32,
+      block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+      layers_per_block: int = 2,
+      norm_num_groups: int = 32,
+      rngs: Optional[nnx.Rngs] = None,
+      dtype: jnp.dtype = jnp.float32,
+      param_dtype: jnp.dtype = jnp.float32,
+  ):
+    self.conv_in = nnx.Conv(
+        in_features=in_channels,
+        out_features=block_out_channels[0],
+        kernel_size=(3, 3),
+        strides=(1, 1),
+        padding=((1, 1), (1, 1)),
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+
+    down_blocks = []
+    output_ch = block_out_channels[0]
+    for i, ch in enumerate(block_out_channels):
+      input_ch = output_ch
+      output_ch = ch
+      is_final = i == len(block_out_channels) - 1
+      down_blocks.append(
+          NNXDownEncoderBlock2D(
+              in_channels=input_ch,
+              out_channels=output_ch,
+              num_layers=layers_per_block,
+              groups=norm_num_groups,
+              add_downsample=not is_final,
+              rngs=rngs,
+              dtype=dtype,
+              param_dtype=param_dtype,
+          )
+      )
+    self.down_blocks = nnx.List(down_blocks)
+
+    self.mid_block = NNXUNetMidBlock2D(
+        in_channels=block_out_channels[-1],
+        groups=norm_num_groups,
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+
+    self.conv_norm_out = nnx.GroupNorm(
+        num_groups=norm_num_groups,
+        num_features=block_out_channels[-1],
+        epsilon=1e-6,
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+    self.conv_out = nnx.Conv(
+        in_features=block_out_channels[-1],
+        out_features=2 * out_channels,  # double_z for Gaussian distribution moments
+        kernel_size=(3, 3),
+        strides=(1, 1),
+        padding=((1, 1), (1, 1)),
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+
+  def __call__(self, x: jax.Array) -> jax.Array:
+    x = self.conv_in(x)
+    for block in self.down_blocks:
+      x = block(x)
+    x = self.mid_block(x)
+    x = self.conv_norm_out(x)
+    x = nnx.silu(x)
+    x = self.conv_out(x)
+    return x
+
+
+class NNXDecoder(nnx.Module):
+  """Complete VAE Decoder in NNX."""
+
+  def __init__(
+      self,
+      in_channels: int = 32,
+      out_channels: int = 3,
+      block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+      layers_per_block: int = 3,
+      norm_num_groups: int = 32,
+      rngs: Optional[nnx.Rngs] = None,
+      dtype: jnp.dtype = jnp.float32,
+      param_dtype: jnp.dtype = jnp.float32,
+  ):
+    reversed_channels = list(reversed(block_out_channels))
+    self.conv_in = nnx.Conv(
+        in_features=in_channels,
+        out_features=reversed_channels[0],
+        kernel_size=(3, 3),
+        strides=(1, 1),
+        padding=((1, 1), (1, 1)),
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+
+    self.mid_block = NNXUNetMidBlock2D(
+        in_channels=reversed_channels[0],
+        groups=norm_num_groups,
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+
+    up_blocks = []
+    output_ch = reversed_channels[0]
+    for i, ch in enumerate(reversed_channels):
+      input_ch = output_ch
+      output_ch = ch
+      is_final = i == len(reversed_channels) - 1
+      up_blocks.append(
+          NNXUpDecoderBlock2D(
+              in_channels=input_ch,
+              out_channels=output_ch,
+              num_layers=layers_per_block,
+              groups=norm_num_groups,
+              add_upsample=not is_final,
+              rngs=rngs,
+              dtype=dtype,
+              param_dtype=param_dtype,
+          )
+      )
+    self.up_blocks = nnx.List(up_blocks)
+
+    self.conv_norm_out = nnx.GroupNorm(
+        num_groups=norm_num_groups,
+        num_features=reversed_channels[-1],
+        epsilon=1e-6,
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+    self.conv_out = nnx.Conv(
+        in_features=reversed_channels[-1],
+        out_features=out_channels,
+        kernel_size=(3, 3),
+        strides=(1, 1),
+        padding=((1, 1), (1, 1)),
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+
+  def __call__(self, x: jax.Array) -> jax.Array:
+    x = self.conv_in(x)
+    x = self.mid_block(x)
+    for block in self.up_blocks:
+      x = block(x)
+    x = self.conv_norm_out(x)
+    x = nnx.silu(x)
+    x = self.conv_out(x)
+    return x
+
+
+class NNXAutoencoderKLFlux2(nnx.Module):
+  """Full FLUX.2-Klein Variational Autoencoder (VAE) in Flax NNX."""
+
+  def __init__(
+      self,
+      in_channels: int = 3,
+      out_channels: int = 3,
+      latent_channels: int = 32,
+      block_out_channels: Tuple[int, ...] = (128, 256, 512, 512),
+      layers_per_block: int = 2,
+      norm_num_groups: int = 32,
+      rngs: Optional[nnx.Rngs] = None,
+      dtype: jnp.dtype = jnp.float32,
+      param_dtype: jnp.dtype = jnp.float32,
+  ):
+    rngs = rngs or nnx.Rngs(0)
+    self.latent_channels = latent_channels
+    self.dtype = dtype
+
+    self.encoder = NNXEncoder(
+        in_channels=in_channels,
+        out_channels=latent_channels,
+        block_out_channels=block_out_channels,
+        layers_per_block=layers_per_block,
+        norm_num_groups=norm_num_groups,
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+    self.quant_conv = nnx.Conv(
+        in_features=2 * latent_channels,
+        out_features=2 * latent_channels,
+        kernel_size=(1, 1),
+        strides=(1, 1),
+        padding="VALID",
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+    self.post_quant_conv = nnx.Conv(
+        in_features=latent_channels,
+        out_features=latent_channels,
+        kernel_size=(1, 1),
+        strides=(1, 1),
+        padding="VALID",
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+    self.decoder = NNXDecoder(
+        in_channels=latent_channels,
+        out_channels=out_channels,
+        block_out_channels=block_out_channels,
+        layers_per_block=layers_per_block + 1,  # 3 resnet blocks in decoder
+        norm_num_groups=norm_num_groups,
+        rngs=rngs,
+        dtype=dtype,
+        param_dtype=param_dtype,
+    )
+
+  def encode(self, sample: jax.Array) -> jax.Array:
+    """Encodes image tensor of shape (B, 3, H, W) to mode latents of shape (B, 32, H/8, W/8)."""
+    # Transpose to channels last (B, H, W, 3)
+    x = jnp.transpose(sample, (0, 2, 3, 1))
+    h = self.encoder(x)
+    moments = self.quant_conv(h)
+    # Extract mean / mode (first latent_channels)
+    mean, _ = jnp.split(moments, 2, axis=-1)  # (B, H/8, W/8, 32)
+    # Transpose back to (B, 32, H/8, W/8)
+    return jnp.transpose(mean, (0, 3, 1, 2))
+
+  def decode(self, latents: jax.Array) -> jax.Array:
+    """Decodes latent tensor of shape (B, 32, H/8, W/8) to image tensor of shape (B, 3, H, W)."""
+    # Transpose to channels last (B, H/8, W/8, 32)
+    z = jnp.transpose(latents, (0, 2, 3, 1))
+    h = self.post_quant_conv(z)
+    img = self.decoder(h)
+    # Transpose back to (B, 3, H, W)
+    return jnp.transpose(img, (0, 3, 1, 2))
+
+
+def load_and_convert_flux2klein_nnx_vae_weights(
+    safetensors_path: str,
+    nnx_vae: NNXAutoencoderKLFlux2,
+    dtype: Optional[jnp.dtype] = None,
+    pt_state_dict: Optional[dict] = None,
+):
+  """Directly loads and maps PyTorch safetensors into NNXAutoencoderKLFlux2 State."""
+  from safetensors.numpy import load_file
+
+  if pt_state_dict is None:
+    pt_state_dict = load_file(safetensors_path)
+
+  target_dtype = dtype if dtype is not None else jnp.float32
+
+  def get_pt_tensor(key, is_norm=False):
+    tensor = pt_state_dict[key]
+    leaf_dtype = jnp.float32 if is_norm else target_dtype
+    return jnp.array(tensor, dtype=leaf_dtype)
+
+  def get_conv_kernel(key):
+    return jnp.array(pt_state_dict[key].transpose(2, 3, 1, 0), dtype=target_dtype)
+
+  def get_linear_kernel(key):
+    return jnp.array(pt_state_dict[key].T, dtype=target_dtype)
+
+  flat_state = dict(nnx.to_flat_state(nnx.state(nnx_vae, nnx.Param)))
+
+  def set_val(var, val):
+    var[...] = val
+
+  # =========================================================================
+  # 1. ENCODER
+  # =========================================================================
+  set_val(flat_state[("encoder", "conv_in", "kernel")], get_conv_kernel("encoder.conv_in.weight"))
+  set_val(flat_state[("encoder", "conv_in", "bias")], get_pt_tensor("encoder.conv_in.bias"))
+
+  for b_idx in range(4):
+    down_block_pt = f"encoder.down_blocks.{b_idx}"
+    for r_idx in range(2):
+      res_pt = f"{down_block_pt}.resnets.{r_idx}"
+      res_path = ("encoder", "down_blocks", b_idx, "resnets", r_idx)
+
+      set_val(flat_state[res_path + ("norm1", "scale")], get_pt_tensor(f"{res_pt}.norm1.weight", is_norm=True))
+      set_val(flat_state[res_path + ("norm1", "bias")], get_pt_tensor(f"{res_pt}.norm1.bias", is_norm=True))
+      set_val(flat_state[res_path + ("conv1", "kernel")], get_conv_kernel(f"{res_pt}.conv1.weight"))
+      set_val(flat_state[res_path + ("conv1", "bias")], get_pt_tensor(f"{res_pt}.conv1.bias"))
+
+      set_val(flat_state[res_path + ("norm2", "scale")], get_pt_tensor(f"{res_pt}.norm2.weight", is_norm=True))
+      set_val(flat_state[res_path + ("norm2", "bias")], get_pt_tensor(f"{res_pt}.norm2.bias", is_norm=True))
+      set_val(flat_state[res_path + ("conv2", "kernel")], get_conv_kernel(f"{res_pt}.conv2.weight"))
+      set_val(flat_state[res_path + ("conv2", "bias")], get_pt_tensor(f"{res_pt}.conv2.bias"))
+
+      shortcut_key = f"{res_pt}.conv_shortcut.weight"
+      if shortcut_key in pt_state_dict:
+        set_val(flat_state[res_path + ("conv_shortcut", "kernel")], get_conv_kernel(shortcut_key))
+        set_val(flat_state[res_path + ("conv_shortcut", "bias")], get_pt_tensor(f"{res_pt}.conv_shortcut.bias"))
+
+    if b_idx < 3:
+      ds_pt = f"{down_block_pt}.downsamplers.0.conv"
+      ds_path = ("encoder", "down_blocks", b_idx, "downsamplers_0", "conv")
+      set_val(flat_state[ds_path + ("kernel",)], get_conv_kernel(f"{ds_pt}.weight"))
+      set_val(flat_state[ds_path + ("bias",)], get_pt_tensor(f"{ds_pt}.bias"))
+
+  # Encoder Mid Block
+  for r_idx in [0, 1]:
+    res_pt = f"encoder.mid_block.resnets.{r_idx}"
+    res_path = ("encoder", "mid_block", f"resnets_{r_idx}")
+    set_val(flat_state[res_path + ("norm1", "scale")], get_pt_tensor(f"{res_pt}.norm1.weight", is_norm=True))
+    set_val(flat_state[res_path + ("norm1", "bias")], get_pt_tensor(f"{res_pt}.norm1.bias", is_norm=True))
+    set_val(flat_state[res_path + ("conv1", "kernel")], get_conv_kernel(f"{res_pt}.conv1.weight"))
+    set_val(flat_state[res_path + ("conv1", "bias")], get_pt_tensor(f"{res_pt}.conv1.bias"))
+    set_val(flat_state[res_path + ("norm2", "scale")], get_pt_tensor(f"{res_pt}.norm2.weight", is_norm=True))
+    set_val(flat_state[res_path + ("norm2", "bias")], get_pt_tensor(f"{res_pt}.norm2.bias", is_norm=True))
+    set_val(flat_state[res_path + ("conv2", "kernel")], get_conv_kernel(f"{res_pt}.conv2.weight"))
+    set_val(flat_state[res_path + ("conv2", "bias")], get_pt_tensor(f"{res_pt}.conv2.bias"))
+
+  attn_pt = "encoder.mid_block.attentions.0"
+  attn_path = ("encoder", "mid_block", "attentions_0")
+  set_val(flat_state[attn_path + ("group_norm", "scale")], get_pt_tensor(f"{attn_pt}.group_norm.weight", is_norm=True))
+  set_val(flat_state[attn_path + ("group_norm", "bias")], get_pt_tensor(f"{attn_pt}.group_norm.bias", is_norm=True))
+  set_val(flat_state[attn_path + ("to_q", "kernel")], get_linear_kernel(f"{attn_pt}.to_q.weight"))
+  set_val(flat_state[attn_path + ("to_q", "bias")], get_pt_tensor(f"{attn_pt}.to_q.bias"))
+  set_val(flat_state[attn_path + ("to_k", "kernel")], get_linear_kernel(f"{attn_pt}.to_k.weight"))
+  set_val(flat_state[attn_path + ("to_k", "bias")], get_pt_tensor(f"{attn_pt}.to_k.bias"))
+  set_val(flat_state[attn_path + ("to_v", "kernel")], get_linear_kernel(f"{attn_pt}.to_v.weight"))
+  set_val(flat_state[attn_path + ("to_v", "bias")], get_pt_tensor(f"{attn_pt}.to_v.bias"))
+  set_val(flat_state[attn_path + ("to_out", "kernel")], get_linear_kernel(f"{attn_pt}.to_out.0.weight"))
+  set_val(flat_state[attn_path + ("to_out", "bias")], get_pt_tensor(f"{attn_pt}.to_out.0.bias"))
+
+  set_val(flat_state[("encoder", "conv_norm_out", "scale")], get_pt_tensor("encoder.conv_norm_out.weight", is_norm=True))
+  set_val(flat_state[("encoder", "conv_norm_out", "bias")], get_pt_tensor("encoder.conv_norm_out.bias", is_norm=True))
+  set_val(flat_state[("encoder", "conv_out", "kernel")], get_conv_kernel("encoder.conv_out.weight"))
+  set_val(flat_state[("encoder", "conv_out", "bias")], get_pt_tensor("encoder.conv_out.bias"))
+
+  # =========================================================================
+  # 2. QUANT CONV & POST QUANT CONV
+  # =========================================================================
+  set_val(flat_state[("quant_conv", "kernel")], get_conv_kernel("quant_conv.weight"))
+  set_val(flat_state[("quant_conv", "bias")], get_pt_tensor("quant_conv.bias"))
+  set_val(flat_state[("post_quant_conv", "kernel")], get_conv_kernel("post_quant_conv.weight"))
+  set_val(flat_state[("post_quant_conv", "bias")], get_pt_tensor("post_quant_conv.bias"))
+
+  # =========================================================================
+  # 3. DECODER
+  # =========================================================================
+  set_val(flat_state[("decoder", "conv_in", "kernel")], get_conv_kernel("decoder.conv_in.weight"))
+  set_val(flat_state[("decoder", "conv_in", "bias")], get_pt_tensor("decoder.conv_in.bias"))
+
+  for r_idx in [0, 1]:
+    res_pt = f"decoder.mid_block.resnets.{r_idx}"
+    res_path = ("decoder", "mid_block", f"resnets_{r_idx}")
+    set_val(flat_state[res_path + ("norm1", "scale")], get_pt_tensor(f"{res_pt}.norm1.weight", is_norm=True))
+    set_val(flat_state[res_path + ("norm1", "bias")], get_pt_tensor(f"{res_pt}.norm1.bias", is_norm=True))
+    set_val(flat_state[res_path + ("conv1", "kernel")], get_conv_kernel(f"{res_pt}.conv1.weight"))
+    set_val(flat_state[res_path + ("conv1", "bias")], get_pt_tensor(f"{res_pt}.conv1.bias"))
+    set_val(flat_state[res_path + ("norm2", "scale")], get_pt_tensor(f"{res_pt}.norm2.weight", is_norm=True))
+    set_val(flat_state[res_path + ("norm2", "bias")], get_pt_tensor(f"{res_pt}.norm2.bias", is_norm=True))
+    set_val(flat_state[res_path + ("conv2", "kernel")], get_conv_kernel(f"{res_pt}.conv2.weight"))
+    set_val(flat_state[res_path + ("conv2", "bias")], get_pt_tensor(f"{res_pt}.conv2.bias"))
+
+  dec_attn_pt = "decoder.mid_block.attentions.0"
+  dec_attn_path = ("decoder", "mid_block", "attentions_0")
+  set_val(
+      flat_state[dec_attn_path + ("group_norm", "scale")], get_pt_tensor(f"{dec_attn_pt}.group_norm.weight", is_norm=True)
+  )
+  set_val(flat_state[dec_attn_path + ("group_norm", "bias")], get_pt_tensor(f"{dec_attn_pt}.group_norm.bias", is_norm=True))
+  set_val(flat_state[dec_attn_path + ("to_q", "kernel")], get_linear_kernel(f"{dec_attn_pt}.to_q.weight"))
+  set_val(flat_state[dec_attn_path + ("to_q", "bias")], get_pt_tensor(f"{dec_attn_pt}.to_q.bias"))
+  set_val(flat_state[dec_attn_path + ("to_k", "kernel")], get_linear_kernel(f"{dec_attn_pt}.to_k.weight"))
+  set_val(flat_state[dec_attn_path + ("to_k", "bias")], get_pt_tensor(f"{dec_attn_pt}.to_k.bias"))
+  set_val(flat_state[dec_attn_path + ("to_v", "kernel")], get_linear_kernel(f"{dec_attn_pt}.to_v.weight"))
+  set_val(flat_state[dec_attn_path + ("to_v", "bias")], get_pt_tensor(f"{dec_attn_pt}.to_v.bias"))
+  set_val(flat_state[dec_attn_path + ("to_out", "kernel")], get_linear_kernel(f"{dec_attn_pt}.to_out.0.weight"))
+  set_val(flat_state[dec_attn_path + ("to_out", "bias")], get_pt_tensor(f"{dec_attn_pt}.to_out.0.bias"))
+
+  for b_idx in range(4):
+    up_block_pt = f"decoder.up_blocks.{b_idx}"
+    for r_idx in range(3):
+      res_pt = f"{up_block_pt}.resnets.{r_idx}"
+      res_path = ("decoder", "up_blocks", b_idx, "resnets", r_idx)
+
+      set_val(flat_state[res_path + ("norm1", "scale")], get_pt_tensor(f"{res_pt}.norm1.weight", is_norm=True))
+      set_val(flat_state[res_path + ("norm1", "bias")], get_pt_tensor(f"{res_pt}.norm1.bias", is_norm=True))
+      set_val(flat_state[res_path + ("conv1", "kernel")], get_conv_kernel(f"{res_pt}.conv1.weight"))
+      set_val(flat_state[res_path + ("conv1", "bias")], get_pt_tensor(f"{res_pt}.conv1.bias"))
+
+      set_val(flat_state[res_path + ("norm2", "scale")], get_pt_tensor(f"{res_pt}.norm2.weight", is_norm=True))
+      set_val(flat_state[res_path + ("norm2", "bias")], get_pt_tensor(f"{res_pt}.norm2.bias", is_norm=True))
+      set_val(flat_state[res_path + ("conv2", "kernel")], get_conv_kernel(f"{res_pt}.conv2.weight"))
+      set_val(flat_state[res_path + ("conv2", "bias")], get_pt_tensor(f"{res_pt}.conv2.bias"))
+
+      shortcut_key = f"{res_pt}.conv_shortcut.weight"
+      if shortcut_key in pt_state_dict:
+        set_val(flat_state[res_path + ("conv_shortcut", "kernel")], get_conv_kernel(shortcut_key))
+        set_val(flat_state[res_path + ("conv_shortcut", "bias")], get_pt_tensor(f"{res_pt}.conv_shortcut.bias"))
+
+    if b_idx < 3:
+      ups_pt = f"{up_block_pt}.upsamplers.0.conv"
+      ups_path = ("decoder", "up_blocks", b_idx, "upsamplers_0", "conv")
+      set_val(flat_state[ups_path + ("kernel",)], get_conv_kernel(f"{ups_pt}.weight"))
+      set_val(flat_state[ups_path + ("bias",)], get_pt_tensor(f"{ups_pt}.bias"))
+
+  set_val(flat_state[("decoder", "conv_norm_out", "scale")], get_pt_tensor("decoder.conv_norm_out.weight", is_norm=True))
+  set_val(flat_state[("decoder", "conv_norm_out", "bias")], get_pt_tensor("decoder.conv_norm_out.bias", is_norm=True))
+  set_val(flat_state[("decoder", "conv_out", "kernel")], get_conv_kernel("decoder.conv_out.weight"))
+  set_val(flat_state[("decoder", "conv_out", "bias")], get_pt_tensor("decoder.conv_out.bias"))
+
+  # Update nnx_vae state
+  nnx.update(nnx_vae, nnx.from_flat_state(flat_state))
+
+  # Extract Batch Normalization running stats
+  bn_mean = jnp.array(get_pt_tensor("bn.running_mean")).reshape(1, -1, 1, 1)
+  bn_var = jnp.array(get_pt_tensor("bn.running_var")).reshape(1, -1, 1, 1)
+  batch_norm_eps = 0.0001
+  bn_std = jnp.sqrt(bn_var + batch_norm_eps)
+
+  return bn_mean, bn_std
diff --git a/src/maxdiffusion/pipelines/flux/__init__.py b/src/maxdiffusion/pipelines/flux/__init__.py
index 39ea05b57..12ed789bc 100644
--- a/src/maxdiffusion/pipelines/flux/__init__.py
+++ b/src/maxdiffusion/pipelines/flux/__init__.py
@@ -14,8 +14,22 @@
 limitations under the License.
 """
 
-_import_structure = {"pipeline_jflux": "JfluxPipeline"}
+from typing import TYPE_CHECKING
+from ...utils import DIFFUSERS_SLOW_IMPORT, _LazyModule
 
-from .flux_pipeline import (
-    FluxPipeline,
-)
+_import_structure = {
+    "flux_pipeline": ["FluxPipeline"],
+    "flux2klein_pipeline": ["FlaxFlux2KleinPipeline"],
+}
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+  from .flux_pipeline import (
+      FluxPipeline,
+  )
+  from .flux2klein_pipeline import (
+      FlaxFlux2KleinPipeline,
+  )
+else:
+  import sys
+
+  sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
diff --git a/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py b/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py
index 4fa257ca2..10414d63f 100644
--- a/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py
+++ b/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py
@@ -14,6 +14,7 @@
 limitations under the License.
 """
 
+import math
 import os
 import time
 from typing import List, Union, Optional, Any
@@ -41,7 +42,9 @@
 
 from ...models.flux.util import (
     pack_latents,
+    patchify_latents,
     prepare_latent_image_ids,
+    prepare_multi_image_ids,
     prepare_text_ids,
 )
 
@@ -105,6 +108,7 @@ def __init__(
     # JIT compilation cache
     self._jitted_qwen3_forward = None
     self._jitted_transformer_step = None
+    self._jitted_vae_encode = None
     self._jitted_vae_decode = None
 
   def _setup_jit_functions(self):
@@ -123,22 +127,57 @@ def qwen3_forward(q_params, ids, mask):
       prompt_embeds = jax.lax.with_sharding_constraint(prompt_embeds, jax.sharding.NamedSharding(self.mesh, context_spec))
       return prompt_embeds
 
-    @jax.jit(static_argnums=(4, 5), donate_argnums=(1,))
-    def vae_decode(v_params, latents_packed, vae_bn_mean, vae_bn_std, height, width):
-      batch_size_val = latents_packed.shape[0]
-      h_latent = height // 8
-      w_latent = width // 8
+    if isinstance(self.vae, nnx.Module):
+      v_graph, _, v_rest = nnx.split(self.vae, nnx.Param, ...)
 
-      vae_bn_mean_seq = vae_bn_mean.reshape(1, 1, 128)
-      vae_bn_std_seq = vae_bn_std.reshape(1, 1, 128)
+      @jax.jit
+      def vae_encode(v_params, img):
+        merged = nnx.merge(v_graph, v_params, v_rest)
+        return merged.encode(img)
+
+      @jax.jit(static_argnums=(4, 5), donate_argnums=(1,))
+      def vae_decode(v_params, latents_packed, vae_bn_mean, vae_bn_std, height, width):
+        batch_size_val = latents_packed.shape[0]
+        h_latent = height // 8
+        w_latent = width // 8
+
+        vae_bn_mean_seq = vae_bn_mean.reshape(1, 1, 128)
+        vae_bn_std_seq = vae_bn_std.reshape(1, 1, 128)
+
+        latents_bn = latents_packed * vae_bn_std_seq + vae_bn_mean_seq
+        latents_unpacked = jnp.reshape(latents_bn, (batch_size_val, h_latent // 2, w_latent // 2, 32, 2, 2))
+        latents_unpacked = jnp.transpose(latents_unpacked, (0, 3, 1, 4, 2, 5))
+        latents_unpacked = jnp.reshape(latents_unpacked, (batch_size_val, 32, h_latent, w_latent))
+
+        merged = nnx.merge(v_graph, v_params, v_rest)
+        res = merged.decode(latents_unpacked)
+        return FlaxDecoderOutput(sample=res)
+
+    else:
+
+      @jax.jit
+      def vae_encode(v_params, img):
+        # FlaxAutoencoderKL expects (B, 3, H, W)
+        res = self.vae.apply({"params": v_params}, sample=img, method=self.vae.encode)
+        moments = res.latent_dist.mode()
+        return jnp.transpose(moments, (0, 3, 1, 2))
+
+      @jax.jit(static_argnums=(4, 5), donate_argnums=(1,))
+      def vae_decode(v_params, latents_packed, vae_bn_mean, vae_bn_std, height, width):
+        batch_size_val = latents_packed.shape[0]
+        h_latent = height // 8
+        w_latent = width // 8
+
+        vae_bn_mean_seq = vae_bn_mean.reshape(1, 1, 128)
+        vae_bn_std_seq = vae_bn_std.reshape(1, 1, 128)
 
-      latents_bn = latents_packed * vae_bn_std_seq + vae_bn_mean_seq
-      latents_unpacked = jnp.reshape(latents_bn, (batch_size_val, h_latent // 2, w_latent // 2, 32, 2, 2))
-      latents_unpacked = jnp.transpose(latents_unpacked, (0, 3, 1, 4, 2, 5))
-      latents_unpacked = jnp.reshape(latents_unpacked, (batch_size_val, 32, h_latent, w_latent))
+        latents_bn = latents_packed * vae_bn_std_seq + vae_bn_mean_seq
+        latents_unpacked = jnp.reshape(latents_bn, (batch_size_val, h_latent // 2, w_latent // 2, 32, 2, 2))
+        latents_unpacked = jnp.transpose(latents_unpacked, (0, 3, 1, 4, 2, 5))
+        latents_unpacked = jnp.reshape(latents_unpacked, (batch_size_val, 32, h_latent, w_latent))
 
-      res = self.vae.apply({"params": v_params}, latents=latents_unpacked, method=self.vae.decode)
-      return FlaxDecoderOutput(sample=res.sample)
+        res = self.vae.apply({"params": v_params}, latents=latents_unpacked, method=self.vae.decode)
+        return FlaxDecoderOutput(sample=res.sample)
 
     if isinstance(self.transformer, nnx.Module):
       g, nnx_state, r = nnx.split(self.transformer, nnx.Param, ...)
@@ -157,8 +196,10 @@ def transformer_step(t_params, latents, img_ids, prompt_embeds, txt_ids, vec, ti
             return_dict=True,
         )
 
-      @jax.jit
-      def fused_denoise_loop(t_params, latents, img_ids, prompt_embeds, txt_ids, vec, timesteps, sigmas, guidance):
+      @jax.jit(static_argnums=(9,))
+      def fused_denoise_loop(
+          t_params, latents, img_ids, prompt_embeds, txt_ids, vec, timesteps, sigmas, guidance, target_len=None
+      ):
         sigmas_padded = jnp.concatenate([sigmas, jnp.array([0.0], dtype=sigmas.dtype)])
         nnx_merged = nnx.merge(g, t_params, r)
 
@@ -177,13 +218,87 @@ def scan_body(cur_latents, step_idx):
           )
           sigma = sigmas_padded[step_idx]
           sigma_next = sigmas_padded[step_idx + 1]
-          prev_sample = cur_latents + model_output.sample * (sigma_next - sigma)
+          dt = sigma_next - sigma
+          v = model_output.sample
+          if target_len is not None and cur_latents.shape[1] > target_len:
+            target_latents = cur_latents[:, :target_len, :]
+            v_target = v[:, :target_len, :]
+            next_target = target_latents + v_target * dt
+            prev_sample = jnp.concatenate([next_target, cur_latents[:, target_len:, :]], axis=1)
+          else:
+            prev_sample = cur_latents + v * dt
           return prev_sample, None
 
         steps = jnp.arange(timesteps.shape[0])
         final_latents, _ = jax.lax.scan(scan_body, latents, steps)
         return final_latents
 
+      @jax.jit(static_argnums=(10, 11))
+      def fused_kv_denoise_loop(
+          t_params,
+          target_latents,
+          ref_latents,
+          target_img_ids,
+          ref_img_ids,
+          prompt_embeds,
+          txt_ids,
+          vec,
+          timesteps,
+          sigmas,
+          guidance=None,
+          num_ref_tokens=0,
+      ):
+        sigmas_padded = jnp.concatenate([sigmas, jnp.array([0.0], dtype=sigmas.dtype)])
+        nnx_merged = nnx.merge(g, t_params, r)
+
+        step0_latents = jnp.concatenate([ref_latents, target_latents], axis=1)
+        step0_img_ids = jnp.concatenate([ref_img_ids, target_img_ids], axis=1)
+        t0_val = timesteps[0]
+        t0_vec = jnp.broadcast_to(t0_val / 1000.0, (target_latents.shape[0],))
+
+        out0, kv_cache = nnx_merged(
+            hidden_states=step0_latents,
+            img_ids=step0_img_ids,
+            encoder_hidden_states=prompt_embeds,
+            txt_ids=txt_ids,
+            pooled_projections=vec,
+            timestep=t0_vec,
+            guidance=guidance,
+            return_dict=True,
+            kv_cache_mode="extract",
+            num_ref_tokens=num_ref_tokens,
+        )
+        dt0 = sigmas_padded[1] - sigmas_padded[0]
+        v0 = out0.sample
+        latents_step1 = target_latents + v0 * dt0
+
+        def scan_body(cur_latents, step_idx):
+          t_val = timesteps[step_idx]
+          t_vec = jnp.broadcast_to(t_val / 1000.0, (cur_latents.shape[0],))
+          model_output = nnx_merged(
+              hidden_states=cur_latents,
+              img_ids=target_img_ids,
+              encoder_hidden_states=prompt_embeds,
+              txt_ids=txt_ids,
+              pooled_projections=vec,
+              timestep=t_vec,
+              guidance=guidance,
+              return_dict=True,
+              kv_cache=kv_cache,
+              kv_cache_mode="cached",
+              num_ref_tokens=num_ref_tokens,
+          )
+          sigma = sigmas_padded[step_idx]
+          sigma_next = sigmas_padded[step_idx + 1]
+          dt = sigma_next - sigma
+          v = model_output.sample
+          next_latents = cur_latents + v * dt
+          return next_latents, None
+
+        steps = jnp.arange(1, timesteps.shape[0])
+        final_latents, _ = jax.lax.scan(scan_body, latents_step1, steps)
+        return final_latents
+
     else:
 
       @jax.jit
@@ -199,8 +314,10 @@ def transformer_step(t_params, latents, img_ids, prompt_embeds, txt_ids, vec, ti
             guidance=guidance,
         )
 
-      @jax.jit
-      def fused_denoise_loop(t_params, latents, img_ids, prompt_embeds, txt_ids, vec, timesteps, sigmas, guidance):
+      @jax.jit(static_argnums=(9,))
+      def fused_denoise_loop(
+          t_params, latents, img_ids, prompt_embeds, txt_ids, vec, timesteps, sigmas, guidance, target_len=None
+      ):
         sigmas_padded = jnp.concatenate([sigmas, jnp.array([0.0], dtype=sigmas.dtype)])
 
         def scan_body(cur_latents, step_idx):
@@ -218,18 +335,70 @@ def scan_body(cur_latents, step_idx):
           )
           sigma = sigmas_padded[step_idx]
           sigma_next = sigmas_padded[step_idx + 1]
-          prev_sample = cur_latents + model_output.sample * (sigma_next - sigma)
+          dt = sigma_next - sigma
+          v = model_output.sample
+          if target_len is not None and cur_latents.shape[1] > target_len:
+            target_latents = cur_latents[:, :target_len, :]
+            v_target = v[:, :target_len, :]
+            next_target = target_latents + v_target * dt
+            prev_sample = jnp.concatenate([next_target, cur_latents[:, target_len:, :]], axis=1)
+          else:
+            prev_sample = cur_latents + v * dt
           return prev_sample, None
 
         steps = jnp.arange(timesteps.shape[0])
         final_latents, _ = jax.lax.scan(scan_body, latents, steps)
         return final_latents
 
+      fused_kv_denoise_loop = fused_denoise_loop
+
     self._jitted_qwen3_forward = qwen3_forward
     self._jitted_transformer_step = transformer_step
     self._jitted_fused_denoise_loop = fused_denoise_loop
+    self._jitted_fused_kv_denoise_loop = fused_kv_denoise_loop
+    self._jitted_vae_encode = vae_encode
     self._jitted_vae_decode = vae_decode
 
+  @staticmethod
+  def preprocess_reference_image(
+      image: Image.Image,
+      max_area: int = 1024 * 1024,
+      multiple_of: int = 16,
+  ) -> np.ndarray:
+    """Validates and preprocesses a single PIL reference image into BCHW [-1, 1] array."""
+    if not isinstance(image, Image.Image):
+      raise TypeError(f"Expected PIL.Image.Image for reference image, got {type(image)}")
+
+    width, height = image.size
+    if width < 64 or height < 64:
+      raise ValueError(f"Image too small: {width}x{height}. Both dimensions must be at least 64px")
+    aspect_ratio = max(width / height, height / width)
+    if aspect_ratio > 8.0:
+      raise ValueError(f"Aspect ratio too extreme: {width}x{height} (ratio: {aspect_ratio:.1f}:1 > 8:1)")
+
+    if image.mode != "RGB":
+      image = image.convert("RGB")
+
+    if width * height > max_area:
+      scale = math.sqrt(max_area / (width * height))
+      width = int(width * scale)
+      height = int(height * scale)
+      image = image.resize((width, height), Image.Resampling.LANCZOS)
+      width, height = image.size
+
+    target_w = (width // multiple_of) * multiple_of
+    target_h = (height // multiple_of) * multiple_of
+
+    left = (width - target_w) // 2
+    top = (height - target_h) // 2
+    right = left + target_w
+    bottom = top + target_h
+    image = image.crop((left, top, right, bottom))
+
+    arr = np.asarray(image, dtype=np.float32) / 127.5 - 1.0
+    arr = np.transpose(arr, (2, 0, 1))  # (3, H, W)
+    return np.expand_dims(arr, axis=0)  # (1, 3, H, W)
+
   def _get_dynamic_batch_sharding(self):
     """Dynamically infers the batch dimension sharding specification from self.mesh."""
     batch_axes = [axis for axis in ("data", "fsdp") if axis in self.mesh.axis_names and self.mesh.shape[axis] > 1]
@@ -237,25 +406,61 @@ def _get_dynamic_batch_sharding(self):
     return jax.sharding.NamedSharding(self.mesh, spec)
 
   def compile_aot_async(
-      self, params, vae_params, qwen3_params, vae_bn_mean, vae_bn_std, batch_size=1, height=1024, width=1024
+      self,
+      params,
+      vae_params,
+      qwen3_params,
+      vae_bn_mean,
+      vae_bn_std,
+      batch_size=1,
+      height=1024,
+      width=1024,
+      images=None,
+      num_conditioning_images=0,
+      use_kv=None,
   ):
     """Triggers AOT compilation for Qwen3, Flux Transformer, and VAE concurrently using ThreadPoolExecutor."""
     self._setup_jit_functions()
     max_logging.log("🚀 Pre-compiling XLA graphs for Qwen3, Flux Transformer, and VAE concurrently...")
     from concurrent.futures import ThreadPoolExecutor
 
+    if images is not None and isinstance(images, Image.Image):
+      images = [images]
+
     seq_len_img = (height // 16) * (width // 16)
+    total_ref_tokens = 0
+    ref_shapes = []
+    if images is not None and len(images) > 0:
+      for img in images:
+        if not isinstance(img, Image.Image):
+          raise TypeError(f"Expected PIL.Image.Image for reference image, got {type(img)}")
+        w, h = img.size
+        if w * h > 1024 * 1024:
+          scale = math.sqrt((1024 * 1024) / (w * h))
+          w = int(w * scale)
+          h = int(h * scale)
+        target_w = (w // 16) * 16
+        target_h = (h // 16) * 16
+        ref_shapes.append((target_h, target_w))
+        total_ref_tokens += (target_h // 16) * (target_w // 16)
+      num_conditioning_images = len(images)
+    elif num_conditioning_images > 0:
+      # If num_conditioning_images specified without images, assume default output size
+      total_ref_tokens = num_conditioning_images * seq_len_img
+
+    total_img_len = seq_len_img + total_ref_tokens
     seq_len_txt = self._config.max_sequence_length
 
     dummy_ids = jnp.zeros((batch_size, seq_len_txt), dtype=jnp.int32)
     dummy_mask = jnp.ones((batch_size, seq_len_txt), dtype=jnp.int32)
 
-    dummy_latents = jnp.zeros((batch_size, seq_len_img, 128), dtype=jnp.float32)
-    dummy_img_ids = jnp.zeros((batch_size, seq_len_img, 4), dtype=jnp.int32)
+    dummy_latents = jnp.zeros((batch_size, total_img_len, 128), dtype=jnp.float32)
+    dummy_img_ids = jnp.zeros((batch_size, total_img_len, 4), dtype=jnp.int32)
     dummy_prompt_embeds = jnp.zeros((batch_size, seq_len_txt, self.transformer.joint_attention_dim), dtype=jnp.bfloat16)
     dummy_txt_ids = jnp.zeros((batch_size, seq_len_txt, 4), dtype=jnp.float32)
     dummy_t_vec = jnp.zeros((batch_size,), dtype=jnp.float32)
 
+    dummy_target_latents = jnp.zeros((batch_size, seq_len_img, 128), dtype=jnp.float32)
     dummy_bn_mean = jnp.array(vae_bn_mean, dtype=jnp.float32)
     dummy_bn_std = jnp.array(vae_bn_std, dtype=jnp.float32)
 
@@ -277,6 +482,7 @@ def put_data_on_devices(x, sharding):
     dummy_prompt_embeds = put_data_on_devices(dummy_prompt_embeds, context_sharding)
     dummy_txt_ids = put_data_on_devices(dummy_txt_ids, data_sharding)
     dummy_t_vec = put_data_on_devices(dummy_t_vec, data_sharding)
+    dummy_target_latents = put_data_on_devices(dummy_target_latents, data_sharding)
     dummy_bn_mean = put_data_on_devices(dummy_bn_mean, replicated_sharding)
     dummy_bn_std = put_data_on_devices(dummy_bn_std, replicated_sharding)
 
@@ -290,35 +496,78 @@ def compile_qwen3():
     dummy_timesteps = put_data_on_devices(jnp.zeros((num_steps,), dtype=jnp.float32), replicated_sharding)
     dummy_sigmas = put_data_on_devices(jnp.zeros((num_steps + 1,), dtype=jnp.float32), replicated_sharding)
 
+    use_kv = getattr(self._config, "use_kv", False) if use_kv is None else use_kv
+    dummy_ref_latents = (
+        put_data_on_devices(jnp.zeros((batch_size, total_ref_tokens, 128), dtype=jnp.float32), data_sharding)
+        if total_ref_tokens > 0
+        else None
+    )
+    dummy_ref_img_ids = (
+        put_data_on_devices(jnp.zeros((batch_size, total_ref_tokens, 4), dtype=jnp.int32), data_sharding)
+        if total_ref_tokens > 0
+        else None
+    )
+    dummy_target_img_ids = put_data_on_devices(jnp.zeros((batch_size, seq_len_img, 4), dtype=jnp.int32), data_sharding)
+
     def compile_transformer():
       t0 = time.perf_counter()
       with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules):
-        self._jitted_fused_denoise_loop.lower(
-            params,
-            dummy_latents,
-            dummy_img_ids,
-            dummy_prompt_embeds,
-            dummy_txt_ids,
-            None,
-            dummy_timesteps,
-            dummy_sigmas,
-            None,
-        ).compile()
-      max_logging.log(f" -> [AOT COMPILED] Fused Flux Transformer Denoise Scan in {time.perf_counter() - t0:.2f}s")
+        if use_kv and total_ref_tokens > 0 and self._jitted_fused_kv_denoise_loop is not None:
+          self._jitted_fused_kv_denoise_loop.lower(
+              params,
+              dummy_target_latents,
+              dummy_ref_latents,
+              dummy_target_img_ids,
+              dummy_ref_img_ids,
+              dummy_prompt_embeds,
+              dummy_txt_ids,
+              None,
+              dummy_timesteps,
+              dummy_sigmas,
+              None,
+              num_ref_tokens=total_ref_tokens,
+          ).compile()
+          max_logging.log(
+              f" -> [AOT COMPILED] Fused Flux Transformer KV Denoise Scan in {time.perf_counter() - t0:.2f}s"
+          )
+        else:
+          self._jitted_fused_denoise_loop.lower(
+              params,
+              dummy_latents,
+              dummy_img_ids,
+              dummy_prompt_embeds,
+              dummy_txt_ids,
+              None,
+              dummy_timesteps,
+              dummy_sigmas,
+              None,
+              seq_len_img,
+          ).compile()
+          max_logging.log(f" -> [AOT COMPILED] Fused Flux Transformer Denoise Scan in {time.perf_counter() - t0:.2f}s")
 
     def compile_vae():
       t0 = time.perf_counter()
       with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules):
-        self._jitted_vae_decode.lower(vae_params, dummy_latents, dummy_bn_mean, dummy_bn_std, height, width).compile()
+        self._jitted_vae_decode.lower(vae_params, dummy_target_latents, dummy_bn_mean, dummy_bn_std, height, width).compile()
       max_logging.log(f" -> [AOT COMPILED] VAE Decoder in {time.perf_counter() - t0:.2f}s")
 
+    def compile_vae_encode():
+      t0 = time.perf_counter()
+      enc_h, enc_w = ref_shapes[0] if ref_shapes else (height, width)
+      dummy_rgb = jnp.zeros((1, 3, enc_h, enc_w), dtype=jnp.float32)
+      with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules):
+        self._jitted_vae_encode.lower(vae_params, dummy_rgb).compile()
+      max_logging.log(f" -> [AOT COMPILED] VAE Encoder in {time.perf_counter() - t0:.2f}s")
+
     t_start = time.perf_counter()
-    with ThreadPoolExecutor(max_workers=3) as executor:
+    with ThreadPoolExecutor(max_workers=4) as executor:
       futures = [
           executor.submit(compile_qwen3),
           executor.submit(compile_transformer),
           executor.submit(compile_vae),
       ]
+      if num_conditioning_images > 0 or (images is not None and len(images) > 0):
+        futures.append(executor.submit(compile_vae_encode))
       for future in futures:
         future.result()
     aot_duration = time.perf_counter() - t_start
@@ -364,8 +613,10 @@ def __call__(
       width: int = 1024,
       num_inference_steps: int = 4,
       batch_size: int = 1,
+      images: Optional[Union[Image.Image, List[Image.Image]]] = None,
       use_latents: bool = False,
       latents: Optional[Any] = None,
+      use_kv: Optional[bool] = None,
       measure_time: bool = False,
       warmup: bool = False,
       output_dir: str = "output/",
@@ -375,6 +626,17 @@ def __call__(
     # 1. Setup JIT functions
     self._setup_jit_functions()
 
+    if images is not None:
+      if isinstance(images, Image.Image):
+        images = [images]
+      elif not isinstance(images, (list, tuple)):
+        raise TypeError(
+            f"Expected reference images to be a PIL.Image.Image or list/tuple of PIL.Image.Image, got {type(images)}"
+        )
+      for idx, img in enumerate(images):
+        if not isinstance(img, Image.Image):
+          raise TypeError(f"Expected reference image at index {idx} to be PIL.Image.Image, got {type(img)}")
+
     # 2. Setup prompts and inputs
     if isinstance(prompt, str):
       prompts = [prompt] * batch_size
@@ -392,6 +654,8 @@ def __call__(
         if C == 32:
           max_logging.log("  [PIPELINE] Unpacked 32-channel latents detected. Packing using pack_latents...")
           latents_jax = pack_latents(latents_jax)
+        elif C == 128:
+          latents_jax = jnp.transpose(jnp.reshape(latents_jax, (B, C, H * W)), (0, 2, 1))
         else:
           latents_jax = jnp.transpose(jnp.reshape(latents_jax, (B, C, H * W)), (0, 2, 1))
     else:
@@ -401,7 +665,9 @@ def __call__(
 
     # RoPE position IDs
     txt_ids_val = prepare_text_ids(batch_size, seq_len_txt)
-    img_ids_val = prepare_latent_image_ids(batch_size, height // 16, width // 16)
+    target_img_ids_val = prepare_latent_image_ids(batch_size, height // 16, width // 16)
+    t_pipeline_start = time.perf_counter()
+    trace = {}
 
     # Scheduler
     mu = compute_empirical_mu(seq_len_img, num_inference_steps)
@@ -414,9 +680,6 @@ def __call__(
         sigmas=sigmas_custom,
     )
 
-    t_pipeline_start = time.perf_counter()
-    trace = {}
-
     with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules):
       proc_id = jax.process_index()
       proc_cnt = jax.process_count()
@@ -432,9 +695,63 @@ def put_data_on_devices(x, sharding):
           return jax.device_put(x, sharding)
         return device_put_replicated(x, sharding)
 
+      # ---------------------------------------------------------------------
+      # PHASE 0: Encode Reference Images (VAE)
+      # ---------------------------------------------------------------------
+      if images is not None and len(images) > 0:
+        t0_vae_enc_start = time.perf_counter()
+        trace["start_to_vae_encode"] = t0_vae_enc_start - t_pipeline_start
+        max_logging.log(f"{host_prefix} [PHASE 0] Encoding {len(images)} reference image(s) using JAX VAE encoder on TPU...")
+        norm_ref_latents = []
+        packed_ref_latents = []
+        bn_mean_arr = jnp.array(vae_bn_mean, dtype=jnp.float32)
+        bn_std_arr = jnp.array(vae_bn_std, dtype=jnp.float32)
+
+        for img in images:
+          img_np = self.preprocess_reference_image(img)
+          img_tensor = jnp.array(img_np, dtype=jnp.float32)
+
+          raw_ref_latents = self._jitted_vae_encode(vae_params, img_tensor)
+          raw_ref_latents.block_until_ready()
+          patchified_ref = patchify_latents(raw_ref_latents)
+          normalized_ref = (patchified_ref - bn_mean_arr) / bn_std_arr
+          norm_ref_latents.append(normalized_ref)
+
+          packed = jnp.transpose(
+              jnp.reshape(normalized_ref, (normalized_ref.shape[0], normalized_ref.shape[1], -1)), (0, 2, 1)
+          )
+          if packed.shape[0] == 1 and batch_size > 1:
+            packed = jnp.repeat(packed, batch_size, axis=0)
+          packed_ref_latents.append(packed)
+
+        ref_img_ids_val = prepare_multi_image_ids(norm_ref_latents, scale=10)
+        if ref_img_ids_val.shape[0] == 1 and batch_size > 1:
+          ref_img_ids_val = jnp.repeat(ref_img_ids_val, batch_size, axis=0)
+        ref_latents_jax = jnp.concatenate(packed_ref_latents, axis=1)
+        num_ref_tokens = ref_latents_jax.shape[1]
+        img_ids_val = jnp.concatenate([target_img_ids_val, ref_img_ids_val], axis=1)
+        latents_jax = jnp.concatenate([latents_jax] + packed_ref_latents, axis=1)
+        max_logging.log(f"  [PIPELINE] Joint latents shape: {latents_jax.shape}, Joint img_ids shape: {img_ids_val.shape}")
+
+        t0_vae_enc_end = time.perf_counter()
+        trace["vae_encode"] = t0_vae_enc_end - t0_vae_enc_start
+        trace["image_encoding"] = trace["vae_encode"]
+        max_logging.log(f" -> [TIMING] Reference Image Encoding (VAE): {trace['vae_encode']:.4f} seconds ⏱️")
+      else:
+        img_ids_val = target_img_ids_val
+        packed_ref_latents = []
+        ref_latents_jax = None
+        num_ref_tokens = 0
+        trace["vae_encode"] = 0.0
+        trace["image_encoding"] = 0.0
+
       t0_qwen3_start = time.perf_counter()
-      trace["start_to_qwen3"] = t0_qwen3_start - t_pipeline_start
-      max_logging.log(f" -> [TIMING] Start to Qwen3: {trace['start_to_qwen3']:.4f} seconds ⏱️")
+      if trace.get("vae_encode", 0.0) > 0:
+        trace["vae_encode_to_qwen3"] = t0_qwen3_start - t0_vae_enc_end
+        max_logging.log(f" -> [TIMING] VAE Encode to Qwen3 Overhead: {trace['vae_encode_to_qwen3']:.4f} seconds ⏱️")
+      else:
+        trace["start_to_qwen3"] = t0_qwen3_start - t_pipeline_start
+        max_logging.log(f" -> [TIMING] Start to Qwen3: {trace['start_to_qwen3']:.4f} seconds ⏱️")
 
       # ---------------------------------------------------------------------
       # PHASE A: Encode Prompt (Qwen3)
@@ -447,10 +764,18 @@ def put_data_on_devices(x, sharding):
       max_logging.log(f"{host_prefix} [PHASE A] Encoding {len(prompts)} prompt(s) using JAX Qwen3 on TPU...")
 
       try:
-        # Tokenize using deterministic explicit template string (version-agnostic across transformers versions)
-        templated_texts = [
-            f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n" for p in prompts
-        ]
+        if hasattr(self.tokenizer, "apply_chat_template"):
+          templated_texts = [
+              self.tokenizer.apply_chat_template(
+                  [{"role": "user", "content": p}],
+                  tokenize=False,
+                  add_generation_prompt=True,
+                  enable_thinking=False,
+              )
+              for p in prompts
+          ]
+        else:
+          templated_texts = [f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n" for p in prompts]
         inputs = self.tokenizer(
             templated_texts, return_tensors="np", padding="max_length", truncation=True, max_length=seq_len_txt
         )
@@ -528,25 +853,44 @@ def put_data_on_devices(x, sharding):
         timesteps_device = put_data_on_devices(scheduler_state.timesteps, replicated_sharding)
         sigmas_device = put_data_on_devices(scheduler_state.sigmas, replicated_sharding)
 
-        do_prof_denoise = profile_target in ("all", "denoise")
-        if do_prof_denoise:
-          tb_dir = getattr(self._config, "tensorboard_dir", "/tmp")
-          jax.profiler.start_trace(os.path.join(tb_dir, "profile_denoise"))
-        with jax.named_scope("fused_flux_denoise_loop"):
-          latents_jax = self._jitted_fused_denoise_loop(
-              params,
-              latents_jax,
-              img_ids_val,
-              prompt_embeds_jax,
-              txt_ids_val,
-              vec_val,
-              timesteps_device,
-              sigmas_device,
-              guidance_vec_val,
-          )
-          latents_jax.block_until_ready()
-        if do_prof_denoise:
-          jax.profiler.stop_trace()
+        use_kv = getattr(self._config, "use_kv", False) if use_kv is None else use_kv
+        if use_kv and len(packed_ref_latents) > 0:
+          ref_latents_device = put_data_on_devices(ref_latents_jax, data_sharding)
+          ref_img_ids_device = put_data_on_devices(ref_img_ids_val, data_sharding)
+          target_img_ids_device = put_data_on_devices(target_img_ids_val, data_sharding)
+          target_latents_device = put_data_on_devices(latents_jax[:, :seq_len_img, :], data_sharding)
+
+          with jax.named_scope("fused_flux_kv_denoise_loop"):
+            latents_jax = self._jitted_fused_kv_denoise_loop(
+                params,
+                target_latents_device,
+                ref_latents_device,
+                target_img_ids_device,
+                ref_img_ids_device,
+                prompt_embeds_jax,
+                txt_ids_val,
+                vec_val,
+                timesteps_device,
+                sigmas_device,
+                guidance_vec_val,
+                num_ref_tokens,
+            )
+            latents_jax.block_until_ready()
+        else:
+          with jax.named_scope("fused_flux_denoise_loop"):
+            latents_jax = self._jitted_fused_denoise_loop(
+                params,
+                latents_jax,
+                img_ids_val,
+                prompt_embeds_jax,
+                txt_ids_val,
+                vec_val,
+                timesteps_device,
+                sigmas_device,
+                guidance_vec_val,
+                seq_len_img,
+            )
+            latents_jax.block_until_ready()
 
       except Exception as e:
         max_logging.log(f"❌ {host_prefix} EXCEPTION IN DENOISE LOOP: {e}")
@@ -569,6 +913,10 @@ def put_data_on_devices(x, sharding):
     # ---------------------------------------------------------------------
     max_logging.log("[PHASE C] Decoding final latents to RGB image using JAX VAE decoder on TPU...")
 
+    # Slice target latents from joint latents if reference images were present
+    if latents_jax.shape[1] > seq_len_img:
+      latents_jax = latents_jax[:, :seq_len_img, :]
+
     # Decode VAE latents to RGB pixels using fused JIT vae_decode
     data_sharding = self._get_dynamic_batch_sharding()
     replicated_sharding = jax.sharding.NamedSharding(self.mesh, P())
diff --git a/src/maxdiffusion/pipelines/pipeline_flax_utils.py b/src/maxdiffusion/pipelines/pipeline_flax_utils.py
index a2d4cc8db..a75bc5a4b 100644
--- a/src/maxdiffusion/pipelines/pipeline_flax_utils.py
+++ b/src/maxdiffusion/pipelines/pipeline_flax_utils.py
@@ -40,8 +40,7 @@
     logging,
 )
 
-
-from transformers import FlaxPreTrainedModel
+import transformers
 
 INDEX_FILE = "diffusion_flax_model.bin"
 
@@ -497,6 +496,7 @@ def load_module(name, value):
         else:
           loaded_sub_model = cached_folder
 
+        flax_pretrained_model_cls = getattr(transformers, "FlaxPreTrainedModel", None)
         if issubclass(class_obj, FlaxModelMixin):
           loaded_sub_model, loaded_params = load_method(
               loadable_folder,
@@ -514,7 +514,7 @@ def load_module(name, value):
               quant=quant,
           )
           params[name] = loaded_params
-        elif issubclass(class_obj, FlaxPreTrainedModel):
+        elif flax_pretrained_model_cls is not None and issubclass(class_obj, flax_pretrained_model_cls):
           if from_pt:
             # TODO(Suraj): Fix this in Transformers. We should be able to use `_do_init=False` here
             loaded_sub_model = load_method(loadable_folder, from_pt=from_pt)
diff --git a/src/maxdiffusion/tests/edit_flux2klein_e2e_test.py b/src/maxdiffusion/tests/edit_flux2klein_e2e_test.py
new file mode 100644
index 000000000..c83604265
--- /dev/null
+++ b/src/maxdiffusion/tests/edit_flux2klein_e2e_test.py
@@ -0,0 +1,348 @@
+"""
+Copyright 2026 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import os
+import gc
+import unittest
+import pytest
+import numpy as np
+from PIL import Image
+from skimage.metrics import structural_similarity as ssim
+import torch
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from jax.sharding import Mesh
+from transformers import AutoConfig, Qwen2TokenizerFast
+
+from maxdiffusion import max_logging, pyconfig
+from maxdiffusion.max_utils import create_device_mesh
+from maxdiffusion.models.flux.transformers.transformer_flux_flax import NNXFlux2KleinTransformer2DModel
+from maxdiffusion.models.flux.vae.autoencoder_kl_flux2_nnx import (
+    NNXAutoencoderKLFlux2,
+    load_and_convert_flux2klein_nnx_vae_weights,
+)
+from maxdiffusion.models.flux.util import load_and_convert_flux_klein_nnx_weights
+from maxdiffusion.models.qwen3_flax import FlaxQwen3Model, FlaxQwen3Config
+from maxdiffusion.models.qwen3_utils import load_and_convert_qwen3_weights
+from maxdiffusion.schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler
+from maxdiffusion.pipelines.flux.flux2klein_pipeline import FlaxFlux2KleinPipeline
+
+IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true"
+THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+PROMPT = "a vibrant artistic painting combining the dog, car, mountain, and fruit bowl in surreal neon lighting"
+
+
+def get_model_snapshot_dir(model_id: str) -> str:
+  """Locates the snapshot directory for the given model_id strictly under HF_HOME."""
+  try:
+    from huggingface_hub import snapshot_download
+
+    return snapshot_download(repo_id=model_id, local_files_only=True)
+  except Exception:
+    pass
+
+  hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
+  escaped_id = model_id.replace("/", "--")
+  cache_dir = os.path.join(hf_home, "hub", f"models--{escaped_id}", "snapshots")
+  if not os.path.exists(cache_dir):
+    cache_dir = os.path.join(hf_home, f"models--{escaped_id}", "snapshots")
+
+  if not os.path.exists(cache_dir):
+    raise FileNotFoundError(f"Hugging Face cache directory not found for '{model_id}' under HF_HOME ({hf_home}).")
+
+  snapshots = [s for s in os.listdir(cache_dir) if not s.startswith(".")]
+  if not snapshots:
+    raise FileNotFoundError(f"No snapshot directory found for '{model_id}' in {cache_dir}.")
+
+  return os.path.join(cache_dir, snapshots[0])
+
+
+class TestFlux2KleinImageEditE2EParity(unittest.TestCase):
+  """End-to-End Parity Test between PyTorch Diffusers CPU and MaxDiffusion TPU."""
+
+  def setUp(self):
+    jax.config.update("jax_use_shardy_partitioner", True)
+
+    if "FLUX2_KLEIN_4B_MODEL_PATH" in os.environ:
+      self.model_dir = os.environ["FLUX2_KLEIN_4B_MODEL_PATH"]
+    else:
+      self.model_dir = get_model_snapshot_dir("black-forest-labs/FLUX.2-klein-4B")
+
+    self.transformer_path = os.path.join(self.model_dir, "transformer")
+    self.vae_path = os.path.join(self.model_dir, "vae", "diffusion_pytorch_model.safetensors")
+    self.text_encoder_path = os.path.join(self.model_dir, "text_encoder")
+    self.tokenizer_path = os.path.join(self.model_dir, "tokenizer")
+
+    self.output_dir = "/tmp/e2e_parity"
+    os.makedirs(self.output_dir, exist_ok=True)
+
+    # Resolve reference images
+    ref_dir = os.path.join(THIS_DIR, "images", "flux2klein")
+    self.ref_images = []
+    if os.path.exists(ref_dir):
+      for i in range(4):
+        p = os.path.join(ref_dir, f"ref_image_{i}.png")
+        if os.path.exists(p):
+          self.ref_images.append(Image.open(p).convert("RGB"))
+
+    if len(self.ref_images) < 4:
+      # Generate synthetic test reference images if not present
+      colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0)]
+      for i, c in enumerate(colors):
+        arr = np.full((512, 512, 3), c, dtype=np.uint8)
+        self.ref_images.append(Image.fromarray(arr))
+
+  @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run on Github Actions (requires TPU and full weights)")
+  def test_e2e_image_edit_parity_vs_diffusers(self):
+    """Generates an image edit on PyTorch Diffusers CPU and MaxDiffusion TPU and asserts SSIM >= 0.75."""
+    from diffusers import Flux2KleinPipeline as DiffusersFlux2KleinPipeline
+
+    max_logging.log("\n" + "=" * 80)
+    max_logging.log("🚀 [STEP 1/3] Running Reference PyTorch Diffusers CPU Pipeline...")
+    max_logging.log("=" * 80)
+
+    diffusers_pipe = DiffusersFlux2KleinPipeline.from_pretrained(self.model_dir, torch_dtype=torch.bfloat16)
+    diffusers_pipe.to("cpu")
+
+    # Generate initial noise latents deterministically on CPU (4D tensor for Diffusers prepare_latents)
+    gen = torch.Generator(device="cpu").manual_seed(42)
+    raw_latents_pt = torch.randn(
+        (1, 128, 512 // 16, 512 // 16),
+        generator=gen,
+        dtype=torch.bfloat16,
+        device="cpu",
+    )
+
+    with torch.no_grad():
+      diffusers_out = diffusers_pipe(
+          prompt=PROMPT,
+          image=self.ref_images,
+          height=512,
+          width=512,
+          num_inference_steps=4,
+          latents=raw_latents_pt,
+          guidance_scale=1.0,
+      )
+
+    diffusers_image = diffusers_out.images[0]
+    diffusers_img_path = os.path.join(self.output_dir, "diffusers_cpu_output.png")
+    diffusers_image.save(diffusers_img_path)
+    max_logging.log(f" -> Saved PyTorch Diffusers output to: {diffusers_img_path}")
+
+    # Free PyTorch pipeline memory before TPU run
+    del diffusers_pipe
+    gc.collect()
+
+    max_logging.log("\n" + "=" * 80)
+    max_logging.log("🚀 [STEP 2/3] Running MaxDiffusion Unified FlaxFlux2KleinPipeline on TPU...")
+    max_logging.log("=" * 80)
+
+    # 1. Device mesh setup
+    active_devices = jax.devices()
+    active_device_count = len(active_devices)
+
+    pyconfig._config = None
+    pyconfig.config = None
+    config_path = os.path.join(THIS_DIR, "..", "configs", "base_flux2klein.yml")
+    args = [
+        None,
+        config_path,
+        "run_name=e2e_parity_test",
+        f"output_dir={self.output_dir}",
+        f"per_device_batch_size={1.0 / active_device_count}",
+        "height=512",
+        "width=512",
+        "seed=42",
+        "weights_dtype=bfloat16",
+        "activations_dtype=bfloat16",
+        "precision=DEFAULT",
+        "text_encoder_attention=dot_product",
+    ]
+    pyconfig.initialize(args)
+    config = pyconfig.config
+
+    if active_device_count > 1:
+      pyconfig._config.keys["ici_tensor_parallelism"] = active_device_count
+      pyconfig._config.keys["ici_data_parallelism"] = 1
+      pyconfig._config.keys["ici_fsdp_parallelism"] = 1
+      pyconfig._config.keys["ici_context_parallelism"] = 1
+
+    devices_array = create_device_mesh(config, devices=active_devices)
+    mesh = Mesh(devices_array, config.mesh_axes)
+
+    # 2. Load NNX Transformer
+    max_logging.log(" -> Loading NNX Transformer weights...")
+    rngs = nnx.Rngs(0)
+    transformer = NNXFlux2KleinTransformer2DModel(
+        rngs=rngs,
+        patch_size=1,
+        in_channels=128,
+        num_layers=5,
+        num_single_layers=20,
+        attention_head_dim=128,
+        num_attention_heads=24,
+        joint_attention_dim=7680,
+        pooled_projection_dim=None,
+        guidance_embeds=False,
+        axes_dim=(32, 32, 32, 32),
+        scale_shift_order="scale_shift",
+        dtype=jnp.bfloat16,
+        weights_dtype=jnp.bfloat16,
+    )
+    t_state = load_and_convert_flux_klein_nnx_weights(
+        self.transformer_path,
+        nnx.state(transformer, nnx.Param),
+        num_double_layers=5,
+        num_single_layers=20,
+        dtype=jnp.bfloat16,
+    )
+    nnx.update(transformer, t_state)
+
+    # 3. Load NNX VAE
+    max_logging.log(" -> Loading NNX VAE weights...")
+    nnx_vae = NNXAutoencoderKLFlux2(dtype=jnp.bfloat16, param_dtype=jnp.bfloat16)
+    bn_mean, bn_std = load_and_convert_flux2klein_nnx_vae_weights(self.vae_path, nnx_vae, dtype=jnp.bfloat16)
+
+    # 4. Load Qwen3
+    max_logging.log(" -> Loading Qwen3 weights...")
+    pt_config = AutoConfig.from_pretrained(self.text_encoder_path)
+    qwen3_config = FlaxQwen3Config(
+        vocab_size=pt_config.vocab_size,
+        hidden_size=pt_config.hidden_size,
+        intermediate_size=pt_config.intermediate_size,
+        num_hidden_layers=pt_config.num_hidden_layers,
+        num_attention_heads=pt_config.num_attention_heads,
+        num_key_value_heads=getattr(pt_config, "num_key_value_heads", pt_config.num_attention_heads),
+        max_position_embeddings=getattr(pt_config, "max_position_embeddings", 32768),
+        rms_norm_eps=getattr(pt_config, "rms_norm_eps", 1e-6),
+        rope_theta=getattr(pt_config, "rope_theta", getattr(pt_config, "rope_base", 1000000.0)),
+        dtype=jnp.bfloat16,
+        max_layer_to_run=27,
+    )
+    text_encoder = FlaxQwen3Model(config=qwen3_config)
+    abstract_q_vars = text_encoder.init(
+        jax.random.PRNGKey(0), jnp.zeros((1, 512), dtype=jnp.int32), jnp.zeros((1, 512), dtype=jnp.int32)
+    )
+    q_params = load_and_convert_qwen3_weights(self.text_encoder_path, abstract_q_vars["params"], qwen3_config)
+
+    tokenizer = Qwen2TokenizerFast.from_pretrained(self.tokenizer_path)
+    scheduler = FlaxFlowMatchScheduler(
+        num_train_timesteps=1000,
+        shift=1.0,
+        sigma_max=1.0,
+        sigma_min=0.001,
+        inverse_timesteps=False,
+        extra_one_step=False,
+        reverse_sigmas=False,
+        use_dynamic_shifting=True,
+        time_shift_type="exponential",
+    )
+
+    # 5. Place parameters on TPU HBM
+    t_params = nnx.state(transformer, nnx.Param)
+    v_params = nnx.state(nnx_vae, nnx.Param)
+
+    t_params = jax.device_put(t_params)
+    v_params = jax.device_put(v_params)
+    q_params = jax.device_put(q_params)
+
+    # 6. Instantiate Unified FlaxFlux2KleinPipeline
+    pipeline = FlaxFlux2KleinPipeline(
+        transformer=transformer,
+        vae=nnx_vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        scheduler=scheduler,
+        config=config,
+        mesh=mesh,
+    )
+
+    # 7. AOT Compile async
+    pipeline.compile_aot_async(
+        params=t_params,
+        vae_params=v_params,
+        qwen3_params=q_params,
+        vae_bn_mean=bn_mean,
+        vae_bn_std=bn_std,
+        batch_size=1,
+        height=512,
+        width=512,
+        images=self.ref_images,
+    )
+
+    # Convert PyTorch initial noise latents to JAX array (shape: 1, 32, 64, 64)
+    initial_latents_jax = jnp.array(raw_latents_pt.detach().float().cpu().numpy())
+
+    # 8. Run pipeline
+    max_logging.log(f" -> Running FlaxFlux2KleinPipeline with {len(self.ref_images)} reference images on TPU...")
+    pipeline(
+        prompt=PROMPT,
+        params=t_params,
+        vae_params=v_params,
+        qwen3_params=q_params,
+        vae_bn_mean=bn_mean,
+        vae_bn_std=bn_std,
+        transformer_shardings=None,
+        vae_shardings=None,
+        qwen3_shardings=None,
+        height=512,
+        width=512,
+        num_inference_steps=4,
+        batch_size=1,
+        images=self.ref_images,
+        use_latents=True,
+        latents=initial_latents_jax,
+        output_dir=self.output_dir,
+        output_name="maxdiffusion_tpu_output.png",
+    )
+
+    maxdiff_img_path = os.path.join(self.output_dir, "maxdiffusion_tpu_output.png")
+    self.assertTrue(os.path.exists(maxdiff_img_path), "MaxDiffusion output image was not saved!")
+    maxdiff_image = Image.open(maxdiff_img_path).convert("RGB")
+
+    max_logging.log("\n" + "=" * 80)
+    max_logging.log("📊 [STEP 3/3] Evaluating End-to-End Parity (SSIM & PSNR)...")
+    max_logging.log("=" * 80)
+
+    diffusers_arr = np.array(diffusers_image).astype(np.uint8)
+    maxdiff_arr = np.array(maxdiff_image).astype(np.uint8)
+
+    self.assertEqual(diffusers_arr.shape, maxdiff_arr.shape)
+
+    ssim_val = ssim(diffusers_arr, maxdiff_arr, channel_axis=-1, data_range=255)
+    mse = np.mean((diffusers_arr.astype(np.float64) - maxdiff_arr.astype(np.float64)) ** 2)
+    psnr_val = 10.0 * np.log10(255.0**2 / (mse + 1e-10))
+
+    max_logging.log(f" -> SSIM (Diffusers CPU vs MaxDiffusion TPU): {ssim_val:.6f}")
+    max_logging.log(f" -> PSNR (Diffusers CPU vs MaxDiffusion TPU): {psnr_val:.2f} dB")
+    max_logging.log(f" -> MSE: {mse:.4f}")
+
+    # Create side-by-side comparison image
+    side_by_side = Image.new("RGB", (1024, 512))
+    side_by_side.paste(diffusers_image, (0, 0))
+    side_by_side.paste(maxdiff_image, (512, 0))
+    comparison_path = os.path.join(self.output_dir, "e2e_parity_diffusers_vs_maxdiffusion.png")
+    side_by_side.save(comparison_path)
+    max_logging.log(f" -> Saved side-by-side comparison to: {comparison_path}")
+
+    self.assertGreaterEqual(ssim_val, 0.75, f"SSIM score {ssim_val:.4f} is below target threshold 0.75!")
+    max_logging.log("🎉 END-TO-END PARITY TEST PASSED! MaxDiffusion matches Diffusers reference!")
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/src/maxdiffusion/tests/flux2klein_image_preprocessing_test.py b/src/maxdiffusion/tests/flux2klein_image_preprocessing_test.py
new file mode 100644
index 000000000..cf57614ba
--- /dev/null
+++ b/src/maxdiffusion/tests/flux2klein_image_preprocessing_test.py
@@ -0,0 +1,129 @@
+"""
+Copyright 2026 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import unittest
+import numpy as np
+import pytest
+from PIL import Image
+
+from maxdiffusion.pipelines.flux.flux2klein_pipeline import FlaxFlux2KleinPipeline
+
+
+class TestFlux2KleinImagePreprocessing(unittest.TestCase):
+  """Unit test suite for canonical reference image preprocessing in FLUX.2-Klein."""
+
+  def test_a_no_duplicate_resize(self):
+    """Test A: Verifies that passing a raw PIL image is preprocessed cleanly without pre-alteration."""
+    img = Image.new("RGB", (512, 512), color=(120, 150, 200))
+    res = FlaxFlux2KleinPipeline.preprocess_reference_image(img)
+    self.assertEqual(res.shape, (1, 3, 512, 512))
+
+  def test_b_aspect_ratio_preserved(self):
+    """Test B: Rectangular images must preserve aspect ratio and not be forced into a square or output shape."""
+    # 1600 x 900 -> total area 1,440,000 > 1024^2 (1,048,576)
+    img = Image.new("RGB", (1600, 900), color=(50, 100, 150))
+    res = FlaxFlux2KleinPipeline.preprocess_reference_image(img)
+    _, _, h, w = res.shape
+    self.assertEqual(w % 16, 0)
+    self.assertEqual(h % 16, 0)
+    self.assertAlmostEqual(w / h, 1600 / 900, delta=0.1)
+    self.assertLessEqual(w * h, 1024 * 1024 + 16 * 1024)
+
+  def test_c_output_size_independence(self):
+    """Test C: Preprocessed reference image must be independent of generation output dimensions."""
+    img = Image.new("RGB", (640, 480), color=(200, 100, 50))
+    res1 = FlaxFlux2KleinPipeline.preprocess_reference_image(img)
+    res2 = FlaxFlux2KleinPipeline.preprocess_reference_image(img)
+    np.testing.assert_array_equal(res1, res2)
+    self.assertEqual(res1.shape, (1, 3, 480, 640))
+
+  def test_d_vae_ready_format(self):
+    """Test D: Verifies output shape is [1, 3, H, W], H%16==0, W%16==0, and values are in [-1, 1]."""
+    img = Image.new("RGB", (345, 678), color=(255, 128, 0))
+    res = FlaxFlux2KleinPipeline.preprocess_reference_image(img)
+    self.assertEqual(res.ndim, 4)
+    self.assertEqual(res.shape[0], 1)
+    self.assertEqual(res.shape[1], 3)
+    self.assertEqual(res.shape[2] % 16, 0)
+    self.assertEqual(res.shape[3] % 16, 0)
+    self.assertTrue(np.all(res >= -1.0) and np.all(res <= 1.0))
+    self.assertEqual(res.dtype, np.float32)
+
+  def test_e_input_validation(self):
+    """Test E: Verifies clear error on unsupported types, small images, and extreme aspect ratios."""
+    # Non-PIL input
+    with self.assertRaises(TypeError):
+      FlaxFlux2KleinPipeline.preprocess_reference_image(np.zeros((512, 512, 3), dtype=np.uint8))
+
+    with self.assertRaises(TypeError):
+      FlaxFlux2KleinPipeline.preprocess_reference_image("path/to/image.png")
+
+    # Too small (min side < 64)
+    small_img = Image.new("RGB", (32, 128))
+    with self.assertRaises(ValueError):
+      FlaxFlux2KleinPipeline.preprocess_reference_image(small_img)
+
+    # Extreme aspect ratio (> 8:1)
+    extreme_img = Image.new("RGB", (1000, 100))
+    with self.assertRaises(ValueError):
+      FlaxFlux2KleinPipeline.preprocess_reference_image(extreme_img)
+
+  def test_f_diffusers_exact_parity(self):
+    """Test F: Direct numerical parity comparison with Diffusers Flux2ImageProcessor."""
+    try:
+      from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor
+    except ImportError:
+      pytest.skip("diffusers not installed or Flux2ImageProcessor not available")
+
+    proc = Flux2ImageProcessor()
+
+    # Test images of various shapes and content
+    test_dims = [(512, 512), (768, 512), (1200, 800), (1600, 900)]
+    for w, h in test_dims:
+      # Generate non-trivial image with color gradient
+      x = np.linspace(0, 255, w, dtype=np.uint8)
+      y = np.linspace(0, 255, h, dtype=np.uint8)
+      xx, yy = np.meshgrid(x, y)
+      arr = np.stack([xx, yy, ((xx + yy) // 2).astype(np.uint8)], axis=-1)
+      pil_img = Image.fromarray(arr)
+
+      # 1. Diffusers preprocessing
+      image_w, image_height = pil_img.size
+      if image_w * image_height > 1024 * 1024:
+        diffusers_pil = proc._resize_to_target_area(pil_img, 1024 * 1024)
+        image_w, image_height = diffusers_pil.size
+      else:
+        diffusers_pil = pil_img
+      image_w = (image_w // 16) * 16
+      image_height = (image_height // 16) * 16
+      diffusers_tensor = proc.preprocess(diffusers_pil, height=image_height, width=image_w, resize_mode="crop")
+      diffusers_np = diffusers_tensor.detach().cpu().float().numpy()
+
+      # 2. MaxDiffusion canonical preprocessing
+      maxdiff_np = FlaxFlux2KleinPipeline.preprocess_reference_image(pil_img)
+
+      # Assert identical shape
+      self.assertEqual(diffusers_np.shape, maxdiff_np.shape)
+
+      # Assert numerical equivalence
+      mae = np.mean(np.abs(diffusers_np - maxdiff_np))
+      max_diff = np.max(np.abs(diffusers_np - maxdiff_np))
+      self.assertLess(mae, 1e-4, f"MAE {mae} exceeded threshold for dims ({w}, {h})")
+      self.assertLess(max_diff, 1e-3, f"Max diff {max_diff} exceeded threshold for dims ({w}, {h})")
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/src/maxdiffusion/tests/flux2klein_kv_pipeline_e2e_test.py b/src/maxdiffusion/tests/flux2klein_kv_pipeline_e2e_test.py
new file mode 100644
index 000000000..5ac5b2409
--- /dev/null
+++ b/src/maxdiffusion/tests/flux2klein_kv_pipeline_e2e_test.py
@@ -0,0 +1,441 @@
+"""
+Copyright 2026 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import gc
+import json
+import os
+import unittest
+import numpy as np
+import pytest
+from PIL import Image
+from skimage.metrics import structural_similarity as ssim
+import torch
+
+import jax
+import jax.numpy as jnp
+import flax
+from flax import nnx
+from flax.linen import partitioning as nn_partitioning
+import flax.linen as nn
+from jax.sharding import Mesh
+from transformers import AutoConfig, Qwen2TokenizerFast
+
+from maxdiffusion import max_utils, pyconfig
+from maxdiffusion.max_utils import create_device_mesh, get_flash_block_sizes
+from maxdiffusion.models.flux.transformers.transformer_flux_flax import NNXFlux2KleinTransformer2DModel
+from maxdiffusion.models.flux.vae.autoencoder_kl_flux2_nnx import (
+    NNXAutoencoderKLFlux2,
+    load_and_convert_flux2klein_nnx_vae_weights,
+)
+from maxdiffusion.models.flux.util import load_and_convert_flux_klein_nnx_weights
+from maxdiffusion.models.qwen3_flax import FlaxQwen3Model, FlaxQwen3Config
+from maxdiffusion.models.qwen3_utils import load_and_convert_qwen3_weights
+from maxdiffusion.schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler
+from maxdiffusion.pipelines.flux.flux2klein_pipeline import FlaxFlux2KleinPipeline
+
+IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true"
+THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+PROMPT = "a vibrant artistic painting combining the dog, car, mountain, and fruit bowl in surreal neon lighting"
+
+
+def compute_psnr(img1: Image.Image, img2: Image.Image) -> float:
+  arr1 = np.array(img1, dtype=np.float64)
+  arr2 = np.array(img2, dtype=np.float64)
+  mse = np.mean((arr1 - arr2) ** 2)
+  if mse == 0:
+    return float("inf")
+  return float(20 * np.log10(255.0 / np.sqrt(mse)))
+
+
+def compute_ssim(img1: Image.Image, img2: Image.Image) -> float:
+  arr1 = np.array(img1.convert("RGB"))
+  arr2 = np.array(img2.convert("RGB"))
+  return float(ssim(arr1, arr2, channel_axis=-1))
+
+
+def find_model_path():
+  if "FLUX2_KLEIN_KV_MODEL_PATH" in os.environ:
+    return os.environ["FLUX2_KLEIN_KV_MODEL_PATH"]
+  hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
+  candidates = [
+      os.path.join(hf_home, "hub/models--black-forest-labs--FLUX.2-klein-9B-KV/snapshots"),
+      os.path.join(hf_home, "hub/models--black-forest-labs--FLUX.2-klein-9b-kv/snapshots"),
+      os.path.join(hf_home, "hub/models--black-forest-labs--FLUX.2-klein-9B/snapshots"),
+      "/mnt/hyperdisk_weights/hub/flux2klein-9b-kv",
+      "/mnt/data/models/flux2klein-9b-kv",
+  ]
+  for c in candidates:
+    if os.path.exists(c):
+      if "snapshots" in c:
+        snaps = os.listdir(c)
+        if snaps:
+          return os.path.join(c, snaps[0])
+      else:
+        return c
+  return "black-forest-labs/FLUX.2-klein-9B-KV"
+
+
+class TestFlux2KleinKVPipelineE2EBF16Parity(unittest.TestCase):
+  """End-to-end parity test comparing Diffusers Flux2KleinKVPipeline vs MaxDiffusion FlaxFlux2KleinPipeline (use_kv=True) in bfloat16."""
+
+  @classmethod
+  def setUpClass(cls):
+    cls.model_path = find_model_path()
+    cls.work_dir = "/tmp/flux2klein_kv_e2e"
+    os.makedirs(cls.work_dir, exist_ok=True)
+
+    cls.height = 256
+    cls.width = 256
+    cls.num_inference_steps = 4
+    cls.seed = int(os.getenv("FLUX2_KLEIN_E2E_SEED", "42"))
+
+    # 1. Load 4 real reference images (256x256)
+    ref_dir = os.path.join(THIS_DIR, "images", "flux2klein")
+    cls.ref_images = []
+    if os.path.exists(ref_dir):
+      for i in range(4):
+        p = os.path.join(ref_dir, f"ref_image_{i}.png")
+        if os.path.exists(p):
+          cls.ref_images.append(Image.open(p).convert("RGB").resize((256, 256), Image.Resampling.BICUBIC))
+
+    if len(cls.ref_images) < 4:
+      colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0)]
+      for i, c in enumerate(colors):
+        arr = np.full((256, 256, 3), c, dtype=np.uint8)
+        cls.ref_images.append(Image.fromarray(arr))
+
+    # 2. Generate shared starting noise latents (1, 32, 32, 32)
+    rng = np.random.RandomState(cls.seed)
+    latents_unpacked = rng.randn(1, 32, cls.height // 8, cls.width // 8).astype(np.float32)
+    cls.latents_unpacked_jax = jnp.array(latents_unpacked)
+
+    # Prepare packed latents for PyTorch: (1, 32, H/16, 2, W/16, 2) -> permute(0, 1, 3, 5, 2, 4) -> reshape(1, 128, H/16, W/16)
+    latents_unpacked_pt = torch.from_numpy(latents_unpacked)
+    latents_pt_packed = latents_unpacked_pt.view(1, 32, cls.height // 16, 2, cls.width // 16, 2)
+    latents_pt_packed = latents_pt_packed.permute(0, 1, 3, 5, 2, 4)
+    cls.latents_pt_packed = latents_pt_packed.reshape(1, 128, cls.height // 16, cls.width // 16)
+
+  @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Requires TPU and full FLUX.2-Klein 9B weights")
+  def test_flux2klein_kv_pipeline_bf16_parity(self):
+    """Executes PyTorch Diffusers Flux2KleinKVPipeline in BF16 and MaxDiffusion FlaxFlux2KleinPipeline in BF16 (use_kv=True) and compares outputs."""
+    print("\n" + "=" * 80)
+    print("🚀 FLUX.2-KLEIN-9B KV CACHE END-TO-END BF16 PARITY TEST")
+    print("=" * 80)
+    print(f"Model Path:          {self.model_path}")
+    print(f"Prompt:              '{PROMPT}'")
+    print(f"Number of Ref Images:{len(self.ref_images)} (256x256)")
+    print(f"Target Resolution:   {self.width}x{self.height}")
+    print(f"Inference Steps:     {self.num_inference_steps}")
+
+    # =========================================================================
+    # LEG 1: PyTorch Diffusers Flux2KleinKVPipeline in BF16
+    # =========================================================================
+    print("\n" + "-" * 80)
+    print("🎬 LEG 1: Running PyTorch Diffusers Flux2KleinKVPipeline (bfloat16 on CPU)...")
+    print("-" * 80)
+    from diffusers import Flux2KleinKVPipeline
+
+    pipe_pt = Flux2KleinKVPipeline.from_pretrained(
+        self.model_path,
+        torch_dtype=torch.bfloat16,
+        local_files_only=True,
+    )
+    pipe_pt.to("cpu")
+
+    with torch.no_grad():
+      pt_out = pipe_pt(
+          prompt=PROMPT,
+          image=self.ref_images,
+          latents=self.latents_pt_packed.to(torch.bfloat16),
+          num_inference_steps=self.num_inference_steps,
+          height=self.height,
+          width=self.width,
+          output_type="pil",
+      ).images[0]
+
+    pt_output_path = os.path.join(self.work_dir, "pt_bf16_output.png")
+    pt_out.save(pt_output_path)
+    print(f" -> Saved PyTorch Diffusers BF16 output: {pt_output_path}")
+
+    del pipe_pt
+    gc.collect()
+
+    # =========================================================================
+    # LEG 2: MaxDiffusion FlaxFlux2KleinPipeline in BF16 (use_kv=True)
+    # =========================================================================
+    print("\n" + "-" * 80)
+    print("🎬 LEG 2: Running MaxDiffusion FlaxFlux2KleinPipeline with use_kv=True (bfloat16 on TPU)...")
+    print("-" * 80)
+
+    # 1. Device mesh setup
+    active_devices = jax.devices()
+    active_device_count = len(active_devices)
+
+    pyconfig._config = None
+    pyconfig.config = None
+    config_path = os.path.join(THIS_DIR, "..", "configs", "base_flux2klein_9B.yml")
+    args = [
+        None,
+        config_path,
+        "run_name=e2e_kv_parity_test",
+        f"output_dir={self.work_dir}",
+        f"per_device_batch_size={1.0 / active_device_count}",
+        f"height={self.height}",
+        f"width={self.width}",
+        f"num_inference_steps={self.num_inference_steps}",
+        f"seed={self.seed}",
+        "use_kv=True",
+        "weights_dtype=bfloat16",
+        "activations_dtype=bfloat16",
+        "precision=DEFAULT",
+        "attention=tokamax_flash",
+        'flash_block_sizes={"block_q": 512, "block_kv": 512, "block_kv_compute": 512}',
+        "text_encoder_attention=dot_product",
+    ]
+    pyconfig.initialize(args, unittest=True)
+    config = pyconfig.config
+
+    if active_device_count > 1:
+      pyconfig._config.keys["ici_tensor_parallelism"] = active_device_count
+      pyconfig._config.keys["ici_data_parallelism"] = 1
+      pyconfig._config.keys["ici_fsdp_parallelism"] = 1
+      pyconfig._config.keys["ici_context_parallelism"] = 1
+
+    pyconfig._config.keys["flash_block_sizes"] = {
+        "block_q": 512,
+        "block_kv": 512,
+        "block_kv_compute": 512,
+    }
+
+    devices_array = create_device_mesh(config, devices=active_devices)
+    mesh = Mesh(devices_array, config.mesh_axes)
+
+    # 2. Text Encoder & Tokenizer
+    text_encoder_path = os.path.join(self.model_path, "text_encoder")
+    tokenizer_path = os.path.join(self.model_path, "tokenizer")
+    pt_config = AutoConfig.from_pretrained(text_encoder_path)
+    qwen3_config = FlaxQwen3Config(
+        vocab_size=pt_config.vocab_size,
+        hidden_size=pt_config.hidden_size,
+        intermediate_size=pt_config.intermediate_size,
+        num_hidden_layers=pt_config.num_hidden_layers,
+        num_attention_heads=pt_config.num_attention_heads,
+        num_key_value_heads=getattr(pt_config, "num_key_value_heads", pt_config.num_attention_heads),
+        max_position_embeddings=getattr(pt_config, "max_position_embeddings", 32768),
+        rms_norm_eps=getattr(pt_config, "rms_norm_eps", 1e-6),
+        rope_theta=getattr(pt_config, "rope_theta", getattr(pt_config, "rope_base", 1000000.0)),
+        dtype=jnp.bfloat16,
+        attention_kernel="dot_product",
+        mesh=mesh,
+        max_layer_to_run=getattr(config, "text_encoder_max_layer", 27),
+        is_causal=True,
+    )
+    text_encoder = FlaxQwen3Model(qwen3_config)
+    tokenizer = Qwen2TokenizerFast.from_pretrained(tokenizer_path)
+
+    # 3. NNX Transformer
+    transformer_path = os.path.join(self.model_path, "transformer")
+    transformer_config_json = os.path.join(transformer_path, "config.json")
+    transformer_pt_cfg = {}
+    if os.path.exists(transformer_config_json):
+      with open(transformer_config_json, "r") as f:
+        transformer_pt_cfg = json.load(f)
+
+    num_double_layers = transformer_pt_cfg.get("num_layers", 8)
+    depth = transformer_pt_cfg.get("num_single_layers", 24)
+    num_attention_heads = transformer_pt_cfg.get("num_attention_heads", 32)
+
+    transformer = NNXFlux2KleinTransformer2DModel(
+        rngs=nnx.Rngs(0),
+        in_channels=128,
+        num_layers=num_double_layers,
+        num_single_layers=depth,
+        attention_head_dim=128,
+        num_attention_heads=num_attention_heads,
+        joint_attention_dim=3 * pt_config.hidden_size,
+        pooled_projection_dim=768,
+        guidance_embeds=transformer_pt_cfg.get("guidance_embeds", False),
+        axes_dim=(32, 32, 32, 32),
+        theta=2000.0,
+        mlp_ratio=3.0,
+        attention_kernel=config.attention,
+        flash_min_seq_length=512,
+        flash_block_sizes=get_flash_block_sizes(config),
+        mesh=mesh,
+        dtype=jnp.bfloat16,
+        weights_dtype=jnp.bfloat16,
+        scale_shift_order="scale_shift",
+        use_base2_exp=True,
+    )
+
+    # 4. NNX VAE
+    vae_path = os.path.join(self.model_path, "vae", "diffusion_pytorch_model.safetensors")
+    if not os.path.exists(vae_path):
+      vae_path = os.path.join(self.model_path, "vae")
+    vae = NNXAutoencoderKLFlux2(
+        in_channels=3,
+        out_channels=3,
+        latent_channels=32,
+        block_out_channels=(128, 256, 512, 512),
+        layers_per_block=2,
+        norm_num_groups=32,
+        dtype=jnp.bfloat16,
+        param_dtype=jnp.bfloat16,
+    )
+
+    # 5. Extract mesh shardings for all models
+    abstract_transformer_state = nnx.state(transformer, nnx.Param)
+    abstract_vae_state = nnx.state(vae, nnx.Param)
+
+    def qwen3_init_fn():
+      return text_encoder.init(
+          jax.random.PRNGKey(0), jnp.zeros((1, 512), dtype=jnp.int32), jnp.zeros((1, 512), dtype=jnp.int32)
+      )
+
+    with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+      logical_transformer_specs = nnx.get_partition_spec(abstract_transformer_state)
+      logical_vae_specs = nnx.get_partition_spec(abstract_vae_state)
+      abstract_qwen3_vars = jax.eval_shape(qwen3_init_fn)
+      logical_qwen3_specs = nn.get_partition_spec(abstract_qwen3_vars)
+
+      transformer_shardings = nn.logical_to_mesh_sharding(logical_transformer_specs, mesh, config.logical_axis_rules)
+      vae_shardings = nn.logical_to_mesh_sharding(logical_vae_specs, mesh, config.logical_axis_rules)
+      qwen3_shardings = flax.core.freeze(
+          nn.logical_to_mesh_sharding(logical_qwen3_specs, mesh, config.logical_axis_rules)["params"]
+      )
+
+    # 6. Load weights on Host CPU and shard across TPU HBM
+    cpu_device = jax.local_devices(backend="cpu")[0]
+    with jax.default_device(cpu_device):
+      t_params = load_and_convert_flux_klein_nnx_weights(
+          transformer_path,
+          abstract_transformer_state,
+          num_double_layers=num_double_layers,
+          num_single_layers=depth,
+          dtype=jnp.bfloat16,
+      )
+      vae_bn_mean, vae_bn_std = load_and_convert_flux2klein_nnx_vae_weights(vae_path, vae, dtype=jnp.bfloat16)
+      v_params = nnx.state(vae, nnx.Param)
+
+      def unbox_fn(x):
+        import flax.linen.spmd as flax_spmd
+
+        return x.unbox() if isinstance(x, flax_spmd.LogicallyPartitioned) else x
+
+      qwen3_params_template = jax.tree_util.tree_map(
+          unbox_fn, abstract_qwen3_vars["params"], is_leaf=lambda k: hasattr(k, "unbox")
+      )
+      qwen3_params_template = flax.core.unfreeze(qwen3_params_template)
+      q_params = load_and_convert_qwen3_weights(text_encoder_path, qwen3_params_template, qwen3_config)
+      q_params = flax.core.freeze(q_params)
+
+    # Shard onto TPU HBM
+    print(" -> Sharding parameters across TPU HBM...")
+    with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+      t_params = jax.tree_util.tree_map(max_utils.device_put_replicated, t_params, transformer_shardings)
+      v_params = jax.tree_util.tree_map(max_utils.device_put_replicated, v_params, vae_shardings)
+      nnx.update(vae, v_params)
+      q_params = jax.tree_util.tree_map(max_utils.device_put_replicated, q_params, qwen3_shardings)
+
+    # 7. Scheduler
+    scheduler = FlaxFlowMatchScheduler(
+        num_train_timesteps=1000,
+        shift=1.0,
+        sigma_max=1.0,
+        sigma_min=0.001,
+        inverse_timesteps=False,
+        extra_one_step=False,
+        reverse_sigmas=False,
+        use_dynamic_shifting=True,
+        time_shift_type="exponential",
+    )
+
+    # 8. Pipeline instantiation & execution
+    pipeline = FlaxFlux2KleinPipeline(
+        transformer=transformer,
+        vae=vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        scheduler=scheduler,
+        config=config,
+        mesh=mesh,
+    )
+
+    pipeline.compile_aot_async(
+        params=t_params,
+        vae_params=v_params,
+        qwen3_params=q_params,
+        vae_bn_mean=vae_bn_mean,
+        vae_bn_std=vae_bn_std,
+        batch_size=1,
+        height=self.height,
+        width=self.width,
+        images=self.ref_images,
+        use_kv=True,
+    )
+
+    jax_output_name = "jax_bf16_output.png"
+    pipeline(
+        prompt=PROMPT,
+        params=t_params,
+        vae_params=v_params,
+        qwen3_params=q_params,
+        vae_bn_mean=vae_bn_mean,
+        vae_bn_std=vae_bn_std,
+        transformer_shardings=transformer_shardings,
+        vae_shardings=vae_shardings,
+        qwen3_shardings=qwen3_shardings,
+        height=self.height,
+        width=self.width,
+        num_inference_steps=self.num_inference_steps,
+        batch_size=1,
+        images=self.ref_images,
+        use_latents=True,
+        latents=self.latents_unpacked_jax,
+        use_kv=True,
+        output_dir=self.work_dir,
+        output_name=jax_output_name,
+    )
+
+    jax_output_path = os.path.join(self.work_dir, jax_output_name)
+    self.assertTrue(os.path.exists(jax_output_path), f"JAX output image not found at {jax_output_path}")
+    print(f" -> Found MaxDiffusion JAX BF16 output: {jax_output_path}")
+
+    # =========================================================================
+    # LEG 3: Compute Parity Metrics (SSIM & PSNR)
+    # =========================================================================
+    print("\n" + "=" * 80)
+    print("📊 CROSS-FRAMEWORK BF16 PARITY EVALUATION REPORT")
+    print("=" * 80)
+
+    img_pt = Image.open(pt_output_path).convert("RGB")
+    img_jax = Image.open(jax_output_path).convert("RGB")
+
+    score_ssim = compute_ssim(img_jax, img_pt)
+    score_psnr = compute_psnr(img_jax, img_pt)
+
+    print(f" -> Structural Similarity (SSIM): {score_ssim:.6f}")
+    print(f" -> Peak Signal-to-Noise Ratio (PSNR): {score_psnr:.2f} dB")
+    print("=" * 80)
+
+    self.assertGreaterEqual(
+        score_ssim, 0.70, f"End-to-End BF16 SSIM {score_ssim:.6f} is below the required acceptance threshold of 0.70"
+    )
+    print("✅ End-to-End BF16 Parity Test PASSED successfully!\n")
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/src/maxdiffusion/tests/generate_flux2klein_e2e_test.py b/src/maxdiffusion/tests/generate_flux2klein_e2e_test.py
index 38f800e4c..f070f84e7 100644
--- a/src/maxdiffusion/tests/generate_flux2klein_e2e_test.py
+++ b/src/maxdiffusion/tests/generate_flux2klein_e2e_test.py
@@ -15,6 +15,7 @@
 """
 
 import os
+import sys
 import subprocess
 import gc
 import numpy as np
@@ -22,10 +23,9 @@
 from PIL import Image
 from skimage.metrics import structural_similarity as ssim
 
-# Set HF_HOME cache path early
+# Ensure default HF_HOME is defined if not set by environment
 if not os.environ.get("HF_HOME"):
-  if os.path.exists("/mnt/data/hf_cache"):
-    os.environ["HF_HOME"] = "/mnt/data/hf_cache"
+  os.environ["HF_HOME"] = os.path.expanduser("~/.cache/huggingface")
 
 
 def compute_psnr(img1, img2):
@@ -43,13 +43,34 @@ def compute_ssim(img1, img2):
   return ssim(img1_gray, img2_gray)
 
 
+def get_model_snapshot_dir(model_id: str) -> str:
+  """Locates the snapshot directory for the given model_id strictly under HF_HOME."""
+  try:
+    from huggingface_hub import snapshot_download
+
+    return snapshot_download(repo_id=model_id, local_files_only=True)
+  except Exception:
+    pass
+
+  hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
+  escaped_id = model_id.replace("/", "--")
+  cache_dir = os.path.join(hf_home, "hub", f"models--{escaped_id}", "snapshots")
+  if not os.path.exists(cache_dir):
+    cache_dir = os.path.join(hf_home, f"models--{escaped_id}", "snapshots")
+
+  if not os.path.exists(cache_dir):
+    raise FileNotFoundError(f"Hugging Face cache directory not found for '{model_id}' under HF_HOME ({hf_home}).")
+
+  snapshots = [s for s in os.listdir(cache_dir) if not s.startswith(".")]
+  if not snapshots:
+    raise FileNotFoundError(f"No snapshot directory found for '{model_id}' in {cache_dir}.")
+
+  return os.path.join(cache_dir, snapshots[0])
+
+
 def run_pytorch_pipeline(model_id, prompt, batch_size, width, height, num_inference_steps, seed, latents_pt_packed, prefix):
   # Locate cached model files
-  cache_dir = f"/mnt/data/hf_cache/hub/models--{model_id.replace('/', '--')}/snapshots"
-  if not os.path.exists(cache_dir):
-    raise FileNotFoundError(f"Hugging Face cache directory not found: {cache_dir}")
-  snapshots = os.listdir(cache_dir)
-  snapshot_dir = os.path.join(cache_dir, snapshots[0])
+  snapshot_dir = get_model_snapshot_dir(model_id)
   print(f"\n[PyTorch] Loading '{model_id}' weights from: {snapshot_dir}")
 
   from diffusers.pipelines.flux2.pipeline_flux2_klein import Flux2KleinPipeline
@@ -153,7 +174,7 @@ def main():
   # 2. Run JAX 4B (via generate_flux2klein.py)
   print("\n[JAX 4B] Executing pipeline script generate_flux2klein.py...")
   cmd_jax_4b = [
-      "python3",
+      sys.executable,
       "src/maxdiffusion/generate_flux2klein.py",
       "src/maxdiffusion/configs/base_flux2klein.yml",
       "skip_jax_distributed_system=True",
@@ -197,7 +218,7 @@ def main():
 
   print("\n[JAX 9B] Executing pipeline script generate_flux2klein.py...")
   cmd_jax_9b = [
-      "python3",
+      sys.executable,
       "src/maxdiffusion/generate_flux2klein.py",
       "src/maxdiffusion/configs/base_flux2klein_9B.yml",
       "skip_jax_distributed_system=True",
diff --git a/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py b/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py
index c61db29f7..6d6664315 100644
--- a/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py
+++ b/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py
@@ -37,11 +37,11 @@ class GenerateFlux2KleinSmokeTest(unittest.TestCase):
   @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions (requires TPU HBM)")
   def test_flux2klein_4b_smoke(self):
     """End-to-end smoke test for Flux.2-klein-4B image generation at 1024x1024."""
-    ref_path = os.path.join(THIS_DIR, "images", "ref_flux2klein_4b.png")
+    ref_path = os.path.join(THIS_DIR, "images", "flux2klein", "ref_flux2klein_4b.png")
     self.assertTrue(os.path.exists(ref_path), f"Reference image not found: {ref_path}")
     base_image = np.array(Image.open(ref_path)).astype(np.uint8)
 
-    output_dir = "/mnt/data/smoke_test_4b" if os.path.exists("/mnt/data") else "/tmp/smoke_test_4b"
+    output_dir = "/tmp/smoke_test_4b"
     os.makedirs(output_dir, exist_ok=True)
     out_path = os.path.join(output_dir, "flux2klein_generated_image.png")
     if os.path.exists(out_path):
@@ -77,16 +77,16 @@ def test_flux2klein_4b_smoke(self):
     self.assertEqual(base_image.shape, test_image.shape)
     ssim_compare = ssim(base_image, test_image, channel_axis=-1, data_range=255)
     print(f"\n[SMOKE TEST 4B] SSIM Score: {ssim_compare:.6f}")
-    self.assertGreaterEqual(ssim_compare, 0.75)
+    self.assertGreaterEqual(ssim_compare, 0.8)
 
   @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions (requires TPU HBM)")
   def test_flux2klein_9b_smoke(self):
     """End-to-end smoke test for Flux.2-klein-9B image generation at 1024x1024."""
-    ref_path = os.path.join(THIS_DIR, "images", "ref_flux2klein_9b.png")
+    ref_path = os.path.join(THIS_DIR, "images", "flux2klein", "ref_flux2klein_9b.png")
     self.assertTrue(os.path.exists(ref_path), f"Reference image not found: {ref_path}")
     base_image = np.array(Image.open(ref_path)).astype(np.uint8)
 
-    output_dir = "/mnt/data/smoke_test_9b" if os.path.exists("/mnt/data") else "/tmp/smoke_test_9b"
+    output_dir = "/tmp/smoke_test_9b"
     os.makedirs(output_dir, exist_ok=True)
     out_path = os.path.join(output_dir, "flux2klein_generated_image.png")
     if os.path.exists(out_path):
@@ -124,6 +124,156 @@ def test_flux2klein_9b_smoke(self):
     print(f"\n[SMOKE TEST 9B] SSIM Score: {ssim_compare:.6f}")
     self.assertGreaterEqual(ssim_compare, 0.8)
 
+  @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions (requires TPU HBM)")
+  def test_flux2klein_4b_image_edit_smoke(self):
+    """End-to-end smoke test for Flux.2-klein-4B image editing at 512x512."""
+    ref_gold_path = os.path.join(THIS_DIR, "images", "flux2klein", "ref_flux2klein_4b_image_edit.png")
+    self.assertTrue(os.path.exists(ref_gold_path), f"Golden reference image not found: {ref_gold_path}")
+    base_image = np.array(Image.open(ref_gold_path)).astype(np.uint8)
+
+    input_img_path = os.path.join(THIS_DIR, "images", "flux2klein", "ref_flux2klein_4b.png")
+    self.assertTrue(os.path.exists(input_img_path), f"Input reference image not found: {input_img_path}")
+
+    output_dir = "/tmp/smoke_test_image_edit_4b"
+    os.makedirs(output_dir, exist_ok=True)
+    out_path = os.path.join(output_dir, "flux2klein_generated_image.png")
+    if os.path.exists(out_path):
+      os.remove(out_path)
+
+    pyconfig._config = None
+    pyconfig.config = None
+    args = [
+        None,
+        os.path.join(THIS_DIR, "..", "configs", "base_flux2klein.yml"),
+        "run_name=smoke_test_image_edit_4b",
+        f"output_dir={output_dir}",
+        "jax_cache_dir=/tmp/cache_dir",
+        f"image_paths=['{input_img_path}']",
+        "prompt=change the lighting to evening",
+        "height=512",
+        "width=512",
+        f"per_device_batch_size={1.0 / jax.device_count()}",
+        "seed=42",
+        "weights_dtype=bfloat16",
+        "activations_dtype=bfloat16",
+        "precision=DEFAULT",
+        "num_reps=5",
+        "text_encoder_attention=dot_product",
+    ]
+
+    generate_flux2klein.main(args)
+
+    rep_out_path = os.path.join(output_dir, "rep_1_flux2klein_generated_image.png")
+    final_out_path = rep_out_path if os.path.exists(rep_out_path) else out_path
+    self.assertTrue(os.path.exists(final_out_path), "Smoke test 4B image edit failed to produce output image!")
+    test_image = np.array(Image.open(final_out_path)).astype(np.uint8)
+
+    self.assertEqual(base_image.shape, test_image.shape)
+    ssim_compare = ssim(base_image, test_image, channel_axis=-1, data_range=255)
+    print(f"\n[SMOKE TEST 4B IMAGE EDIT] SSIM Score: {ssim_compare:.6f}")
+    self.assertGreaterEqual(ssim_compare, 0.8)
+
+  @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions (requires TPU HBM)")
+  def test_flux2klein_9b_image_edit_smoke(self):
+    """End-to-end smoke test for Flux.2-klein-9B image editing at 512x512."""
+    ref_gold_path = os.path.join(THIS_DIR, "images", "flux2klein", "ref_flux2klein_9b_image_edit.png")
+    self.assertTrue(os.path.exists(ref_gold_path), f"Golden reference image not found: {ref_gold_path}")
+    base_image = np.array(Image.open(ref_gold_path)).astype(np.uint8)
+
+    input_img_path = os.path.join(THIS_DIR, "images", "flux2klein", "ref_flux2klein_4b.png")
+    self.assertTrue(os.path.exists(input_img_path), f"Input reference image not found: {input_img_path}")
+
+    output_dir = "/tmp/smoke_test_image_edit_9b"
+    os.makedirs(output_dir, exist_ok=True)
+    out_path = os.path.join(output_dir, "flux2klein_generated_image.png")
+    if os.path.exists(out_path):
+      os.remove(out_path)
+
+    pyconfig._config = None
+    pyconfig.config = None
+    args = [
+        None,
+        os.path.join(THIS_DIR, "..", "configs", "base_flux2klein_9B.yml"),
+        "run_name=smoke_test_image_edit_9b",
+        f"output_dir={output_dir}",
+        "jax_cache_dir=/tmp/cache_dir",
+        f"image_paths=['{input_img_path}']",
+        "prompt=change the lighting to evening",
+        "height=512",
+        "width=512",
+        f"per_device_batch_size={1.0 / jax.device_count()}",
+        "seed=42",
+        "weights_dtype=bfloat16",
+        "activations_dtype=bfloat16",
+        "precision=DEFAULT",
+        "num_reps=5",
+        "text_encoder_attention=dot_product",
+    ]
+
+    generate_flux2klein.main(args)
+
+    rep_out_path = os.path.join(output_dir, "rep_1_flux2klein_generated_image.png")
+    final_out_path = rep_out_path if os.path.exists(rep_out_path) else out_path
+    self.assertTrue(os.path.exists(final_out_path), "Smoke test 9B image edit failed to produce output image!")
+    test_image = np.array(Image.open(final_out_path)).astype(np.uint8)
+
+    self.assertEqual(base_image.shape, test_image.shape)
+    ssim_compare = ssim(base_image, test_image, channel_axis=-1, data_range=255)
+    print(f"\n[SMOKE TEST 9B IMAGE EDIT] SSIM Score: {ssim_compare:.6f}")
+    self.assertGreaterEqual(ssim_compare, 0.8)
+
+  @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions (requires TPU HBM)")
+  def test_flux2klein_9b_kv_image_edit_smoke(self):
+    """End-to-end smoke test for Flux.2-klein-9B KV cache image editing at 512x512."""
+    ref_gold_path = os.path.join(THIS_DIR, "images", "flux2klein", "ref_flux2klein_9b_kv_image_edit.png")
+    self.assertTrue(os.path.exists(ref_gold_path), f"Golden reference image not found: {ref_gold_path}")
+    base_image = np.array(Image.open(ref_gold_path)).astype(np.uint8)
+
+    input_img_path = os.path.join(THIS_DIR, "images", "flux2klein", "ref_flux2klein_4b.png")
+    self.assertTrue(os.path.exists(input_img_path), f"Input reference image not found: {input_img_path}")
+
+    output_dir = "/tmp/smoke_test_kv_image_edit_9b"
+    os.makedirs(output_dir, exist_ok=True)
+    out_path = os.path.join(output_dir, "flux2klein_generated_image.png")
+    if os.path.exists(out_path):
+      os.remove(out_path)
+
+    pyconfig._config = None
+    pyconfig.config = None
+    args = [
+        None,
+        os.path.join(THIS_DIR, "..", "configs", "base_flux2klein_9B.yml"),
+        "run_name=smoke_test_kv_image_edit_9b",
+        f"output_dir={output_dir}",
+        "jax_cache_dir=/tmp/cache_dir",
+        f"image_paths=['{input_img_path}']",
+        "prompt=change the lighting to evening",
+        "height=512",
+        "width=512",
+        f"per_device_batch_size={1.0 / jax.device_count()}",
+        "seed=42",
+        "use_kv=True",
+        "attention=tokamax_flash",
+        'flash_block_sizes={"block_q": 512, "block_kv": 512, "block_kv_compute": 512}',
+        "weights_dtype=bfloat16",
+        "activations_dtype=bfloat16",
+        "precision=DEFAULT",
+        "num_reps=5",
+        "text_encoder_attention=dot_product",
+    ]
+
+    generate_flux2klein.main(args)
+
+    rep_out_path = os.path.join(output_dir, "rep_1_flux2klein_generated_image.png")
+    final_out_path = rep_out_path if os.path.exists(rep_out_path) else out_path
+    self.assertTrue(os.path.exists(final_out_path), "Smoke test 9B KV image edit failed to produce output image!")
+    test_image = np.array(Image.open(final_out_path)).astype(np.uint8)
+
+    self.assertEqual(base_image.shape, test_image.shape)
+    ssim_compare = ssim(base_image, test_image, channel_axis=-1, data_range=255)
+    print(f"\n[SMOKE TEST 9B KV IMAGE EDIT] SSIM Score: {ssim_compare:.6f}")
+    self.assertGreaterEqual(ssim_compare, 0.8)
+
 
 if __name__ == "__main__":
   unittest.main()
diff --git a/src/maxdiffusion/tests/images/ref_flux2klein_4b.png b/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_4b.png
similarity index 100%
rename from src/maxdiffusion/tests/images/ref_flux2klein_4b.png
rename to src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_4b.png
diff --git a/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_4b_image_edit.png b/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_4b_image_edit.png
new file mode 100644
index 000000000..f24599b58
Binary files /dev/null and b/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_4b_image_edit.png differ
diff --git a/src/maxdiffusion/tests/images/ref_flux2klein_9b.png b/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b.png
similarity index 100%
rename from src/maxdiffusion/tests/images/ref_flux2klein_9b.png
rename to src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b.png
diff --git a/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b_image_edit.png b/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b_image_edit.png
new file mode 100644
index 000000000..7d938fb17
Binary files /dev/null and b/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b_image_edit.png differ
diff --git a/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b_kv_image_edit.png b/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b_kv_image_edit.png
new file mode 100644
index 000000000..992159ab9
Binary files /dev/null and b/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b_kv_image_edit.png differ
diff --git a/src/maxdiffusion/tests/images/flux2klein/ref_image_0.png b/src/maxdiffusion/tests/images/flux2klein/ref_image_0.png
new file mode 100644
index 000000000..476ba5984
Binary files /dev/null and b/src/maxdiffusion/tests/images/flux2klein/ref_image_0.png differ
diff --git a/src/maxdiffusion/tests/images/flux2klein/ref_image_1.png b/src/maxdiffusion/tests/images/flux2klein/ref_image_1.png
new file mode 100644
index 000000000..d14acbe2a
Binary files /dev/null and b/src/maxdiffusion/tests/images/flux2klein/ref_image_1.png differ
diff --git a/src/maxdiffusion/tests/images/flux2klein/ref_image_2.png b/src/maxdiffusion/tests/images/flux2klein/ref_image_2.png
new file mode 100644
index 000000000..cb379a6e8
Binary files /dev/null and b/src/maxdiffusion/tests/images/flux2klein/ref_image_2.png differ
diff --git a/src/maxdiffusion/tests/images/flux2klein/ref_image_3.png b/src/maxdiffusion/tests/images/flux2klein/ref_image_3.png
new file mode 100644
index 000000000..bffdbe23c
Binary files /dev/null and b/src/maxdiffusion/tests/images/flux2klein/ref_image_3.png differ