Adding Image Editing by amepas · Pull Request #460 · AI-Hypercomputer/maxdiffusion

amepas · 2026-08-17T21:51:46Z

Early version of adding image editing support.

KV cache not supported currently.

…weight loading, causal splash attention, and end-to-end optimizations - Implemented Flax NNX Transformer architecture (NNXFlux2KleinTransformer2DModel) with support for both 4B (5 double / 20 single layers) and 9B (8 double / 24 single layers) configurations. - Integrated Flax Qwen3 text encoder with 3-layer intermediate hidden states extraction (layers 9, 18, 27), custom causal splash attention, and proper sharding constraints. - Implemented FlaxAutoencoderKL VAE decoder with fused batch normalization unscaling and channel re-layout. - Added fused end-to-end denoising loop scan with Flow Match Euler scheduler. - Added concurrent AOT XLA compilation across Qwen3, Flux transformer, and VAE. - Implemented fast host-memory streaming weight converter for safetensors shards directly into NNX State PyTree. - Optimized splash attention block sizes and Ulysses context parallelism sharding. - Added comprehensive unit tests (nnx_flux2klein_test.py) and end-to-end smoke test suite (generate_flux2klein_smoke_test.py).

… & 9B) - Implement NNXAutoencoderKLFlux2 (pure Flax NNX VAE encoder and decoder) with zero-copy weight loading and verified float32 parity - Implement multi-image conditioning helpers (prepare_multi_image_ids, prepare_image_latents, patchify/unpatchify) - Implement FlaxFlux2KleinImageEditPipeline for end-to-end multi-image editing inference - Implement generate_flux2klein_image_edit.py CLI runner and default YAML configs for 4B and 9B - Add comprehensive unit tests and parity test suite across all 7 implementation phases

…o single pipeline and CLI - Refactored FlaxFlux2KleinPipeline to conditionally handle multi-image editing when images are provided, eliminating duplicate pipeline definitions. - Unified generate_flux2klein.py and base configs (base_flux2klein.yml, base_flux2klein_9B.yml) to accept image_paths for multi-image conditioning. - Added test_flux2klein_4b_image_edit_smoke to generate_flux2klein_smoke_test.py. - Added test_flux2klein_image_edit_e2e_parity.py for end-to-end parity verification against PyTorch Diffusers reference.

github-actions · 2026-08-17T21:52:22Z

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

gemini-code-assist

Code Review

This pull request introduces a unified end-to-end multi-image editing pipeline for FLUX.2-Klein (4B and 9B models) in pure Flax NNX, adding new configuration files, a dedicated generation script, and the pipeline implementation, alongside supporting tests and model updates. Feedback on these changes focuses on improving robustness and portability: specifically, making image normalization more reliable by checking for integer data types, removing hardcoded absolute local paths from configurations, wrapping image loading operations in error-handling blocks, and utilizing official Hugging Face utilities rather than manually constructing cache paths.

gemini-code-assist · 2026-08-17T21:53:15Z

+      if image.max() > 1.0:
+        image = image / 127.5 - 1.0


Using image.max() > 1.0 to determine whether to normalize is unreliable. If an image is extremely dark (e.g., max pixel value is 1.0 or less), it will skip normalization entirely, leading to incorrect pixel values. It is safer to check if the image has an integer data type (like uint8) to decide if normalization is needed.

Suggested change

if image.max() > 1.0:

image = image / 127.5 - 1.0

if np.issubdtype(image.dtype, np.integer):

image = image.astype(np.float32) / 127.5 - 1.0

elif image.max() > 1.0:

image = image / 127.5 - 1.0

gemini-code-assist · 2026-08-17T21:53:15Z

+          if img.max() > 1.0:
+            img = img / 127.5 - 1.0


Using img.max() > 1.0 to determine whether to normalize is unreliable. If an image is extremely dark (e.g., max pixel value is 1.0 or less), it will skip normalization entirely, leading to incorrect pixel values. It is safer to check if the image has an integer data type (like uint8) to decide if normalization is needed.

Suggested change

if img.max() > 1.0:

img = img / 127.5 - 1.0

if np.issubdtype(img.dtype, np.integer):

img = img.astype(np.float32) / 127.5 - 1.0

elif img.max() > 1.0:

img = img / 127.5 - 1.0

gemini-code-assist · 2026-08-17T21:53:15Z

+image_paths: [
+  "/mnt/data/golden_image_edit_data/ref_images/ref_image_0.png",
+  "/mnt/data/golden_image_edit_data/ref_images/ref_image_1.png",
+  "/mnt/data/golden_image_edit_data/ref_images/ref_image_2.png",
+  "/mnt/data/golden_image_edit_data/ref_images/ref_image_3.png",
+]


Hardcoding absolute local paths like /mnt/data/... makes the configuration non-portable and will cause failures on other environments. Please use an empty list [] as the default and let users specify their paths via CLI or custom configs.

image_paths: []

gemini-code-assist · 2026-08-17T21:53:15Z

+image_paths: [
+  "/mnt/data/golden_image_edit_data/ref_images/ref_image_0.png",
+  "/mnt/data/golden_image_edit_data/ref_images/ref_image_1.png",
+  "/mnt/data/golden_image_edit_data/ref_images/ref_image_2.png",
+  "/mnt/data/golden_image_edit_data/ref_images/ref_image_3.png",
+]


Hardcoding absolute local paths like /mnt/data/... makes the configuration non-portable and will cause failures on other environments. Please use an empty list [] as the default and let users specify their paths via CLI or custom configs.

