merge: integrate Z-Image i2L (Image-to-LoRA) from upstream PR filipstrand#361

Adds mflux-z-image-i2l CLI command that generates LoRA adapters from
reference images using SigLIP2 + DINOv3 + i2L decoder, entirely on-device.
Also adds .default. LoRA naming support to ZImageLoRAMapping.

Author: Andrea

CHANGELOG.md

@@ -5,6 +5,15 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.16.4] - 2026-02-15
+
+### 🐛 Bug Fixes
+
+- **Training preview stability**: Always offload optimizer state during preview generation to avoid memory pressure and improve preview reliability.
+- **Apple Silicon compile guard**: Narrow the M1/M2 compile fallback so it excludes Max and Ultra variants, preserving expected optimized behavior on those chips.
+
+---
+
 ## [0.16.3] - 2026-02-14
 
 ### 🐛 Bug Fixes

pyproject.toml

@@ -14,7 +14,7 @@ source-exclude = [
 
 [project]
 name = "mflux"
-version = "0.16.3"
+version = "0.16.4"
 description = "MLX native implementations of state-of-the-art generative image models."
 readme = "README.md"
 keywords = ["flux", "ai", "ml", "transformers", "mlx", "huggingface", "apple-silicon", "diffusers", "qwen", "qwen-image", "seedvr2", "z-image"]
@@ -89,6 +89,7 @@ mflux-generate-qwen-edit = "mflux.models.qwen.cli.qwen_image_edit_generate:main"
 mflux-generate-fibo = "mflux.models.fibo.cli.fibo_generate:main"
 mflux-generate-z-image = "mflux.models.z_image.cli.z_image_generate:main"
 mflux-generate-z-image-turbo = "mflux.models.z_image.cli.z_image_turbo_generate:main"
+mflux-z-image-i2l = "mflux.models.z_image.cli.z_image_i2l:main"
 mflux-refine-fibo = "mflux.models.fibo_vlm.cli.fibo_refine:main"
 mflux-inspire-fibo = "mflux.models.fibo_vlm.cli.fibo_inspire:main"
 mflux-concept = "mflux.models.flux.cli.flux_concept:main"

src/mflux/models/common/training/trainer.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
+import gc
 import random
+import tempfile
 from pathlib import Path
 
 import mlx.core as mx
 from mlx import nn
+from mlx.utils import tree_unflatten
 from PIL import Image as PILImage
 from tqdm import tqdm
 
@@ -120,7 +123,7 @@ def train(
         )
 
         if training_spec.monitoring is not None and training_state.iterator.num_iterations == 0:
-            TrainingTrainer._generate_previews(adapter, training_spec, training_state)
+            TrainingTrainer._generate_previews_with_optimizer_offload(adapter, training_spec, training_state)
             validation_batch = training_state.iterator.get_validation_batch()
             validation_loss = TrainingTrainer.compute_loss(adapter, training_spec, base_config, validation_batch)
             training_state.statistics.append_values(step=training_state.iterator.num_iterations, loss=float(validation_loss))  # fmt: off
@@ -147,7 +150,7 @@ def train(
                 del validation_loss
 
             if training_state.should_generate_image(training_spec):
-                TrainingTrainer._generate_previews(adapter, training_spec, training_state)
+                TrainingTrainer._generate_previews_with_optimizer_offload(adapter, training_spec, training_state)
 
             if training_state.should_save(training_spec):
                 training_state.save(adapter, training_spec)
@@ -227,3 +230,25 @@ def _generate_previews(
                 )
             )
             del image
+
+    @staticmethod
+    def _generate_previews_with_optimizer_offload(
+        adapter: TrainingAdapter,
+        training_spec: TrainingSpec,
+        training_state: TrainingState,
+    ) -> None:
+        optimizer = training_state.optimizer
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            offload_path = Path(tmp_dir) / "optimizer_offload.safetensors"
+            optimizer.save(offload_path)
+            optimizer.optimizer.state = []
+
+            gc.collect()
+            mx.clear_cache()
+            try:
+                TrainingTrainer._generate_previews(adapter, training_spec, training_state)
+            finally:
+                restored_state = tree_unflatten(list(mx.load(str(offload_path)).items()))
+                optimizer.optimizer.state = restored_state
+                gc.collect()
+                mx.clear_cache()

src/mflux/models/z_image/cli/z_image_i2l.py

@@ -0,0 +1,92 @@
+"""CLI entrypoint for Z-Image Image-to-LoRA (i2L).
+
+Usage:
+    mflux-z-image-i2l --image-path ./style_images --output style_lora.safetensors
+    mflux-z-image-i2l --image-path img1.jpg img2.jpg --output style_lora.safetensors
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+from PIL import Image
+
+IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tiff", ".tif"}
+
+
+def _collect_images(paths: list[str]) -> list[Path]:
+    """Resolve a mix of files and directories into a sorted list of image paths."""
+    result = []
+    for p_str in paths:
+        p = Path(p_str)
+        if not p.exists():
+            print(f"Error: Path not found: {p_str}", file=sys.stderr)
+            sys.exit(1)
+        if p.is_dir():
+            found = sorted(f for f in p.iterdir() if f.suffix.lower() in IMAGE_EXTENSIONS)
+            if not found:
+                print(f"Error: No images found in directory: {p_str}", file=sys.stderr)
+                sys.exit(1)
+            result.extend(found)
+        elif p.suffix.lower() in IMAGE_EXTENSIONS:
+            result.append(p)
+        else:
+            print(f"Error: Unsupported file type: {p_str}", file=sys.stderr)
+            sys.exit(1)
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate LoRA weights from style reference images using Z-Image i2L.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  mflux-z-image-i2l --image-path ./my_style
+  mflux-z-image-i2l --image-path ./my_style --output my_style.safetensors
+  mflux-z-image-i2l --image-path img1.jpg img2.jpg img3.jpg img4.jpg
+  mflux-z-image-i2l --image-path ./style_a ./style_b/photo.png
+
+The generated LoRA can then be used with mflux-generate-z-image-turbo:
+  mflux-generate-z-image-turbo --prompt "a cat" --lora-paths style.safetensors
+        """,
+    )
+    parser.add_argument(
+        "--image-path",
+        "-i",
+        nargs="+",
+        required=True,
+        type=str,
+        help="Image files or directories containing style reference images.",
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=str,
+        default="lora.safetensors",
+        help="Output path for the generated LoRA file. Default: lora.safetensors",
+    )
+
+    args = parser.parse_args()
+
+    # Collect image paths from files and directories
+    image_paths = _collect_images(args.image_path)
+
+    # Load images
+    print(f"Loading {len(image_paths)} image(s)...")
+    images = []
+    for p in image_paths:
+        img = Image.open(p).convert("RGB")
+        images.append(img)
+        print(f"  {p.name}: {img.size[0]}x{img.size[1]}")
+
+    # Import here to avoid slow startup for --help
+    from mflux.models.z_image.model.z_image_i2l.i2l_pipeline import ZImageI2LPipeline
+
+    # Create pipeline and generate LoRA
+    pipeline = ZImageI2LPipeline.from_pretrained()
+    pipeline.generate_lora(images=images, output_path=args.output)
+
+
+if __name__ == "__main__":
+    main()

src/mflux/models/z_image/model/z_image_i2l/__init__.py

@@ -0,0 +1,45 @@
+import mlx.core as mx
+from mlx import nn
+from mlx.core.fast import scaled_dot_product_attention
+
+from mflux.models.z_image.model.z_image_i2l.dinov3.dinov3_rope import apply_dinov3_rope
+
+
+class DINOv3Attention(nn.Module):
+    """DINOv3 attention with RoPE on patch tokens.
+
+    hidden_size=4096, num_heads=32, head_dim=128.
+    Bias config: q=False, k=False, v=False, o=True.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.num_heads = 32
+        self.head_dim = 128  # 4096 / 32
+        dim = 4096
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.o_proj = nn.Linear(dim, dim, bias=True)
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        cos: mx.array,
+        sin: mx.array,
+        num_prefix_tokens: int = 5,
+    ) -> mx.array:
+        B, N, _ = hidden_states.shape
+
+        q = self.q_proj(hidden_states).reshape(B, N, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
+        k = self.k_proj(hidden_states).reshape(B, N, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
+        v = self.v_proj(hidden_states).reshape(B, N, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
+
+        # Apply RoPE only to patch tokens
+        q, k = apply_dinov3_rope(q, k, cos, sin, num_prefix_tokens=num_prefix_tokens)
+
+        scale = 1.0 / mx.sqrt(mx.array(self.head_dim, dtype=q.dtype))
+        out = scaled_dot_product_attention(q, k, v, scale=scale)
+
+        out = out.transpose(0, 2, 1, 3).reshape(B, N, -1)
+        return self.o_proj(out)

src/mflux/models/z_image/model/z_image_i2l/dinov3/__init__.py

@@ -0,0 +1,50 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class DINOv3Embeddings(nn.Module):
+    """DINOv3 embeddings: CLS token + register tokens + patch embeddings.
+
+    image_size=224, patch_size=16, hidden_size=4096, num_register_tokens=4.
+    Sequence: [CLS, reg0, reg1, reg2, reg3, patch0, patch1, ...]
+    Total prefix tokens = 5 (1 CLS + 4 registers).
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.hidden_size = 4096
+        self.patch_size = 16
+        self.image_size = 224
+        self.num_register_tokens = 4
+
+        self.cls_token = mx.random.normal(shape=(1, 1, self.hidden_size))
+        self.register_tokens = mx.random.normal(shape=(1, self.num_register_tokens, self.hidden_size))
+        self.patch_embeddings = nn.Conv2d(
+            in_channels=3,
+            out_channels=self.hidden_size,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            bias=True,
+        )
+
+    def __call__(self, pixel_values: mx.array) -> mx.array:
+        """
+        Args:
+            pixel_values: (B, 3, 224, 224)
+        Returns:
+            (B, 1 + 4 + 196, 4096) = (B, 201, 4096)
+        """
+        B = pixel_values.shape[0]
+
+        # Patch embedding: (B, C, H, W) -> (B, H', W', hidden) via Conv2d (needs NHWC)
+        x = mx.transpose(pixel_values, (0, 2, 3, 1))  # (B, H, W, C)
+        x = self.patch_embeddings(x)  # (B, H', W', hidden)
+        # Flatten spatial: (B, num_patches, hidden)
+        x = x.reshape(B, -1, self.hidden_size)
+
+        # Prepend CLS and register tokens
+        cls_tokens = mx.broadcast_to(self.cls_token, (B, 1, self.hidden_size))
+        reg_tokens = mx.broadcast_to(self.register_tokens, (B, self.num_register_tokens, self.hidden_size))
+        x = mx.concatenate([cls_tokens, reg_tokens, x], axis=1)
+
+        return x

src/mflux/models/z_image/model/z_image_i2l/dinov3/dinov3_attention.py

@@ -0,0 +1,13 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class DINOv3LayerScale(nn.Module):
+    """Learnable per-channel scaling, same pattern as DINOv2."""
+
+    def __init__(self, dims: int = 4096, init_values: float = 1.0):
+        super().__init__()
+        self.gamma = init_values * mx.ones((dims,))
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return x * self.gamma

src/mflux/models/z_image/model/z_image_i2l/dinov3/dinov3_embeddings.py

@@ -0,0 +1,19 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class DINOv3GatedMLP(nn.Module):
+    """DINOv3 gated MLP with SiLU activation.
+
+    hidden_size=4096, intermediate_size=8192. All projections have bias.
+    Formula: down_proj(silu(gate_proj(x)) * up_proj(x))
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.gate_proj = nn.Linear(4096, 8192, bias=True)
+        self.up_proj = nn.Linear(4096, 8192, bias=True)
+        self.down_proj = nn.Linear(8192, 4096, bias=True)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))