Adding Image Editing to Flux2.Klein Models

src/maxdiffusion/configs/base_flux2klein.yml

@@ -40,6 +40,8 @@ max_sequence_length: 512
 time_shift: True
 base_shift: 0.5
 max_shift: 1.15
+image_paths: []
+use_base2_exp: True
 
 
 unet_checkpoint: ''

src/maxdiffusion/configs/base_flux2klein_9B.yml

@@ -40,6 +40,8 @@ max_sequence_length: 512
 time_shift: True
 base_shift: 0.5
 max_shift: 1.15
+image_paths: []
+use_base2_exp: True
 
 
 unet_checkpoint: ''

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
@@ -321,6 +322,7 @@ 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
@@ -464,6 +466,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").resize((config.width, config.height), Image.Resampling.BICUBIC)
+            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 +531,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 +559,7 @@ def unbox_fn(x):
         batch_size=config.batch_size,
         height=config.height,
         width=config.width,
+        images=images,
     )
 
     max_logging.log("\n" + "=" * 80)
@@ -547,14 +579,16 @@ 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,
         output_name="flux2klein_warmup.png",
         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 +623,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 +644,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 +653,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 +685,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)
 

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:

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

@@ -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")
@@ -1504,6 +1506,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 +1525,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 +1564,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 +1621,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(
@@ -1703,6 +1709,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 +1760,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__(
@@ -1834,6 +1842,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 +1923,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 +1945,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)
         ]