image_paths: []

gemini-code-assist · 2026-08-17T21:53:15Z

+      images = [
+          Image.open(p).convert("RGB").resize((config.width, config.height), Image.Resampling.BICUBIC)
+          for p in image_paths
+      ]


Calling Image.open(p) directly can raise UnidentifiedImageError or OSError if the file is corrupted or permissions are restricted, crashing the entire generation script. Wrap the image loading in a try-except block to handle failures gracefully.

Suggested change

images = [

Image.open(p).convert("RGB").resize((config.width, config.height), Image.Resampling.BICUBIC)

for p in image_paths

]

images = []

for p in image_paths:

try:

img = Image.open(p).convert("RGB").resize((config.width, config.height), Image.Resampling.BICUBIC)

images.append(img)

except Exception as e:

max_logging.log(f"⚠️ Error loading image {p}: {e}")

gemini-code-assist · 2026-08-17T21:53:15Z

+  for p in image_paths:
+    if os.path.exists(p):
+      images.append(Image.open(p))
+    else:
+      max_logging.log(f"⚠️ Warning: Reference image path not found: {p}")


Even if os.path.exists(p) is true, Image.open(p) can still raise UnidentifiedImageError or OSError (e.g., due to corruption or permission issues). Wrapping this in a try-except block ensures the script doesn't crash unexpectedly.

Suggested change

for p in image_paths:

if os.path.exists(p):

images.append(Image.open(p))

else:

max_logging.log(f"⚠️ Warning: Reference image path not found: {p}")

for p in image_paths:

try:

images.append(Image.open(p).convert("RGB"))

except Exception as e:

max_logging.log(f"⚠️ Error loading reference image {p}: {e}")

gemini-code-assist · 2026-08-17T21:53:15Z

+      hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
+      repo_cache = os.path.join(
+          hf_home,
+          "hub",
+          f"models--{getattr(config, 'pretrained_model_name_or_path', '').replace('/', '--')}",
+          "snapshots",
+      )
+      if os.path.exists(repo_cache) and os.listdir(repo_cache):
+        tokenizer_path = os.path.join(repo_cache, os.listdir(repo_cache)[0])


Manually constructing the Hugging Face cache path is fragile and bypasses user-configured cache directories (like HF_HUB_CACHE). Use the official snapshot_download with local_files_only=True to resolve the local snapshot directory robustly.

Suggested change

hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))

repo_cache = os.path.join(

hf_home,

"hub",

f"models--{getattr(config, 'pretrained_model_name_or_path', '').replace('/', '--')}",

"snapshots",

)

if os.path.exists(repo_cache) and os.listdir(repo_cache):

tokenizer_path = os.path.join(repo_cache, os.listdir(repo_cache)[0])

from huggingface_hub import snapshot_download

try:

tokenizer_path = snapshot_download(

repo_id=getattr(config, "pretrained_model_name_or_path", ""),

local_files_only=True

)

except Exception:

pass

… obsolete dev tests - Deleted separate YAML configs (base_flux2klein_image_edit.yml, base_flux2klein_image_edit_9B.yml). - Deleted separate CLI runner (generate_flux2klein_image_edit.py) and separate pipeline (flux2klein_image_edit_pipeline.py). - Removed obsolete intermediate parity/scratch scripts in favor of test_flux2klein_image_edit_e2e_parity.py and generate_flux2klein_smoke_test.py. - Cleaned up export in maxdiffusion.pipelines.flux.__init__.

- Renamed test_flux2klein_image_edit_e2e_parity.py to edit_flux2klein_e2e_test.py. - Added test_flux2klein_4b_image_edit_smoke and test_flux2klein_9b_image_edit_smoke in generate_flux2klein_smoke_test.py. - Added golden reference images ref_flux2klein_4b_image_edit.png and ref_flux2klein_9b_image_edit.png conditioned on ref_flux2klein_4b.png with prompt 'change the lighting to evening'. - Verified SSIM >= 0.95 assertion on both 4B and 9B image edit smoke tests on TPU VM.

amepas added 3 commits August 14, 2026 01:44

gemini-code-assist Bot reviewed Aug 17, 2026

View reviewed changes

amepas added 2 commits August 17, 2026 22:00

amepas closed this Aug 17, 2026

Navigation Menu

Uh oh!

Adding Image Editing - #460

Adding Image Editing#460
amepas wants to merge 5 commits into
mainfrom
onboarding-imageedit-flux2klein

amepas commented Aug 17, 2026

Uh oh!

github-actions Bot commented Aug 17, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Aug 17, 2026

Uh oh!

gemini-code-assist Bot Aug 17, 2026

Uh oh!

gemini-code-assist Bot Aug 17, 2026

Uh oh!

gemini-code-assist Bot Aug 17, 2026

Uh oh!

gemini-code-assist Bot Aug 17, 2026

Uh oh!

gemini-code-assist Bot Aug 17, 2026

Uh oh!

gemini-code-assist Bot Aug 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

-      images = [
-          Image.open(p).convert("RGB").resize((config.width, config.height), Image.Resampling.BICUBIC)
-          for p in image_paths
-      ]
+      images = []
+      for p in image_paths:
+        try:
+          img = Image.open(p).convert("RGB").resize((config.width, config.height), Image.Resampling.BICUBIC)
+          images.append(img)
+        except Exception as e:
+          max_logging.log(f"⚠️ Error loading image {p}: {e}")

Uh oh!

Conversation

amepas commented Aug 17, 2026

Uh oh!

github-actions Bot commented Aug 17, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant