diff --git a/src/maxdiffusion/configs/base_flux2klein.yml b/src/maxdiffusion/configs/base_flux2klein.yml index f1a7e0fa4..16a8a5cfa 100644 --- a/src/maxdiffusion/configs/base_flux2klein.yml +++ b/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: '' diff --git a/src/maxdiffusion/configs/base_flux2klein_9B.yml b/src/maxdiffusion/configs/base_flux2klein_9B.yml index a3a0afeac..c6dc2689b 100644 --- a/src/maxdiffusion/configs/base_flux2klein_9B.yml +++ b/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: '' diff --git a/src/maxdiffusion/generate_flux2klein.py b/src/maxdiffusion/generate_flux2klein.py index b1427f937..887bcfdca 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 @@ -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,6 +579,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 +587,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 +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) 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..3183f462c 100644 --- a/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py +++ b/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) ] diff --git a/src/maxdiffusion/models/flux/util.py b/src/maxdiffusion/models/flux/util.py index aa43609a1..5eb25572c 100644 --- a/src/maxdiffusion/models/flux/util.py +++ b/src/maxdiffusion/models/flux/util.py @@ -734,6 +734,74 @@ 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 prepare_image_latents(vae, images, bn_mean, bn_std, scale=10): + """Encodes, patchifies, normalizes, packs, and generates 4D RoPE IDs for a list of reference images. + + images: list of arrays of shape (1, 3, H_k, W_k) or (3, H_k, W_k) in range [-1, 1]. + Returns: + image_latents_concat: shape (1, total_ref_tokens, 128) + image_latent_ids: shape (1, total_ref_tokens, 4) + """ + import jax.numpy as jnp + from einops import rearrange + + norm_latents = [] + for img in images: + if img.ndim == 3: + img = jnp.expand_dims(img, axis=0) + raw_latents = vae.encode(img) # (1, 32, H/8, W/8) + patchified = patchify_latents(raw_latents) # (1, 128, H/16, W/16) + normalized = (patchified - bn_mean) / bn_std + norm_latents.append(normalized) + + image_latent_ids = prepare_multi_image_ids(norm_latents, scale=scale) + + packed_latents = [] + for latent in norm_latents: + packed = rearrange(latent, "b c h w -> b (h w) c") + packed_latents.append(packed) + + image_latents_concat = jnp.concatenate(packed_latents, axis=1) + return image_latents_concat, image_latent_ids + + 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 +824,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...") + # 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")) - # 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")) - - # 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 +878,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 "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")) + + 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")) - if b_idx < 3: - upsampler_jax = up_block_jax["upsamplers_0"] - upsampler_pt = f"{up_block_pt}.upsamplers.0" + for idx in [0, 1]: + res_jax = dec_jax["mid_block"][f"resnets_{idx}"] + res_pt_prefix = f"decoder.mid_block.resnets.{idx}" - upsampler_jax["conv"]["kernel"] = jnp.array( - get_pytorch_weight_tensor(f"{upsampler_pt}.conv.weight").transpose(2, 3, 1, 0) + 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..e69de29bb 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..c94ebd8a3 100644 --- a/src/maxdiffusion/pipelines/flux/__init__.py +++ b/src/maxdiffusion/pipelines/flux/__init__.py @@ -19,3 +19,6 @@ from .flux_pipeline import ( FluxPipeline, ) +from .flux2klein_pipeline import ( + FlaxFlux2KleinPipeline, +) diff --git a/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py b/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py index 4fa257ca2..8ff868a0b 100644 --- a/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py +++ b/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py @@ -41,7 +41,9 @@ from ...models.flux.util import ( pack_latents, + patchify_latents, prepare_latent_image_ids, + prepare_multi_image_ids, prepare_text_ids, ) @@ -105,6 +107,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 +126,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 - 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)) + vae_bn_mean_seq = vae_bn_mean.reshape(1, 1, 128) + vae_bn_std_seq = vae_bn_std.reshape(1, 1, 128) - res = self.vae.apply({"params": v_params}, latents=latents_unpacked, method=self.vae.decode) - return FlaxDecoderOutput(sample=res.sample) + 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) if isinstance(self.transformer, nnx.Module): g, nnx_state, r = nnx.split(self.transformer, nnx.Param, ...) @@ -157,8 +195,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,7 +217,15 @@ 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]) @@ -199,8 +247,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,7 +268,15 @@ 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]) @@ -228,6 +286,7 @@ def scan_body(cur_latents, step_idx): self._jitted_qwen3_forward = qwen3_forward self._jitted_transformer_step = transformer_step self._jitted_fused_denoise_loop = fused_denoise_loop + self._jitted_vae_encode = vae_encode self._jitted_vae_decode = vae_decode def _get_dynamic_batch_sharding(self): @@ -237,25 +296,44 @@ 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, + image=None, + num_conditioning_images=0, ): """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 None and image is not None: + images = image if isinstance(image, list) else [image] + + if images is not None and len(images) > 0: + num_conditioning_images = len(images) + seq_len_img = (height // 16) * (width // 16) + total_img_len = (1 + num_conditioning_images) * seq_len_img 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 +355,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) @@ -303,22 +382,32 @@ def compile_transformer(): 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() + dummy_rgb = jnp.zeros((1, 3, height, width), 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,6 +453,8 @@ def __call__( width: int = 1024, num_inference_steps: int = 4, batch_size: int = 1, + images: Optional[List[Any]] = None, + image: Optional[Union[Any, List[Any]]] = None, use_latents: bool = False, latents: Optional[Any] = None, measure_time: bool = False, @@ -375,6 +466,9 @@ def __call__( # 1. Setup JIT functions self._setup_jit_functions() + if images is None and image is not None: + images = image if isinstance(image, list) else [image] + # 2. Setup prompts and inputs if isinstance(prompt, str): prompts = [prompt] * batch_size @@ -392,6 +486,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 +497,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 +512,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 +527,89 @@ 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: + if isinstance(img, Image.Image): + img = img.convert("RGB").resize((width, height), Image.Resampling.BICUBIC) + arr = np.array(img, dtype=np.float32) / 127.5 - 1.0 + arr = np.transpose(arr, (2, 0, 1)) + img_tensor = jnp.expand_dims(jnp.array(arr), axis=0) + elif isinstance(img, np.ndarray): + if img.ndim == 3: + img = np.expand_dims(img, axis=0) + if img.shape[-1] == 3: + img = np.transpose(img, (0, 3, 1, 2)) + if np.issubdtype(img.dtype, np.integer): + img = img.astype(np.float32) / 127.5 - 1.0 + elif np.issubdtype(img.dtype, np.floating): + if img.max() > 1.0: + img = img / 127.5 - 1.0 + elif img.min() >= 0.0: + img = img * 2.0 - 1.0 + img_tensor = jnp.array(img, dtype=np.float32) + elif isinstance(img, jnp.ndarray): + if img.ndim == 3: + img = jnp.expand_dims(img, axis=0) + if img.shape[-1] == 3: + img = jnp.transpose(img, (0, 3, 1, 2)) + if jnp.issubdtype(img.dtype, jnp.integer): + img = img.astype(jnp.float32) / 127.5 - 1.0 + elif jnp.issubdtype(img.dtype, jnp.floating): + if img.max() > 1.0: + img = img / 127.5 - 1.0 + elif img.min() >= 0.0: + img = img * 2.0 - 1.0 + img_tensor = img + else: + raise ValueError(f"Unsupported image type: {type(img)}") + + 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) + 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 + 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) @@ -543,6 +718,7 @@ def put_data_on_devices(x, sharding): timesteps_device, sigmas_device, guidance_vec_val, + seq_len_img, ) latents_jax.block_until_ready() if do_prof_denoise: @@ -569,6 +745,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/tests/edit_flux2klein_e2e_test.py b/src/maxdiffusion/tests/edit_flux2klein_e2e_test.py new file mode 100644 index 000000000..18a4ff074 --- /dev/null +++ b/src/maxdiffusion/tests/edit_flux2klein_e2e_test.py @@ -0,0 +1,349 @@ +""" +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 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" + + +class TestFlux2KleinImageEditE2EParity(unittest.TestCase): + """End-to-End Parity Test between PyTorch Diffusers CPU and MaxDiffusion TPU.""" + + def setUp(self): + jax.config.update("jax_default_matmul_precision", "highest") + 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: + 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-4B/snapshots"), + os.path.join(hf_home, "hub/models--black-forest-labs--FLUX.2-klein-4b/snapshots"), + "/mnt/hyperdisk_weights/hub/models--black-forest-labs--FLUX.2-klein-4B/snapshots", + "/mnt/data/models/flux2klein-4b", + ] + self.model_dir = None + for c in candidates: + if os.path.exists(c): + if "snapshots" in c: + snaps = os.listdir(c) + if snaps: + self.model_dir = os.path.join(c, snaps[0]) + else: + self.model_dir = c + if self.model_dir: + 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") + if os.path.exists(self.transformer_path) and os.path.exists(self.vae_path): + break + if self.model_dir is None: + self.model_dir = "black-forest-labs/FLUX.2-klein-4B" + + if hasattr(self, "model_dir") and self.model_dir and not hasattr(self, "transformer_path"): + 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 + + print("\n" + "=" * 80) + print("🚀 [STEP 1/3] Running Reference PyTorch Diffusers CPU Pipeline...") + print("=" * 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) + print(f" -> Saved PyTorch Diffusers output to: {diffusers_img_path}") + + # Free PyTorch pipeline memory before TPU run + del diffusers_pipe + gc.collect() + + print("\n" + "=" * 80) + print("🚀 [STEP 2/3] Running MaxDiffusion Unified FlaxFlux2KleinPipeline 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.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 + print(" -> 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 + print(" -> 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 + print(" -> 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=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, + 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 + print(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") + + print("\n" + "=" * 80) + print("📊 [STEP 3/3] Evaluating End-to-End Parity (SSIM & PSNR)...") + print("=" * 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)) + + print(f" -> SSIM (Diffusers CPU vs MaxDiffusion TPU): {ssim_val:.6f}") + print(f" -> PSNR (Diffusers CPU vs MaxDiffusion TPU): {psnr_val:.2f} dB") + print(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) + print(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!") + print("🎉 END-TO-END PARITY TEST PASSED! MaxDiffusion matches Diffusers reference!") + + +if __name__ == "__main__": + unittest.main() diff --git a/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py b/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py index c61db29f7..b0c4f3f2f 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,104 @@ 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) + 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_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