ollama source for Momentry Core verification

x/imagegen/.gitignore

@@ -0,0 +1,38 @@
+# Build directories
+build/
+dist/
+
+# CMake
+CMakeCache.txt
+CMakeFiles/
+cmake_install.cmake
+Makefile
+*.cmake
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+*~
+
+# macOS
+.DS_Store
+*.dSYM/
+
+# Go
+*.exe
+*.exe~
+*.dll
+*.so
+*.dylib
+
+# Python
+*.npy
+
+/engine
+weights
+outputs
+
+prompt.txt
+negative.txt

x/imagegen/README.md

@@ -0,0 +1,250 @@
+# Image Generation in Ollama (Experimental)
+
+Generate images from text prompts using local AI models.
+
+## Quick Start
+
+```bash
+# Run with a prompt
+ollama run z-image "a sunset over mountains"
+Generating: step 30/30
+Image saved to: /tmp/ollama-image-1704067200.png
+```
+
+On macOS, the generated image will automatically open in Preview.
+
+## Supported Models
+
+| Model | VRAM Required | Notes |
+|-------|---------------|-------|
+| z-image | ~12GB | Based on Flux architecture |
+
+## CLI Usage
+
+```bash
+# Generate an image
+ollama run z-image "a cat playing piano"
+
+# Check if model is running
+ollama ps
+
+# Stop the model
+ollama stop z-image
+```
+
+## API
+
+### OpenAI-Compatible Endpoint
+
+```bash
+POST /v1/images/generations
+```
+
+**Request:**
+```json
+{
+  "model": "z-image",
+  "prompt": "a sunset over mountains",
+  "size": "1024x1024",
+  "response_format": "b64_json"
+}
+```
+
+**Response:**
+```json
+{
+  "created": 1704067200,
+  "data": [
+    {
+      "b64_json": "iVBORw0KGgo..."
+    }
+  ]
+}
+```
+
+### Example: cURL
+
+```bash
+curl http://localhost:11434/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "z-image",
+    "prompt": "a white cat",
+    "size": "1024x1024"
+  }'
+```
+
+### Example: Save to File
+
+```bash
+curl -s http://localhost:11434/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "z-image",
+    "prompt": "a white cat",
+    "size": "1024x1024"
+  }' | jq -r '.data[0].b64_json' | base64 -d > image.png
+```
+
+### Streaming Progress
+
+Enable streaming to receive progress updates via SSE:
+
+```bash
+curl http://localhost:11434/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{"model": "z-image", "prompt": "a sunset", "stream": true}'
+```
+
+Events:
+```
+event: progress
+data: {"step": 1, "total": 30}
+
+event: progress
+data: {"step": 2, "total": 30}
+...
+
+event: done
+data: {"created": 1704067200, "data": [{"b64_json": "..."}]}
+```
+
+## Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| model | string | required | Model name |
+| prompt | string | required | Text description of image |
+| size | string | "1024x1024" | Image dimensions (WxH) |
+| n | int | 1 | Number of images (currently only 1 supported) |
+| response_format | string | "b64_json" | "b64_json" or "url" |
+| stream | bool | false | Enable progress streaming |
+
+## Requirements
+
+- macOS with Apple Silicon (M1/M2/M3/M4)
+- CUDA: tested on CUDA 12 Blackwell, more testing coming soon
+- Sufficient VRAM (see model table above)
+- Ollama built with MLX support
+
+## Limitations
+
+- macOS only (uses MLX backend)
+- Single image per request
+- Fixed step count (30 steps)
+- Modelfiles not yet supported (use `ollama create` from model directory)
+
+---
+
+# Tensor Model Storage Format
+
+Tensor models store each tensor as a separate blob with metadata in the manifest. This enables faster downloads (parallel fetching) and deduplication (shared tensors are stored once).
+
+## Manifest Structure
+
+The manifest follows the standard ollama format with tensor-specific layer metadata:
+
+```json
+{
+  "schemaVersion": 2,
+  "mediaType": "application/vnd.docker.distribution.manifest.v2+json",
+  "config": { "digest": "sha256:...", "size": 1234 },
+  "layers": [
+    {
+      "mediaType": "application/vnd.ollama.image.tensor",
+      "digest": "sha256:25b36eed...",
+      "size": 49807448,
+      "name": "text_encoder/model.layers.0.mlp.down_proj.weight",
+      "dtype": "BF16",
+      "shape": [2560, 9728]
+    },
+    {
+      "mediaType": "application/vnd.ollama.image.json",
+      "digest": "sha256:abc123...",
+      "size": 512,
+      "name": "text_encoder/config.json"
+    }
+  ]
+}
+```
+
+Each tensor layer includes:
+- `name`: Path-style tensor name (e.g., `text_encoder/model.layers.0.mlp.down_proj.weight`)
+- `dtype`: Data type (BF16, F32, etc.)
+- `shape`: Tensor dimensions
+
+Config layers use the same path-style naming (e.g., `tokenizer/tokenizer.json`).
+
+## Blob Format
+
+Each tensor blob is a minimal safetensors file:
+
+```
+[8 bytes: header size (uint64 LE)]
+[~80 bytes: JSON header, padded to 8-byte alignment]
+[N bytes: raw tensor data]
+```
+
+Header contains a single tensor named `"data"`:
+
+```json
+{"data":{"dtype":"BF16","shape":[2560,9728],"data_offsets":[0,49807360]}}
+```
+
+## Why Include the Header?
+
+The ~88 byte safetensors header enables MLX's native `mlx_load_safetensors` function, which:
+
+1. **Uses mmap** - Maps file directly into memory, no copies
+2. **Zero-copy to GPU** - MLX reads directly from mapped pages
+3. **No custom code** - Standard MLX API, battle-tested
+
+Without the header, we'd need custom C++ code to create MLX arrays from raw mmap'd data. MLX's public API doesn't expose this - it always copies when creating arrays from external pointers.
+
+The overhead is negligible: 88 bytes per tensor = ~100KB total for a 13GB model (0.0007%).
+
+## Why Per-Tensor Blobs?
+
+**Deduplication**: Blobs are content-addressed by SHA256. If two models share identical tensors (same weights, dtype, shape), they share the same blob file.
+
+Example: Model A and Model B both use the same text encoder. The text encoder's 400 tensors are stored once, referenced by both manifests.
+
+```
+~/.ollama/models/
+  blobs/
+    sha256-25b36eed...  <- shared by both models
+    sha256-abc123...
+  manifests/
+    library/model-a/latest  <- references sha256-25b36eed
+    library/model-b/latest  <- references sha256-25b36eed
+```
+
+## Import Flow
+
+```
+cd ./weights/Z-Image-Turbo
+ollama create z-image
+
+1. Scan component directories (text_encoder/, transformer/, vae/)
+2. For each .safetensors file:
+   - Extract individual tensors
+   - Wrap each in minimal safetensors format (88B header + data)
+   - Write to blob store (SHA256 content-addressed)
+   - Add layer entry to manifest with path-style name
+3. Copy config files (*.json) as config layers
+4. Write manifest
+```
+
+## FP8 Quantization
+
+Z-Image supports FP8 quantization to reduce memory usage by ~50% while maintaining image quality.
+
+### Usage
+
+```bash
+cd ./weights/Z-Image-Turbo
+ollama create z-image-fp8 --quantize fp8
+```
+
+This quantizes weights during import. The resulting model will be ~15GB instead of ~31GB.
+

x/imagegen/cache/cache.go

@@ -0,0 +1,170 @@
+package cache
+
+import "github.com/ollama/ollama/x/imagegen/mlx"
+
+type Cache interface {
+	Update(k, v *mlx.Array, seqLen int) (*mlx.Array, *mlx.Array)
+	Offset() int
+	Len() int
+	State() []*mlx.Array
+	Reset()
+}
+
+type KVCache struct {
+	keys, values *mlx.Array
+	offset       int
+	step         int
+}
+
+func NewKVCache() *KVCache {
+	return &KVCache{step: 256}
+}
+
+func (c *KVCache) Update(k, v *mlx.Array, seqLen int) (*mlx.Array, *mlx.Array) {
+	prev := c.offset
+	shape := k.Shape()
+	B, H, Dk := shape[0], shape[1], shape[3]
+	Dv := v.Shape()[3]
+
+	// Grow buffer if needed
+	if c.keys == nil || (prev+seqLen) > int(c.keys.Shape()[2]) {
+		nSteps := (c.step + seqLen - 1) / c.step
+		newK := mlx.Zeros([]int32{B, H, int32(nSteps * c.step), Dk}, k.Dtype())
+		newV := mlx.Zeros([]int32{B, H, int32(nSteps * c.step), Dv}, v.Dtype())
+
+		if c.keys != nil {
+			if prev%c.step != 0 {
+				c.keys = mlx.Slice(c.keys, []int32{0, 0, 0, 0}, []int32{B, H, int32(prev), Dk})
+				c.values = mlx.Slice(c.values, []int32{0, 0, 0, 0}, []int32{B, H, int32(prev), Dv})
+			}
+			c.keys = mlx.Concatenate([]*mlx.Array{c.keys, newK}, 2)
+			c.values = mlx.Concatenate([]*mlx.Array{c.values, newV}, 2)
+		} else {
+			c.keys, c.values = newK, newV
+		}
+	}
+
+	c.offset += seqLen
+	c.keys = mlx.SliceUpdateInplace(c.keys, k, []int32{0, 0, int32(prev), 0}, []int32{B, H, int32(c.offset), Dk})
+	c.values = mlx.SliceUpdateInplace(c.values, v, []int32{0, 0, int32(prev), 0}, []int32{B, H, int32(c.offset), Dv})
+
+	return mlx.Slice(c.keys, []int32{0, 0, 0, 0}, []int32{B, H, int32(c.offset), Dk}),
+		mlx.Slice(c.values, []int32{0, 0, 0, 0}, []int32{B, H, int32(c.offset), Dv})
+}
+
+func (c *KVCache) State() []*mlx.Array {
+	if c.keys == nil {
+		return nil
+	}
+	return []*mlx.Array{c.keys, c.values}
+}
+
+func (c *KVCache) Offset() int { return c.offset }
+func (c *KVCache) Len() int    { return c.offset }
+
+// Reset clears the cache state for a new generation session
+func (c *KVCache) Reset() {
+	c.keys = nil
+	c.values = nil
+	c.offset = 0
+}
+
+// RotatingKVCache implements sliding window attention with bounded memory
+type RotatingKVCache struct {
+	keys, values *mlx.Array
+	offset       int
+	maxSize      int
+	step         int
+	idx          int
+}
+
+func NewRotatingKVCache(maxSize int) *RotatingKVCache {
+	return &RotatingKVCache{maxSize: maxSize, step: 256}
+}
+
+func (c *RotatingKVCache) Update(k, v *mlx.Array, seqLen int) (*mlx.Array, *mlx.Array) {
+	if seqLen > 1 {
+		return c.updateConcat(k, v, seqLen)
+	}
+	return c.updateInPlace(k, v)
+}
+
+func (c *RotatingKVCache) updateInPlace(k, v *mlx.Array) (*mlx.Array, *mlx.Array) {
+	shape := k.Shape()
+	B, H, Dk := shape[0], shape[1], shape[3]
+	Dv := v.Shape()[3]
+
+	// Grow buffer if not yet at max
+	if c.keys == nil || (c.idx >= int(c.keys.Shape()[2]) && int(c.keys.Shape()[2]) < c.maxSize) {
+		var cap int
+		if c.keys != nil {
+			cap = int(c.keys.Shape()[2])
+		}
+		newSize := min(c.step, c.maxSize-cap)
+		newK := mlx.Zeros([]int32{B, H, int32(newSize), Dk}, k.Dtype())
+		newV := mlx.Zeros([]int32{B, H, int32(newSize), Dv}, v.Dtype())
+		if c.keys != nil {
+			c.keys = mlx.Concatenate([]*mlx.Array{c.keys, newK}, 2)
+			c.values = mlx.Concatenate([]*mlx.Array{c.values, newV}, 2)
+		} else {
+			c.keys, c.values = newK, newV
+		}
+	}
+
+	// Rotate when hitting max
+	if c.idx >= c.maxSize {
+		c.idx = 0
+	}
+
+	c.keys = mlx.SliceUpdateInplace(c.keys, k, []int32{0, 0, int32(c.idx), 0}, []int32{B, H, int32(c.idx + 1), Dk})
+	c.values = mlx.SliceUpdateInplace(c.values, v, []int32{0, 0, int32(c.idx), 0}, []int32{B, H, int32(c.idx + 1), Dv})
+
+	c.offset++
+	c.idx++
+
+	validLen := int32(min(c.offset, c.maxSize))
+	return mlx.Slice(c.keys, []int32{0, 0, 0, 0}, []int32{B, H, validLen, Dk}),
+		mlx.Slice(c.values, []int32{0, 0, 0, 0}, []int32{B, H, validLen, Dv})
+}
+
+func (c *RotatingKVCache) updateConcat(k, v *mlx.Array, seqLen int) (*mlx.Array, *mlx.Array) {
+	shape := k.Shape()
+	B, H, Dk := shape[0], shape[1], shape[3]
+	Dv := v.Shape()[3]
+
+	if c.keys == nil {
+		c.keys, c.values = k, v
+	} else {
+		c.keys = mlx.Concatenate([]*mlx.Array{c.keys, k}, 2)
+		c.values = mlx.Concatenate([]*mlx.Array{c.values, v}, 2)
+	}
+	c.offset += seqLen
+
+	// Trim to max_size to maintain sliding window
+	cap := int(c.keys.Shape()[2])
+	if trim := cap - c.maxSize; trim > 0 {
+		c.keys = mlx.Slice(c.keys, []int32{0, 0, int32(trim), 0}, []int32{B, H, int32(cap), Dk})
+		c.values = mlx.Slice(c.values, []int32{0, 0, int32(trim), 0}, []int32{B, H, int32(cap), Dv})
+	}
+
+	c.idx = int(c.keys.Shape()[2])
+	return c.keys, c.values
+}
+
+func (c *RotatingKVCache) State() []*mlx.Array {
+	if c.keys == nil {
+		return nil
+	}
+	return []*mlx.Array{c.keys, c.values}
+}
+
+func (c *RotatingKVCache) Offset() int { return c.offset }
+func (c *RotatingKVCache) Len() int    { return min(c.offset, c.maxSize) }
+
+// Reset clears the cache state for a new generation session
+func (c *RotatingKVCache) Reset() {
+	c.keys = nil
+	c.values = nil
+	c.offset = 0
+	c.idx = 0
+}

x/imagegen/cache/step.go

@@ -0,0 +1,162 @@
+package cache
+
+import "github.com/ollama/ollama/x/imagegen/mlx"
+
+// StepCache caches layer outputs across diffusion denoising steps.
+// Based on DeepCache (CVPR 2024) and Learning-to-Cache (NeurIPS 2024):
+// shallow layers change little between consecutive steps, so we can
+// cache their outputs and skip recomputation on non-refresh steps.
+//
+// Supports both single-stream and dual-stream architectures:
+//   - Single-stream: use Get/Set for the single output per layer
+//   - Dual-stream: use Get/Set for stream 1 (imgH), Get2/Set2 for stream 2 (txtH)
+//
+// Usage (single-stream):
+//
+//	cache := NewStepCache(15)  // cache first 15 layers
+//	for step := 0; step < numSteps; step++ {
+//	    refresh := cache.ShouldRefresh(step, 3)  // refresh every 3 steps
+//	    for i, layer := range layers {
+//	        if i < 15 && !refresh && cache.Get(i) != nil {
+//	            output = cache.Get(i)  // reuse cached
+//	        } else {
+//	            output = layer.Forward(input)
+//	            if i < 15 && refresh {
+//	                cache.Set(i, output)
+//	            }
+//	        }
+//	    }
+//	}
+//	cache.Free()  // cleanup when done
+//
+// Usage (dual-stream):
+//
+//	cache := NewStepCache(15)
+//	for step := 0; step < numSteps; step++ {
+//	    refresh := cache.ShouldRefresh(step, 3)
+//	    for i, layer := range layers {
+//	        if i < 15 && !refresh && cache.Get(i) != nil {
+//	            imgH, txtH = cache.Get(i), cache.Get2(i)
+//	        } else {
+//	            imgH, txtH = layer.Forward(imgH, txtH, ...)
+//	            if i < 15 && refresh {
+//	                cache.Set(i, imgH)
+//	                cache.Set2(i, txtH)
+//	            }
+//	        }
+//	    }
+//	}
+type StepCache struct {
+	layers   []*mlx.Array // cached layer outputs (stream 1)
+	layers2  []*mlx.Array // cached layer outputs (stream 2, for dual-stream models)
+	constant *mlx.Array   // optional constant (e.g., text embeddings)
+}
+
+// NewStepCache creates a cache for the given number of layers.
+func NewStepCache(numLayers int) *StepCache {
+	return &StepCache{
+		layers:  make([]*mlx.Array, numLayers),
+		layers2: make([]*mlx.Array, numLayers),
+	}
+}
+
+// ShouldRefresh returns true if the cache should be refreshed at this step.
+// Refresh happens on step 0, interval, 2*interval, etc.
+func (c *StepCache) ShouldRefresh(step, interval int) bool {
+	return step%interval == 0
+}
+
+// Get returns the cached output for a layer, or nil if not cached.
+func (c *StepCache) Get(layer int) *mlx.Array {
+	if layer < len(c.layers) {
+		return c.layers[layer]
+	}
+	return nil
+}
+
+// Set stores a layer output (stream 1), freeing any previous value.
+func (c *StepCache) Set(layer int, arr *mlx.Array) {
+	if layer < len(c.layers) {
+		if c.layers[layer] != nil {
+			c.layers[layer].Free()
+		}
+		c.layers[layer] = arr
+	}
+}
+
+// Get2 returns the cached output for a layer (stream 2), or nil if not cached.
+// Used for dual-stream architectures.
+func (c *StepCache) Get2(layer int) *mlx.Array {
+	if layer < len(c.layers2) {
+		return c.layers2[layer]
+	}
+	return nil
+}
+
+// Set2 stores a layer output (stream 2), freeing any previous value.
+// Used for dual-stream architectures.
+func (c *StepCache) Set2(layer int, arr *mlx.Array) {
+	if layer < len(c.layers2) {
+		if c.layers2[layer] != nil {
+			c.layers2[layer].Free()
+		}
+		c.layers2[layer] = arr
+	}
+}
+
+// GetConstant returns the cached constant value.
+func (c *StepCache) GetConstant() *mlx.Array {
+	return c.constant
+}
+
+// SetConstant stores a constant value, freeing any previous value.
+func (c *StepCache) SetConstant(arr *mlx.Array) {
+	if c.constant != nil {
+		c.constant.Free()
+	}
+	c.constant = arr
+}
+
+// Arrays returns all non-nil cached arrays (for pool.Keep).
+func (c *StepCache) Arrays() []*mlx.Array {
+	var result []*mlx.Array
+	if c.constant != nil {
+		result = append(result, c.constant)
+	}
+	for _, arr := range c.layers {
+		if arr != nil {
+			result = append(result, arr)
+		}
+	}
+	for _, arr := range c.layers2 {
+		if arr != nil {
+			result = append(result, arr)
+		}
+	}
+	return result
+}
+
+// Free releases all cached arrays. Call when generation completes.
+func (c *StepCache) Free() {
+	if c.constant != nil {
+		c.constant.Free()
+		c.constant = nil
+	}
+	for i, arr := range c.layers {
+		if arr != nil {
+			arr.Free()
+			c.layers[i] = nil
+		}
+	}
+	for i, arr := range c.layers2 {
+		if arr != nil {
+			arr.Free()
+			c.layers2[i] = nil
+		}
+	}
+}
+
+// NumLayers returns the number of layers this cache can store.
+func (c *StepCache) NumLayers() int {
+	return len(c.layers)
+}

x/imagegen/cache/teacache.go

@@ -0,0 +1,195 @@
+// Package cache provides caching mechanisms for diffusion model inference.
+package cache
+
+import (
+	"github.com/ollama/ollama/x/imagegen/mlx"
+)
+
+// TeaCache implements Timestep Embedding Aware Caching for diffusion models.
+// It caches the transformer output and reuses it when timestep values
+// are similar between consecutive steps.
+//
+// For CFG (classifier-free guidance), it caches pos and neg predictions
+// separately and always computes CFG fresh to avoid error amplification.
+//
+// Reference: "Timestep Embedding Tells: It's Time to Cache for Video Diffusion Model"
+// https://github.com/ali-vilab/TeaCache
+type TeaCache struct {
+	// Cached transformer output from last computed step (non-CFG mode)
+	cachedOutput *mlx.Array
+
+	// Cached CFG outputs (pos and neg separately)
+	cachedPosOutput *mlx.Array
+	cachedNegOutput *mlx.Array
+
+	// Previous timestep value for difference calculation
+	prevTimestep float32
+
+	// Accumulated difference for rescaling
+	accumulatedDiff float32
+
+	// Configuration
+	threshold      float32 // Threshold for recomputation decision
+	rescaleFactor  float32 // Model-specific rescaling factor
+	skipEarlySteps int     // Number of early steps to never cache
+
+	// Statistics
+	cacheHits   int
+	cacheMisses int
+}
+
+// TeaCacheConfig holds configuration for TeaCache.
+type TeaCacheConfig struct {
+	// Threshold for recomputation. Lower = more cache hits, potential quality loss.
+	// Recommended: 0.05-0.15 for image models
+	Threshold float32
+
+	// Rescale factor to adjust timestep embedding differences.
+	// Model-specific, typically 1.0-2.0
+	RescaleFactor float32
+
+	// SkipEarlySteps: number of early steps to always compute (never cache).
+	// Set to 2-3 for CFG mode to preserve structure. 0 = no skipping.
+	SkipEarlySteps int
+}
+
+// DefaultTeaCacheConfig returns default configuration for TeaCache.
+func DefaultTeaCacheConfig() *TeaCacheConfig {
+	return &TeaCacheConfig{
+		Threshold:     0.1,
+		RescaleFactor: 1.0,
+	}
+}
+
+// NewTeaCache creates a new TeaCache instance.
+func NewTeaCache(cfg *TeaCacheConfig) *TeaCache {
+	if cfg == nil {
+		cfg = DefaultTeaCacheConfig()
+	}
+	return &TeaCache{
+		threshold:      cfg.Threshold,
+		rescaleFactor:  cfg.RescaleFactor,
+		skipEarlySteps: cfg.SkipEarlySteps,
+	}
+}
+
+// ShouldCompute determines if we should compute the full forward pass
+// or reuse the cached output based on timestep similarity.
+//
+// Algorithm:
+// 1. First step always computes
+// 2. Subsequent steps compare |currTimestep - prevTimestep| * rescaleFactor
+// 3. If accumulated difference > threshold, compute new output
+// 4. Otherwise, reuse cached output
+func (tc *TeaCache) ShouldCompute(step int, timestep float32) bool {
+	// Always compute early steps (critical for structure)
+	// Check both regular cache and CFG cache
+	hasCachedOutput := tc.cachedOutput != nil || tc.HasCFGCache()
+	if step < tc.skipEarlySteps || step == 0 || !hasCachedOutput {
+		return true
+	}
+
+	// Compute absolute difference between current and previous timestep
+	diff := timestep - tc.prevTimestep
+	if diff < 0 {
+		diff = -diff
+	}
+
+	// Apply rescaling factor
+	scaledDiff := diff * tc.rescaleFactor
+
+	// Accumulate difference (helps track drift over multiple cached steps)
+	tc.accumulatedDiff += scaledDiff
+
+	// Decision based on accumulated difference
+	if tc.accumulatedDiff > tc.threshold {
+		tc.accumulatedDiff = 0 // Reset accumulator
+		return true
+	}
+
+	return false
+}
+
+// UpdateCache stores the computed output for potential reuse (non-CFG mode).
+func (tc *TeaCache) UpdateCache(output *mlx.Array, timestep float32) {
+	// Free previous cached output
+	if tc.cachedOutput != nil {
+		tc.cachedOutput.Free()
+	}
+
+	// Store new cached values
+	tc.cachedOutput = output
+	tc.prevTimestep = timestep
+	tc.cacheMisses++
+}
+
+// UpdateCFGCache stores pos and neg outputs separately for CFG mode.
+// This allows CFG to be computed fresh each step, avoiding error amplification.
+func (tc *TeaCache) UpdateCFGCache(posOutput, negOutput *mlx.Array, timestep float32) {
+	// Free previous cached outputs
+	if tc.cachedPosOutput != nil {
+		tc.cachedPosOutput.Free()
+	}
+	if tc.cachedNegOutput != nil {
+		tc.cachedNegOutput.Free()
+	}
+
+	// Store new cached values
+	tc.cachedPosOutput = posOutput
+	tc.cachedNegOutput = negOutput
+	tc.prevTimestep = timestep
+	tc.cacheMisses++
+}
+
+// GetCached returns the cached output (non-CFG mode).
+func (tc *TeaCache) GetCached() *mlx.Array {
+	tc.cacheHits++
+	return tc.cachedOutput
+}
+
+// GetCFGCached returns cached pos and neg outputs for CFG mode.
+func (tc *TeaCache) GetCFGCached() (pos, neg *mlx.Array) {
+	tc.cacheHits++
+	return tc.cachedPosOutput, tc.cachedNegOutput
+}
+
+// HasCFGCache returns true if CFG cache is available.
+func (tc *TeaCache) HasCFGCache() bool {
+	return tc.cachedPosOutput != nil && tc.cachedNegOutput != nil
+}
+
+// Arrays returns all arrays that should be kept alive.
+func (tc *TeaCache) Arrays() []*mlx.Array {
+	var arrays []*mlx.Array
+	if tc.cachedOutput != nil {
+		arrays = append(arrays, tc.cachedOutput)
+	}
+	if tc.cachedPosOutput != nil {
+		arrays = append(arrays, tc.cachedPosOutput)
+	}
+	if tc.cachedNegOutput != nil {
+		arrays = append(arrays, tc.cachedNegOutput)
+	}
+	return arrays
+}
+
+// Stats returns cache hit/miss statistics.
+func (tc *TeaCache) Stats() (hits, misses int) {
+	return tc.cacheHits, tc.cacheMisses
+}
+
+// Free releases all cached arrays.
+func (tc *TeaCache) Free() {
+	if tc.cachedOutput != nil {
+		tc.cachedOutput.Free()
+		tc.cachedOutput = nil
+	}
+	if tc.cachedPosOutput != nil {
+		tc.cachedPosOutput.Free()
+		tc.cachedPosOutput = nil
+	}
+	if tc.cachedNegOutput != nil {
+		tc.cachedNegOutput.Free()
+		tc.cachedNegOutput = nil
+	}
+}

x/imagegen/cli.go

@@ -0,0 +1,576 @@
+// cli.go provides CLI commands for image generation models.
+//
+// TODO (jmorganca): Integrate these commands into cmd/cmd.go when stable.
+// Currently these are separate to keep experimental code isolated.
+
+package imagegen
+
+import (
+	"encoding/base64"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"regexp"
+	"slices"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/spf13/cobra"
+
+	"github.com/ollama/ollama/api"
+	"github.com/ollama/ollama/envconfig"
+	"github.com/ollama/ollama/progress"
+	"github.com/ollama/ollama/readline"
+)
+
+// ImageGenOptions holds options for image generation.
+// These can be set via environment variables or interactive commands.
+type ImageGenOptions struct {
+	Width          int
+	Height         int
+	Steps          int
+	Seed           int
+	NegativePrompt string
+}
+
+// DefaultOptions returns the default image generation options.
+func DefaultOptions() ImageGenOptions {
+	return ImageGenOptions{
+		Width:  1024,
+		Height: 1024,
+		Steps:  0, // 0 means model default
+		Seed:   0, // 0 means random
+	}
+}
+
+// RegisterFlags adds image generation flags to the given command.
+// Flags are hidden since they only apply to image generation models.
+func RegisterFlags(cmd *cobra.Command) {
+	cmd.Flags().Int("width", 1024, "Image width")
+	cmd.Flags().Int("height", 1024, "Image height")
+	cmd.Flags().Int("steps", 0, "Denoising steps (0 = model default)")
+	cmd.Flags().Int("seed", 0, "Random seed (0 for random)")
+	cmd.Flags().String("negative", "", "Negative prompt")
+	// Hide from main flags section - shown in separate section via AppendFlagsDocs
+	cmd.Flags().MarkHidden("width")
+	cmd.Flags().MarkHidden("height")
+	cmd.Flags().MarkHidden("steps")
+	cmd.Flags().MarkHidden("seed")
+	cmd.Flags().MarkHidden("negative")
+}
+
+// AppendFlagsDocs appends image generation flags documentation to the command's usage template.
+func AppendFlagsDocs(cmd *cobra.Command) {
+	usage := `
+Image Generation Flags (experimental):
+      --width int      Image width
+      --height int     Image height
+      --steps int      Denoising steps
+      --seed int       Random seed
+      --negative str   Negative prompt
+`
+	cmd.SetUsageTemplate(cmd.UsageTemplate() + usage)
+}
+
+// RunCLI handles the CLI for image generation models.
+// Returns true if it handled the request, false if the caller should continue with normal flow.
+// Supports flags: --width, --height, --steps, --seed, --negative
+// Image paths can be included in the prompt and will be extracted automatically.
+func RunCLI(cmd *cobra.Command, name string, prompt string, interactive bool, keepAlive *api.Duration) error {
+	// Get options from flags (with env var defaults)
+	opts := DefaultOptions()
+	if cmd != nil && cmd.Flags() != nil {
+		if v, err := cmd.Flags().GetInt("width"); err == nil && v > 0 {
+			opts.Width = v
+		}
+		if v, err := cmd.Flags().GetInt("height"); err == nil && v > 0 {
+			opts.Height = v
+		}
+		if v, err := cmd.Flags().GetInt("steps"); err == nil && v > 0 {
+			opts.Steps = v
+		}
+		if v, err := cmd.Flags().GetInt("seed"); err == nil && v != 0 {
+			opts.Seed = v
+		}
+		if v, err := cmd.Flags().GetString("negative"); err == nil && v != "" {
+			opts.NegativePrompt = v
+		}
+	}
+
+	if interactive {
+		return runInteractive(cmd, name, keepAlive, opts)
+	}
+
+	// One-shot generation
+	return generateImageWithOptions(cmd, name, prompt, keepAlive, opts)
+}
+
+// generateImageWithOptions generates an image with the given options.
+func generateImageWithOptions(cmd *cobra.Command, modelName, prompt string, keepAlive *api.Duration, opts ImageGenOptions) error {
+	client, err := api.ClientFromEnvironment()
+	if err != nil {
+		return err
+	}
+
+	// Extract any image paths from the prompt
+	prompt, images, err := extractFileData(prompt)
+	if err != nil {
+		return err
+	}
+
+	req := &api.GenerateRequest{
+		Model:  modelName,
+		Prompt: prompt,
+		Images: images,
+		Width:  int32(opts.Width),
+		Height: int32(opts.Height),
+		Steps:  int32(opts.Steps),
+	}
+	if opts.Seed != 0 {
+		req.Options = map[string]any{"seed": opts.Seed}
+	}
+	if keepAlive != nil {
+		req.KeepAlive = keepAlive
+	}
+
+	// Show loading spinner until generation starts
+	p := progress.NewProgress(os.Stderr)
+	spinner := progress.NewSpinner("")
+	p.Add("", spinner)
+
+	var stepBar *progress.StepBar
+	var imageBase64 string
+	err = client.Generate(cmd.Context(), req, func(resp api.GenerateResponse) error {
+		// Handle progress updates using structured fields
+		if resp.Total > 0 {
+			if stepBar == nil {
+				spinner.Stop()
+				stepBar = progress.NewStepBar("Generating", int(resp.Total))
+				p.Add("", stepBar)
+			}
+			stepBar.Set(int(resp.Completed))
+		}
+
+		// Handle final response with image data
+		if resp.Done && resp.Image != "" {
+			imageBase64 = resp.Image
+		}
+
+		return nil
+	})
+
+	p.StopAndClear()
+	if err != nil {
+		return err
+	}
+
+	if imageBase64 != "" {
+		// Decode base64 and save to CWD
+		imageData, err := base64.StdEncoding.DecodeString(imageBase64)
+		if err != nil {
+			return fmt.Errorf("failed to decode image: %w", err)
+		}
+
+		// Create filename from prompt
+		safeName := sanitizeFilename(prompt)
+		if len(safeName) > 50 {
+			safeName = safeName[:50]
+		}
+		timestamp := time.Now().Format("20060102-150405")
+		filename := fmt.Sprintf("%s-%s.png", safeName, timestamp)
+
+		if err := os.WriteFile(filename, imageData, 0o644); err != nil {
+			return fmt.Errorf("failed to save image: %w", err)
+		}
+
+		displayImageInTerminal(filename)
+		fmt.Printf("Image saved to: %s\n", filename)
+	}
+
+	return nil
+}
+
+// runInteractive runs an interactive REPL for image generation.
+func runInteractive(cmd *cobra.Command, modelName string, keepAlive *api.Duration, opts ImageGenOptions) error {
+	client, err := api.ClientFromEnvironment()
+	if err != nil {
+		return err
+	}
+
+	// Preload the model with the specified keepalive
+	p := progress.NewProgress(os.Stderr)
+	spinner := progress.NewSpinner("")
+	p.Add("", spinner)
+
+	preloadReq := &api.GenerateRequest{
+		Model:     modelName,
+		KeepAlive: keepAlive,
+	}
+	if err := client.Generate(cmd.Context(), preloadReq, func(resp api.GenerateResponse) error {
+		return nil
+	}); err != nil {
+		p.StopAndClear()
+		return fmt.Errorf("failed to load model: %w", err)
+	}
+	p.StopAndClear()
+
+	scanner, err := readline.New(readline.Prompt{
+		Prompt:      ">>> ",
+		Placeholder: "Describe an image to generate (/help for commands)",
+	})
+	if err != nil {
+		return err
+	}
+
+	if envconfig.NoHistory() {
+		scanner.HistoryDisable()
+	}
+
+	for {
+		line, err := scanner.Readline()
+		switch {
+		case errors.Is(err, io.EOF):
+			fmt.Println()
+			return nil
+		case errors.Is(err, readline.ErrInterrupt):
+			if line == "" {
+				fmt.Println("\nUse Ctrl + d or /bye to exit.")
+			}
+			continue
+		case err != nil:
+			return err
+		}
+
+		line = strings.TrimSpace(line)
+		if line == "" {
+			continue
+		}
+
+		// Handle commands
+		switch {
+		case strings.HasPrefix(line, "/bye"):
+			return nil
+		case strings.HasPrefix(line, "/?"), strings.HasPrefix(line, "/help"):
+			printInteractiveHelp()
+			continue
+		case strings.HasPrefix(line, "/set "):
+			if err := handleSetCommand(line[5:], &opts); err != nil {
+				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
+			}
+			continue
+		case strings.HasPrefix(line, "/show"):
+			printCurrentSettings(opts)
+			continue
+		case strings.HasPrefix(line, "/"):
+			// Check if it's a file path, not a command
+			args := strings.Fields(line)
+			isFile := false
+			for _, f := range extractFileNames(line) {
+				if strings.HasPrefix(f, args[0]) {
+					isFile = true
+					break
+				}
+			}
+			if !isFile {
+				fmt.Fprintf(os.Stderr, "Unknown command: %s (try /help)\n", args[0])
+				continue
+			}
+		}
+
+		// Extract any image paths from the input
+		prompt, images, err := extractFileData(line)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
+			continue
+		}
+
+		// Generate image with current options
+		req := &api.GenerateRequest{
+			Model:  modelName,
+			Prompt: prompt,
+			Images: images,
+			Width:  int32(opts.Width),
+			Height: int32(opts.Height),
+			Steps:  int32(opts.Steps),
+		}
+		if opts.Seed != 0 {
+			req.Options = map[string]any{"seed": opts.Seed}
+		}
+		if keepAlive != nil {
+			req.KeepAlive = keepAlive
+		}
+
+		// Show loading spinner until generation starts
+		p := progress.NewProgress(os.Stderr)
+		spinner := progress.NewSpinner("")
+		p.Add("", spinner)
+
+		var stepBar *progress.StepBar
+		var imageBase64 string
+
+		err = client.Generate(cmd.Context(), req, func(resp api.GenerateResponse) error {
+			// Handle progress updates using structured fields
+			if resp.Total > 0 {
+				if stepBar == nil {
+					spinner.Stop()
+					stepBar = progress.NewStepBar("Generating", int(resp.Total))
+					p.Add("", stepBar)
+				}
+				stepBar.Set(int(resp.Completed))
+			}
+
+			// Handle final response with image data
+			if resp.Done && resp.Image != "" {
+				imageBase64 = resp.Image
+			}
+
+			return nil
+		})
+
+		p.StopAndClear()
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
+			continue
+		}
+
+		// Save image to current directory with descriptive name
+		if imageBase64 != "" {
+			// Decode base64 image data
+			imageData, err := base64.StdEncoding.DecodeString(imageBase64)
+			if err != nil {
+				fmt.Fprintf(os.Stderr, "Error decoding image: %v\n", err)
+				continue
+			}
+
+			// Create filename from prompt (sanitized)
+			safeName := sanitizeFilename(line)
+			if len(safeName) > 50 {
+				safeName = safeName[:50]
+			}
+			timestamp := time.Now().Format("20060102-150405")
+			filename := fmt.Sprintf("%s-%s.png", safeName, timestamp)
+
+			if err := os.WriteFile(filename, imageData, 0o644); err != nil {
+				fmt.Fprintf(os.Stderr, "Error saving image: %v\n", err)
+				continue
+			}
+
+			displayImageInTerminal(filename)
+			fmt.Printf("Image saved to: %s\n", filename)
+		}
+
+		fmt.Println()
+	}
+}
+
+// sanitizeFilename removes characters that aren't safe for filenames.
+func sanitizeFilename(s string) string {
+	s = strings.ToLower(s)
+	s = strings.ReplaceAll(s, " ", "-")
+	// Remove any character that's not alphanumeric or hyphen
+	var result strings.Builder
+	for _, r := range s {
+		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '-' {
+			result.WriteRune(r)
+		}
+	}
+	return result.String()
+}
+
+// printInteractiveHelp prints help for interactive mode commands.
+// TODO: reconcile /set commands with /set parameter in text gen REPL (cmd/cmd.go)
+func printInteractiveHelp() {
+	fmt.Fprintln(os.Stderr, "Commands:")
+	fmt.Fprintln(os.Stderr, "  /set width <n>     Set image width")
+	fmt.Fprintln(os.Stderr, "  /set height <n>    Set image height")
+	fmt.Fprintln(os.Stderr, "  /set steps <n>     Set denoising steps")
+	fmt.Fprintln(os.Stderr, "  /set seed <n>      Set random seed")
+	fmt.Fprintln(os.Stderr, "  /set negative <s>  Set negative prompt")
+	fmt.Fprintln(os.Stderr, "  /show              Show current settings")
+	fmt.Fprintln(os.Stderr, "  /bye               Exit")
+	fmt.Fprintln(os.Stderr)
+	fmt.Fprintln(os.Stderr, "Or type a prompt to generate an image.")
+	fmt.Fprintln(os.Stderr)
+}
+
+// printCurrentSettings prints the current image generation settings.
+func printCurrentSettings(opts ImageGenOptions) {
+	fmt.Fprintf(os.Stderr, "Current settings:\n")
+	fmt.Fprintf(os.Stderr, "  width:    %d\n", opts.Width)
+	fmt.Fprintf(os.Stderr, "  height:   %d\n", opts.Height)
+	fmt.Fprintf(os.Stderr, "  steps:    %d\n", opts.Steps)
+	fmt.Fprintf(os.Stderr, "  seed:     %d (0=random)\n", opts.Seed)
+	if opts.NegativePrompt != "" {
+		fmt.Fprintf(os.Stderr, "  negative: %s\n", opts.NegativePrompt)
+	}
+	fmt.Fprintln(os.Stderr)
+}
+
+// handleSetCommand handles /set commands to change options.
+func handleSetCommand(args string, opts *ImageGenOptions) error {
+	parts := strings.SplitN(args, " ", 2)
+	if len(parts) < 2 {
+		return fmt.Errorf("usage: /set <option> <value>")
+	}
+
+	key := strings.ToLower(parts[0])
+	value := strings.TrimSpace(parts[1])
+
+	switch key {
+	case "width", "w":
+		v, err := strconv.Atoi(value)
+		if err != nil || v <= 0 {
+			return fmt.Errorf("width must be a positive integer")
+		}
+		opts.Width = v
+		fmt.Fprintf(os.Stderr, "Set width to %d\n", v)
+	case "height", "h":
+		v, err := strconv.Atoi(value)
+		if err != nil || v <= 0 {
+			return fmt.Errorf("height must be a positive integer")
+		}
+		opts.Height = v
+		fmt.Fprintf(os.Stderr, "Set height to %d\n", v)
+	case "steps", "s":
+		v, err := strconv.Atoi(value)
+		if err != nil || v <= 0 {
+			return fmt.Errorf("steps must be a positive integer")
+		}
+		opts.Steps = v
+		fmt.Fprintf(os.Stderr, "Set steps to %d\n", v)
+	case "seed":
+		v, err := strconv.Atoi(value)
+		if err != nil {
+			return fmt.Errorf("seed must be an integer")
+		}
+		opts.Seed = v
+		fmt.Fprintf(os.Stderr, "Set seed to %d\n", v)
+	case "negative", "neg", "n":
+		opts.NegativePrompt = value
+		if value == "" {
+			fmt.Fprintln(os.Stderr, "Cleared negative prompt")
+		} else {
+			fmt.Fprintf(os.Stderr, "Set negative prompt to: %s\n", value)
+		}
+	default:
+		return fmt.Errorf("unknown option: %s (try /help)", key)
+	}
+	return nil
+}
+
+// displayImageInTerminal attempts to render an image inline in the terminal.
+// Supports iTerm2, Kitty, WezTerm, Ghostty, and other terminals with inline image support.
+// Returns true if the image was displayed, false otherwise.
+func displayImageInTerminal(imagePath string) bool {
+	// Check if terminal supports inline images
+	termProgram := os.Getenv("TERM_PROGRAM")
+	kittyWindowID := os.Getenv("KITTY_WINDOW_ID")
+	weztermPane := os.Getenv("WEZTERM_PANE")
+	ghostty := os.Getenv("GHOSTTY_RESOURCES_DIR")
+
+	// Read the image file
+	data, err := os.ReadFile(imagePath)
+	if err != nil {
+		return false
+	}
+
+	encoded := base64.StdEncoding.EncodeToString(data)
+
+	switch {
+	case termProgram == "iTerm.app" || termProgram == "WezTerm" || weztermPane != "":
+		// iTerm2/WezTerm inline image protocol
+		// ESC ] 1337 ; File = [arguments] : base64 BEL
+		fmt.Printf("\033]1337;File=inline=1;preserveAspectRatio=1:%s\a\n", encoded)
+		return true
+
+	case kittyWindowID != "" || ghostty != "" || termProgram == "ghostty":
+		// Kitty graphics protocol (also used by Ghostty)
+		// Send in chunks for large images
+		const chunkSize = 4096
+		for i := 0; i < len(encoded); i += chunkSize {
+			end := min(i+chunkSize, len(encoded))
+			chunk := encoded[i:end]
+
+			if i == 0 {
+				// First chunk: a=T (transmit), f=100 (PNG), m=1 (more chunks follow) or m=0 (last chunk)
+				more := 1
+				if end >= len(encoded) {
+					more = 0
+				}
+				fmt.Printf("\033_Ga=T,f=100,m=%d;%s\033\\", more, chunk)
+			} else if end >= len(encoded) {
+				// Last chunk
+				fmt.Printf("\033_Gm=0;%s\033\\", chunk)
+			} else {
+				// Middle chunk
+				fmt.Printf("\033_Gm=1;%s\033\\", chunk)
+			}
+		}
+		fmt.Println()
+		return true
+
+	default:
+		return false
+	}
+}
+
+// extractFileNames finds image file paths in the input string.
+func extractFileNames(input string) []string {
+	// Regex to match file paths with image extensions
+	regexPattern := `(?:[a-zA-Z]:)?(?:\./|/|\\)[\S\\ ]+?\.(?i:jpg|jpeg|png|webp)\b`
+	re := regexp.MustCompile(regexPattern)
+	return re.FindAllString(input, -1)
+}
+
+// extractFileData extracts image data from file paths found in the input.
+// Returns the cleaned prompt (with file paths removed) and the image data.
+func extractFileData(input string) (string, []api.ImageData, error) {
+	filePaths := extractFileNames(input)
+	var imgs []api.ImageData
+
+	for _, fp := range filePaths {
+		// Normalize shell escapes
+		nfp := strings.ReplaceAll(fp, "\\ ", " ")
+		nfp = strings.ReplaceAll(nfp, "\\(", "(")
+		nfp = strings.ReplaceAll(nfp, "\\)", ")")
+		nfp = strings.ReplaceAll(nfp, "%20", " ")
+
+		data, err := getImageData(nfp)
+		if errors.Is(err, os.ErrNotExist) {
+			continue
+		} else if err != nil {
+			return "", nil, err
+		}
+		fmt.Fprintf(os.Stderr, "Added image '%s'\n", nfp)
+		input = strings.ReplaceAll(input, fp, "")
+		imgs = append(imgs, data)
+	}
+	return strings.TrimSpace(input), imgs, nil
+}
+
+// getImageData reads and validates image data from a file.
+func getImageData(filePath string) ([]byte, error) {
+	file, err := os.Open(filePath)
+	if err != nil {
+		return nil, err
+	}
+	defer file.Close()
+
+	buf := make([]byte, 512)
+	_, err = file.Read(buf)
+	if err != nil {
+		return nil, err
+	}
+
+	contentType := http.DetectContentType(buf)
+	allowedTypes := []string{"image/jpeg", "image/jpg", "image/png", "image/webp"}
+	if !slices.Contains(allowedTypes, contentType) {
+		return nil, fmt.Errorf("invalid image type: %s", contentType)
+	}
+
+	// Re-read the full file
+	return os.ReadFile(filePath)
+}

x/imagegen/cmd/engine/README.md

@@ -0,0 +1,25 @@
+# MLX Engine
+
+Experimental MLX backend for running models on Apple Silicon and CUDA.
+
+## Build
+
+```bash
+go build -o engine ./x/imagegen/cmd/engine
+```
+
+## Text Generation
+
+Text generation models are no longer supported by this engine.
+
+## Image Generation
+
+```bash
+./engine -zimage -model /path/to/z-image -prompt "a cat" -output cat.png
+```
+
+Options:
+
+- `-width`, `-height` - image dimensions (default 1024x1024)
+- `-steps` - denoising steps (default 9)
+- `-seed` - random seed (default 42)

x/imagegen/cmd/engine/generate.go

@@ -0,0 +1,357 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"time"
+	"unicode/utf8"
+
+	"github.com/ollama/ollama/x/imagegen/cache"
+	"github.com/ollama/ollama/x/imagegen/mlx"
+	"github.com/ollama/ollama/x/imagegen/tokenizer"
+)
+
+// Dedicated stream for generation (like mlx-lm's generation_stream)
+var generationStream *mlx.Stream
+
+// utf8Streamer buffers decoded text and emits only complete UTF-8 characters.
+// This handles cases where tokenizers output partial multi-byte sequences.
+type utf8Streamer struct {
+	buffer []byte
+}
+
+// Write adds decoded text to the buffer and returns complete UTF-8 characters.
+func (s *utf8Streamer) Write(text string) string {
+	s.buffer = append(s.buffer, text...)
+
+	// Find the last position that ends with a complete UTF-8 character
+	validLen := 0
+	for i := 0; i < len(s.buffer); {
+		r, size := utf8.DecodeRune(s.buffer[i:])
+		if r == utf8.RuneError && size == 1 {
+			// Invalid or incomplete UTF-8 sequence at this position
+			// Check if it could be a valid start of a multi-byte sequence
+			if len(s.buffer)-i < 4 {
+				// Might be incomplete, keep it in buffer
+				break
+			}
+			// Definitely invalid, skip this byte
+			i++
+			validLen = i
+		} else {
+			i += size
+			validLen = i
+		}
+	}
+
+	if validLen == 0 {
+		return ""
+	}
+
+	result := string(s.buffer[:validLen])
+	s.buffer = s.buffer[validLen:]
+	return result
+}
+
+// Flush returns any remaining buffered bytes (may be incomplete UTF-8).
+func (s *utf8Streamer) Flush() string {
+	if len(s.buffer) == 0 {
+		return ""
+	}
+	result := string(s.buffer)
+	s.buffer = nil
+	return result
+}
+
+// withStream runs fn with the generation stream as default
+func withStream(fn func()) {
+	// Lazy initialization of generationStream
+	if generationStream == nil {
+		generationStream = mlx.NewStream()
+	}
+	orig := mlx.GetDefaultStream()
+	mlx.SetDefaultStream(generationStream)
+	fn()
+	mlx.SetDefaultStream(orig)
+}
+
+type Model interface {
+	Tokenizer() *tokenizer.Tokenizer
+	VocabSize() int32
+	NewCache(maxSeqLen int32) []cache.Cache
+	Forward(input *mlx.Array, caches []cache.Cache) *mlx.Array
+}
+
+// ChatModel is an optional interface for models that support chat formatting
+type ChatModel interface {
+	FormatPrompt(prompt string) string
+}
+
+// MultimodalModel is for models that support image input
+type MultimodalModel interface {
+	Model
+	FormatPromptWithImage(prompt string) string
+	ExpandImageTokens(tokens []int32) []int32
+	ForwardWithImage(tokens *mlx.Array, image *mlx.Array, caches []cache.Cache) *mlx.Array
+	ImageSize() int32 // Returns expected image size for preprocessing
+}
+
+// ImageLoader loads and preprocesses an image for multimodal models
+// Returns nil if path is empty
+type ImageLoader func(path string, imageSize int32) (*mlx.Array, error)
+
+type input struct {
+	Prompt       string
+	Image        *mlx.Array // Optional preprocessed image for multimodal models
+	MaxTokens    int
+	Temperature  float32
+	TopP         float32
+	TopK         int
+	WiredLimitGB int // Metal wired memory limit in GB (default 32)
+}
+
+type output struct {
+	Text          string
+	Done          bool
+	PrefillTokSec float64
+	GenTokSec     float64
+}
+
+// Decoder wraps model + cache for autoregressive generation.
+type Decoder struct {
+	model         Model
+	caches        []cache.Cache
+	vocabSize     int32
+	temp          float32
+	topK          int
+	topP          float32
+	token         *mlx.Array   // Current token (kept across pools)
+	oldCacheState []*mlx.Array // Preallocated slice for old cache state
+	image         *mlx.Array   // Optional image for multimodal prefill
+}
+
+func NewDecoder(m Model, temp float32, topK int, topP float32) *Decoder {
+	caches := m.NewCache(0)
+	return &Decoder{
+		model:         m,
+		caches:        caches,
+		vocabSize:     m.VocabSize(),
+		temp:          temp,
+		topK:          topK,
+		topP:          topP,
+		oldCacheState: make([]*mlx.Array, 0, len(caches)*2),
+	}
+}
+
+// SetImage sets the image for multimodal prefill (call before prefill)
+func (d *Decoder) SetImage(img *mlx.Array) {
+	d.image = img
+}
+
+func (d *Decoder) prefill(inputIDs []int32) int {
+	processed := 0
+
+	// Track old cache state to free after each chunk
+	var oldCacheState []*mlx.Array
+
+	// For multimodal models with an image, we need to process all tokens together
+	// in the first forward pass so the image embeddings can be inserted properly.
+	// Skip chunking for multimodal prefill.
+	isMultimodal := d.image != nil
+
+	// Process all-but-1 tokens in chunks, eval cache state for memory management
+	// Skip chunking for multimodal - process everything in the final step
+	if !isMultimodal {
+		for len(inputIDs) > 1 {
+			chunkSize := min(2048, len(inputIDs)-1)
+			if chunkSize <= 0 {
+				break
+			}
+			chunk := inputIDs[:chunkSize]
+
+			// Save old cache state before forward
+			oldCacheState = oldCacheState[:0]
+			for _, c := range d.caches {
+				oldCacheState = append(oldCacheState, c.State()...)
+			}
+
+			var cacheState []*mlx.Array
+			withStream(func() {
+				x := mlx.NewArrayInt32(chunk, []int32{1, int32(len(chunk))})
+				d.model.Forward(x, d.caches)
+				for _, c := range d.caches {
+					cacheState = append(cacheState, c.State()...)
+				}
+			})
+			mlx.Eval(cacheState...)
+
+			// Free old cache state
+			for _, arr := range oldCacheState {
+				if arr != nil {
+					arr.Free()
+				}
+			}
+
+			inputIDs = inputIDs[chunkSize:]
+			processed += chunkSize
+		}
+	}
+
+	// Save old cache state before final step
+	oldCacheState = oldCacheState[:0]
+	for _, c := range d.caches {
+		oldCacheState = append(oldCacheState, c.State()...)
+	}
+
+	// Final token + sampling (or all tokens for multimodal)
+	withStream(func() {
+		x := mlx.NewArrayInt32(inputIDs, []int32{1, int32(len(inputIDs))})
+		mlx.Eval(x) // Materialize before any other evals
+
+		var logits *mlx.Array
+		// Use ForwardWithImage if we have an image and model supports it
+		if d.image != nil {
+			if mm, ok := d.model.(MultimodalModel); ok {
+				logits = mm.ForwardWithImage(x, d.image, d.caches)
+				d.image = nil // Only use image for first forward
+			} else {
+				logits = d.model.Forward(x, d.caches)
+			}
+		} else {
+			logits = d.model.Forward(x, d.caches)
+		}
+		d.token = sample(logits, d.temp, d.topK, d.topP, d.vocabSize)
+	})
+	// Keep cache state (token auto-kept by AsyncEval)
+	for _, c := range d.caches {
+		mlx.Keep(c.State()...)
+	}
+	mlx.AsyncEval(d.token)
+
+	// Free old cache state from before final step
+	for _, arr := range oldCacheState {
+		if arr != nil {
+			arr.Free()
+		}
+	}
+
+	mlx.ClearCache()
+
+	return processed + len(inputIDs)
+}
+
+func (d *Decoder) step() int32 {
+	prevToken := d.token
+
+	// Save old cache state (reuse preallocated slice)
+	d.oldCacheState = d.oldCacheState[:0]
+	for _, c := range d.caches {
+		d.oldCacheState = append(d.oldCacheState, c.State()...)
+	}
+
+	withStream(func() {
+		logits := d.model.Forward(mlx.Reshape(prevToken, 1, 1), d.caches)
+		d.token = sample(logits, d.temp, d.topK, d.topP, d.vocabSize)
+	})
+	// Keep token and new cache state so they survive cleanup
+	mlx.Keep(d.token)
+	for _, c := range d.caches {
+		mlx.Keep(c.State()...)
+	}
+	mlx.AsyncEval(d.token)
+
+	// Sync on previous token (GPU already working on next step)
+	val := prevToken.ItemInt32()
+
+	// Free old token and old cache state
+	prevToken.Free()
+	for _, arr := range d.oldCacheState {
+		arr.Free()
+	}
+	return val
+}
+
+func generate(ctx context.Context, m Model, in input, cb func(output)) error {
+	mlx.EnableCompile()
+	wiredLimit := in.WiredLimitGB
+	if wiredLimit <= 0 {
+		wiredLimit = 32 // default 32GB
+	}
+	mlx.MetalSetWiredLimit(uint64(wiredLimit) << 30)
+
+	temp := in.Temperature
+	if temp < 0 {
+		temp = 0.7
+	}
+
+	tok := m.Tokenizer()
+	dec := NewDecoder(m, temp, in.TopK, in.TopP)
+
+	// Apply chat template - use image template if we have an image
+	prompt := in.Prompt
+	var tokens []int32
+	if mm, ok := m.(MultimodalModel); ok && in.Image != nil {
+		prompt = mm.FormatPromptWithImage(prompt)
+		tokens = tok.Encode(prompt, true)
+		tokens = mm.ExpandImageTokens(tokens) // Expand <start_of_image> to 256 image tokens
+		dec.SetImage(in.Image)
+	} else if cm, ok := m.(ChatModel); ok {
+		prompt = cm.FormatPrompt(prompt)
+		tokens = tok.Encode(prompt, true)
+	} else {
+		tokens = tok.Encode(prompt, true)
+	}
+
+	prefillStart := time.Now()
+	prefillTokens := dec.prefill(tokens)
+	// Prefill measurement should include time to first token (like mlx-lm)
+	// Step() waits for prefill to complete and returns first token
+	firstToken := dec.step()
+	prefillTokSec := float64(prefillTokens) / time.Since(prefillStart).Seconds()
+
+	genStart := time.Now()
+	maxTokens := max(in.MaxTokens, 100)
+	var genTokens int
+
+	// UTF-8 streamer to handle partial multi-byte characters
+	streamer := &utf8Streamer{}
+
+	// Handle first token
+	genTokens++
+	if tok.IsEOS(firstToken) {
+		cb(output{Done: true, PrefillTokSec: prefillTokSec, GenTokSec: 0})
+		return nil
+	}
+	if text := streamer.Write(tok.Decode([]int32{firstToken})); text != "" {
+		cb(output{Text: text})
+	}
+
+	for n := 1; n < maxTokens; n++ {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+		token := dec.step()
+		genTokens++
+
+		if tok.IsEOS(token) {
+			break
+		}
+		if text := streamer.Write(tok.Decode([]int32{token})); text != "" {
+			cb(output{Text: text})
+		}
+
+		if n%256 == 0 {
+			mlx.ClearCache()
+		}
+	}
+
+	// Flush any remaining buffered bytes
+	if text := streamer.Flush(); text != "" {
+		cb(output{Text: text})
+	}
+
+	fmt.Printf("\nPeak memory: %.2fGB\n", float64(mlx.MetalGetPeakMemory())/(1<<30))
+	cb(output{Done: true, PrefillTokSec: prefillTokSec,
+		GenTokSec: float64(genTokens) / time.Since(genStart).Seconds()})
+	return nil
+}

x/imagegen/cmd/engine/image.go

@@ -0,0 +1,87 @@
+package main
+
+import (
+	"fmt"
+	"image"
+	"image/png"
+	"os"
+	"path/filepath"
+
+	"github.com/ollama/ollama/x/imagegen/mlx"
+)
+
+// saveImageArray saves an MLX array as a PNG image.
+// Expected format: [B, C, H, W] with values in [0, 1] range and C=3 (RGB).
+func saveImageArray(arr *mlx.Array, path string) error {
+	img, err := arrayToImage(arr)
+	if err != nil {
+		return err
+	}
+	return savePNG(img, path)
+}
+
+func savePNG(img *image.RGBA, path string) error {
+	if filepath.Ext(path) != ".png" {
+		path = path + ".png"
+	}
+	f, err := os.Create(path)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+	return png.Encode(f, img)
+}
+
+func arrayToImage(arr *mlx.Array) (*image.RGBA, error) {
+	shape := arr.Shape()
+	if len(shape) != 4 {
+		return nil, fmt.Errorf("expected 4D array [B, C, H, W], got %v", shape)
+	}
+
+	// Transform to [H, W, C] for image conversion
+	img := mlx.Squeeze(arr, 0)
+	arr.Free()
+	img = mlx.Transpose(img, 1, 2, 0)
+	img = mlx.Contiguous(img)
+	mlx.Eval(img)
+
+	imgShape := img.Shape()
+	H := int(imgShape[0])
+	W := int(imgShape[1])
+	C := int(imgShape[2])
+
+	if C != 3 {
+		img.Free()
+		return nil, fmt.Errorf("expected 3 channels (RGB), got %d", C)
+	}
+
+	// Copy to CPU and free GPU memory
+	data := img.Data()
+	img.Free()
+
+	// Write directly to Pix slice (faster than SetRGBA)
+	goImg := image.NewRGBA(image.Rect(0, 0, W, H))
+	pix := goImg.Pix
+	for y := 0; y < H; y++ {
+		for x := 0; x < W; x++ {
+			srcIdx := (y*W + x) * C
+			dstIdx := (y*W + x) * 4
+			pix[dstIdx+0] = uint8(clampF(data[srcIdx+0]*255+0.5, 0, 255))
+			pix[dstIdx+1] = uint8(clampF(data[srcIdx+1]*255+0.5, 0, 255))
+			pix[dstIdx+2] = uint8(clampF(data[srcIdx+2]*255+0.5, 0, 255))
+			pix[dstIdx+3] = 255
+		}
+	}
+
+	return goImg, nil
+}
+
+func clampF(v, min, max float32) float32 {
+	if v < min {
+		return min
+	}
+	if v > max {
+		return max
+	}
+	return v
+}

x/imagegen/cmd/engine/main.go

@@ -0,0 +1,287 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"image"
+	_ "image/jpeg"
+	_ "image/png"
+	"log"
+	"os"
+	"path/filepath"
+	"runtime/pprof"
+
+	"github.com/ollama/ollama/x/imagegen"
+	"github.com/ollama/ollama/x/imagegen/mlx"
+	"github.com/ollama/ollama/x/imagegen/models/flux2"
+	"github.com/ollama/ollama/x/imagegen/models/zimage"
+	"github.com/ollama/ollama/x/imagegen/safetensors"
+)
+
+// stringSlice is a flag type that accumulates multiple values
+type stringSlice []string
+
+func (s *stringSlice) String() string {
+	return fmt.Sprintf("%v", *s)
+}
+
+func (s *stringSlice) Set(value string) error {
+	*s = append(*s, value)
+	return nil
+}
+
+func main() {
+	modelPath := flag.String("model", "", "Model directory")
+	prompt := flag.String("prompt", "Hello", "Prompt")
+
+	// Text generation params
+	maxTokens := flag.Int("max-tokens", 100, "Max tokens")
+	temperature := flag.Float64("temperature", 0.7, "Temperature")
+	topP := flag.Float64("top-p", 0.9, "Top-p sampling")
+	topK := flag.Int("top-k", 40, "Top-k sampling")
+	imagePath := flag.String("image", "", "Image path for multimodal models")
+
+	// Image generation params
+	width := flag.Int("width", 0, "Image width (0 = auto from input or 1024)")
+	height := flag.Int("height", 0, "Image height (0 = auto from input or 1024)")
+	steps := flag.Int("steps", 0, "Denoising steps (0 = model default)")
+	seed := flag.Int64("seed", 42, "Random seed")
+	out := flag.String("output", "output.png", "Output path")
+
+	// Utility flags
+	listTensors := flag.Bool("list", false, "List tensors only")
+	cpuProfile := flag.String("cpuprofile", "", "Write CPU profile to file")
+	gpuCapture := flag.String("gpu-capture", "", "Capture GPU trace to .gputrace file (run with MTL_CAPTURE_ENABLED=1)")
+	wiredLimitGB := flag.Int("wired-limit", 32, "Metal wired memory limit in GB")
+
+	// Legacy mode flags
+	zimageFlag := flag.Bool("zimage", false, "Z-Image generation")
+	flux2Flag := flag.Bool("flux2", false, "FLUX.2 Klein generation")
+	var inputImages stringSlice
+	flag.Var(&inputImages, "input-image", "Input image for image editing (can be specified multiple times)")
+	negativePrompt := flag.String("negative-prompt", "", "Negative prompt for CFG (empty = no CFG, matching Python)")
+	cfgScale := flag.Float64("cfg-scale", 4.0, "CFG scale for image editing")
+	teaCache := flag.Bool("teacache", false, "Enable TeaCache for faster inference")
+	teaCacheThreshold := flag.Float64("teacache-threshold", 0.1, "TeaCache threshold (lower = more aggressive caching)")
+	fusedQKV := flag.Bool("fused-qkv", false, "Enable fused QKV projection for faster attention")
+
+	flag.Parse()
+
+	if *modelPath == "" {
+		flag.Usage()
+		return
+	}
+
+	// Check if MLX initialized successfully
+	if !mlx.IsMLXAvailable() {
+		log.Fatalf("MLX initialization failed: %v", mlx.GetMLXInitError())
+	}
+
+	// Restore strict error handling now that we know MLX is working.
+	// During init(), a safe handler prevented exit(-1) on GPU errors.
+	mlx.RestoreDefaultErrorHandler()
+
+	// CPU profiling
+	if *cpuProfile != "" {
+		f, err := os.Create(*cpuProfile)
+		if err != nil {
+			log.Fatal(err)
+		}
+		defer f.Close()
+		if err := pprof.StartCPUProfile(f); err != nil {
+			log.Fatal(err)
+		}
+		defer pprof.StopCPUProfile()
+	}
+
+	var err error
+
+	// Handle legacy mode flags that aren't unified yet
+	switch {
+	case *zimageFlag:
+		m := &zimage.Model{}
+		if loadErr := m.Load(*modelPath); loadErr != nil {
+			log.Fatal(loadErr)
+		}
+		var img *mlx.Array
+		img, err = m.GenerateFromConfig(context.Background(), &zimage.GenerateConfig{
+			Prompt:            *prompt,
+			NegativePrompt:    *negativePrompt,
+			CFGScale:          float32(*cfgScale),
+			Width:             int32(*width),
+			Height:            int32(*height),
+			Steps:             *steps,
+			Seed:              *seed,
+			CapturePath:       *gpuCapture,
+			TeaCache:          *teaCache,
+			TeaCacheThreshold: float32(*teaCacheThreshold),
+			FusedQKV:          *fusedQKV,
+		})
+		if err == nil {
+			err = saveImageArray(img, *out)
+		}
+	case *flux2Flag:
+		m := &flux2.Model{}
+		if loadErr := m.Load(*modelPath); loadErr != nil {
+			log.Fatal(loadErr)
+		}
+		// Load input images with EXIF orientation correction
+		var loadedImages []image.Image
+		for _, path := range inputImages {
+			img, loadErr := loadImageWithEXIF(path)
+			if loadErr != nil {
+				log.Fatalf("Failed to load image %s: %v", path, loadErr)
+			}
+			loadedImages = append(loadedImages, img)
+		}
+		// When input images provided and user didn't override dimensions, use 0 to match input
+		fluxWidth := int32(*width)
+		fluxHeight := int32(*height)
+		if len(loadedImages) > 0 && *width == 0 && *height == 0 {
+			// Both unset, will auto-detect from input
+		} else if len(loadedImages) > 0 && *width == 0 {
+			fluxWidth = 0 // Compute from height + aspect ratio
+		} else if len(loadedImages) > 0 && *height == 0 {
+			fluxHeight = 0 // Compute from width + aspect ratio
+		}
+		var img *mlx.Array
+		img, err = m.GenerateFromConfig(context.Background(), &flux2.GenerateConfig{
+			Prompt:        *prompt,
+			Width:         fluxWidth,
+			Height:        fluxHeight,
+			Steps:         *steps,
+			GuidanceScale: float32(*cfgScale),
+			Seed:          *seed,
+			CapturePath:   *gpuCapture,
+			InputImages:   loadedImages,
+		})
+		if err == nil {
+			err = saveImageArray(img, *out)
+		}
+	case *listTensors:
+		err = listModelTensors(*modelPath)
+	default:
+		// llm path
+		m, err := load(*modelPath)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		// Load image if provided and model supports it.
+		var image *mlx.Array
+		if *imagePath != "" {
+			if mm, ok := m.(interface{ ImageSize() int32 }); ok {
+				image, err = imagegen.ProcessImage(*imagePath, mm.ImageSize())
+				if err != nil {
+					log.Fatal("load image:", err)
+				}
+			} else {
+				log.Fatal("model does not support image input")
+			}
+		}
+
+		err = generate(context.Background(), m, input{
+			Prompt:       *prompt,
+			Image:        image,
+			MaxTokens:    *maxTokens,
+			Temperature:  float32(*temperature),
+			TopP:         float32(*topP),
+			TopK:         *topK,
+			WiredLimitGB: *wiredLimitGB,
+		}, func(out output) {
+			if out.Text != "" {
+				fmt.Print(out.Text)
+			}
+			if out.Done {
+				fmt.Printf("\n\n[prefill: %.1f tok/s, gen: %.1f tok/s]\n", out.PrefillTokSec, out.GenTokSec)
+			}
+		})
+	}
+
+	if err != nil {
+		log.Fatal(err)
+	}
+}
+
+func listModelTensors(modelPath string) error {
+	weights, err := safetensors.LoadModelWeights(modelPath)
+	if err != nil {
+		return err
+	}
+	for _, name := range weights.ListTensors() {
+		info, _ := weights.GetTensorInfo(name)
+		fmt.Printf("%s: %v (%s)\n", name, info.Shape, info.Dtype)
+	}
+	return nil
+}
+
+// loadModel builds and evaluates a model using the common load pattern.
+// Release safetensors BEFORE eval - lazy arrays have captured their data,
+// and this reduces peak memory by ~6GB (matches mlx-lm behavior).
+func loadModel[T Model](build func() T, cleanup func()) T {
+	m := build()
+	weights := mlx.Collect(m)
+	cleanup()
+	mlx.Eval(weights...)
+	return m
+}
+
+func load(modelPath string) (Model, error) {
+	kind, err := detectModelKind(modelPath)
+	if err != nil {
+		return nil, fmt.Errorf("detect model kind: %w", err)
+	}
+
+	switch kind {
+	default:
+		return nil, fmt.Errorf("model type %q is not supported by x/imagegen/cmd/engine", kind)
+	}
+}
+
+func detectModelKind(modelPath string) (string, error) {
+	indexPath := filepath.Join(modelPath, "model_index.json")
+	if _, err := os.Stat(indexPath); err == nil {
+		data, err := os.ReadFile(indexPath)
+		if err != nil {
+			return "zimage", nil
+		}
+		var index struct {
+			ClassName string `json:"_class_name"`
+		}
+		if err := json.Unmarshal(data, &index); err == nil {
+			switch index.ClassName {
+			case "FluxPipeline", "ZImagePipeline":
+				return "zimage", nil
+			case "Flux2KleinPipeline":
+				return "flux2", nil
+			}
+		}
+		return "zimage", nil
+	}
+
+	configPath := filepath.Join(modelPath, "config.json")
+	data, err := os.ReadFile(configPath)
+	if err != nil {
+		return "", fmt.Errorf("no config.json or model_index.json found: %w", err)
+	}
+
+	var cfg struct {
+		ModelType string `json:"model_type"`
+	}
+	if err := json.Unmarshal(data, &cfg); err != nil {
+		return "", fmt.Errorf("parse config.json: %w", err)
+	}
+
+	return cfg.ModelType, nil
+}
+
+// loadImageWithEXIF loads an image from a file path with EXIF orientation correction.
+func loadImageWithEXIF(path string) (image.Image, error) {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return nil, fmt.Errorf("read file: %w", err)
+	}
+	return imagegen.DecodeImage(data)
+}

x/imagegen/cmd/engine/sample.go

@@ -0,0 +1,47 @@
+package main
+
+import "github.com/ollama/ollama/x/imagegen/mlx"
+
+// sampleTopK samples from top-k logits using global random state
+func sampleTopK(scaledLogits *mlx.Array, k int) *mlx.Array {
+	neg := mlx.Neg(scaledLogits)
+	indices := mlx.Argpartition(neg, k-1, -1)
+	topKIdx := mlx.Slice(indices, []int32{0}, []int32{int32(k)})
+	values := mlx.TakeAlongAxis(scaledLogits, topKIdx, -1)
+	sampled := mlx.RandomCategorical(values, -1, 1)
+	return mlx.Take(topKIdx, sampled, -1)
+}
+
+// sampleTopP samples using nucleus sampling with global random state
+func sampleTopP(scaledLogits *mlx.Array, p float32, vocabSize int32) *mlx.Array {
+	sorted := mlx.Argsort(mlx.Neg(scaledLogits), -1)
+	sortedLogits := mlx.TakeAlongAxis(scaledLogits, sorted, -1)
+	probs := mlx.Softmax(sortedLogits, -1)
+	cumProbs := mlx.Cumsum(probs, -1)
+	mask := mlx.LessScalar(cumProbs, p)
+	negInf := mlx.FullDtype(float32(-1e9), scaledLogits.Dtype(), vocabSize)
+	masked := mlx.Where(mask, sortedLogits, negInf)
+	sampled := mlx.RandomCategorical(masked, -1, 1)
+	return mlx.Take(sorted, sampled, -1)
+}
+
+// sample samples from logits at the last position
+func sample(logits *mlx.Array, temp float32, topK int, topP float32, vocab int32) *mlx.Array {
+	// Get last position logits: [1, L, vocab] -> [vocab]
+	shape := logits.Shape()
+	seqLen := shape[1]
+	lastLogits := mlx.Slice(logits, []int32{0, seqLen - 1, 0}, []int32{1, seqLen, vocab})
+	lastLogits = mlx.Reshape(lastLogits, vocab)
+
+	if temp == 0 {
+		return mlx.Argmax(lastLogits, -1, false)
+	}
+	scaled := mlx.DivScalar(lastLogits, temp)
+	if topK > 0 && topK < int(vocab) {
+		return sampleTopK(scaled, topK)
+	}
+	if topP > 0 && topP < 1.0 {
+		return sampleTopP(scaled, topP, vocab)
+	}
+	return mlx.RandomCategorical(scaled, -1, 1)
+}

x/imagegen/docs/blob-format.md

@@ -0,0 +1,158 @@
+# Tensor Blob Format
+
+Ollama stores model tensors as individual blobs in the safetensors format. Each blob contains a logical tensor (or a combined quantized tensor with its scale/bias components), or a group of logical tensors (e.g. shared experts for a given layer along with the scale/bias components for that tensor).
+
+## Safetensors File Format
+
+Every blob follows the [safetensors](https://github.com/huggingface/safetensors) layout:
+
+```
+[8 bytes: header_size (uint64 LE)] [header_size bytes: JSON header] [tensor data region]
+```
+
+The JSON header maps tensor names to their dtype, shape, and byte offsets within the data region. A special `__metadata__` key holds string-to-string metadata.
+
+## Unquantized Blobs
+
+An unquantized blob stores a single tensor keyed by its name:
+
+```json
+{
+  "model.layers.0.self_attn.q_proj.weight": {
+    "dtype": "BF16",
+    "shape": [2560, 2560],
+    "data_offsets": [0, 13107200]
+  }
+}
+```
+
+The tensor key is the full tensor name. Dtype is typically `BF16` or `F32`.
+
+## Quantized Blobs (Combined Format)
+
+A quantized blob stores the packed weight, scaling factors, and optional zero-point biases in a single file. Tensor keys use the tensor name, with `.scale` and `.bias` suffixes for the auxiliary tensors:
+
+```json
+{
+  "__metadata__": {
+    "quant_type": "int4",
+    "group_size": "32"
+  },
+  "model.layers.0.mlp.up_proj.weight": {
+    "dtype": "U32",
+    "shape": [2560, 320],
+    "data_offsets": [0, 3276800]
+  },
+  "model.layers.0.mlp.up_proj.weight.scale": {
+    "dtype": "BF16",
+    "shape": [2560, 80],
+    "data_offsets": [3276800, 3686400]
+  },
+  "model.layers.0.mlp.up_proj.weight.bias": {
+    "dtype": "BF16",
+    "shape": [2560, 80],
+    "data_offsets": [3686400, 4096000]
+  }
+}
+```
+
+### Metadata Fields
+
+| Field | Description |
+|---|---|
+| `quant_type` | Quantization type: `int4`, `int8`, `nvfp4`, or `mxfp8` |
+| `group_size` | Number of elements per quantization group (e.g., `32`, `64`) |
+
+### Tensor Keys
+
+| Key | Description |
+|---|---|
+| `{name}` | Packed quantized weights (dtype `U32`) |
+| `{name}.scale` | Per-group scaling factors |
+| `{name}.bias` | Per-group zero-point offsets (affine modes only) |
+
+## Quantization Types
+
+| Type | Bits | Group Size | Mode | Has Bias |
+|---|---|---|---|---|
+| `int4` | 4 | 32 | affine | yes |
+| `int8` | 8 | 64 | affine | yes |
+| `nvfp4` | 4 | 16 | nvfp4 | no |
+| `mxfp8` | 8 | 32 | mxfp8 | no |
+
+**Affine modes** (`int4`, `int8`) use `scale + bias` for dequantization. The bias tensor provides the zero-point offset.
+
+**Non-affine modes** (`nvfp4`, `mxfp8`) use only `scale` with specialized E4M3 scale formats.
+
+### Packed Weight Shape
+
+Quantized weights are packed into `uint32` values:
+- **4-bit** (int4, nvfp4): 8 values per uint32, so `packed_cols = original_cols / 8`
+- **8-bit** (int8, mxfp8): 4 values per uint32, so `packed_cols = original_cols / 4`
+
+Scale shape: `[rows, original_cols / group_size]`
+
+## Manifest References
+
+Blobs are referenced from the model manifest as layers:
+
+```json
+{
+  "mediaType": "application/vnd.ollama.image.tensor",
+  "digest": "sha256:abc123...",
+  "size": 4096150,
+  "name": "model.layers.0.mlp.up_proj.weight"
+}
+```
+
+Each tensor (quantized or not) is one layer in the manifest. The layer name matches the tensor key in the blob header.
+
+## Packed Blobs (Expert Groups)
+
+For MoE (Mixture of Experts) models, expert tensors from the same layer are packed into a single blob to reduce blob count and improve loading efficiency. A packed blob is a standard safetensors file containing multiple tensor entries:
+
+```json
+{
+  "model.layers.1.mlp.experts.0.down_proj.weight": {
+    "dtype": "U32",
+    "shape": [2560, 640],
+    "data_offsets": [0, 6553600]
+  },
+  "model.layers.1.mlp.experts.0.down_proj.weight.scale": {
+    "dtype": "BF16",
+    "shape": [2560, 40],
+    "data_offsets": [6553600, 6963200]
+  },
+  "model.layers.1.mlp.experts.0.gate_proj.weight": {
+    "dtype": "U32",
+    "shape": [10240, 320],
+    "data_offsets": [6963200, 20070400]
+  },
+  "model.layers.1.mlp.experts.0.gate_proj.weight.scale": { "..." : "..." }
+}
+```
+
+### Grouping Rules
+
+- `model.layers.{L}.mlp.experts.*` tensors are packed into one blob per layer
+- `model.layers.{L}.mlp.shared_experts.*` tensors are packed into one blob per layer
+- All other tensors remain as individual blobs
+
+### Manifest Representation
+
+One manifest layer per packed group, using the group prefix as the layer name:
+
+```json
+{
+  "mediaType": "application/vnd.ollama.image.tensor",
+  "digest": "sha256:...",
+  "size": 123456789,
+  "name": "model.layers.1.mlp.experts"
+}
+```
+
+## Loading
+
+At load time, `mlx_load_safetensors` opens each blob via mmap for zero-copy access. For combined quantized blobs, the loader extracts `{name}`, `{name}.scale`, and `{name}.bias` tensors and caches them as `name`, `name + "_scale"`, and `name + "_qbias"` respectively, maintaining compatibility with the weight loading interface.
+
+For packed blobs, if the manifest layer name (group prefix) is not found as a tensor key, the loader parses the blob header to discover all tensor names and loads each individually.

x/imagegen/image.go

@@ -0,0 +1,293 @@
+package imagegen
+
+import (
+	"bytes"
+	"encoding/base64"
+	"fmt"
+	"image"
+	"image/color"
+	"image/draw"
+	_ "image/jpeg"
+	"image/png"
+	"os"
+	"path/filepath"
+
+	"github.com/ollama/ollama/x/imagegen/mlx"
+)
+
+// SaveImage saves an MLX array as a PNG image file.
+// Expected format: [B, C, H, W] with values in [0, 1] range and C=3 (RGB).
+func SaveImage(arr *mlx.Array, path string) error {
+	img, err := ArrayToImage(arr)
+	if err != nil {
+		return err
+	}
+
+	if filepath.Ext(path) != ".png" {
+		path = path + ".png"
+	}
+
+	f, err := os.Create(path)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	return png.Encode(f, img)
+}
+
+// EncodeImageBase64 encodes an MLX array as a base64-encoded PNG.
+// Expected format: [B, C, H, W] with values in [0, 1] range and C=3 (RGB).
+func EncodeImageBase64(arr *mlx.Array) (string, error) {
+	img, err := ArrayToImage(arr)
+	if err != nil {
+		return "", err
+	}
+
+	var buf bytes.Buffer
+	if err := png.Encode(&buf, img); err != nil {
+		return "", err
+	}
+
+	return base64.StdEncoding.EncodeToString(buf.Bytes()), nil
+}
+
+// ArrayToImage converts an MLX array to a Go image.RGBA.
+// Expected format: [B, C, H, W] with values in [0, 1] range and C=3 (RGB).
+func ArrayToImage(arr *mlx.Array) (*image.RGBA, error) {
+	shape := arr.Shape()
+	if len(shape) != 4 {
+		return nil, fmt.Errorf("expected 4D array [B, C, H, W], got %v", shape)
+	}
+
+	// Transform to [H, W, C] for image conversion
+	// Free intermediate arrays to avoid memory leak
+	squeezed := mlx.Squeeze(arr, 0)
+	transposed := mlx.Transpose(squeezed, 1, 2, 0)
+	squeezed.Free()
+	img := mlx.Contiguous(transposed)
+	transposed.Free()
+	mlx.Eval(img)
+
+	imgShape := img.Shape()
+	H := int(imgShape[0])
+	W := int(imgShape[1])
+	C := int(imgShape[2])
+
+	if C != 3 {
+		img.Free()
+		return nil, fmt.Errorf("expected 3 channels (RGB), got %d", C)
+	}
+
+	// Copy to CPU and free GPU memory
+	data := img.Data()
+	img.Free()
+
+	// Write directly to Pix slice (faster than SetRGBA)
+	goImg := image.NewRGBA(image.Rect(0, 0, W, H))
+	pix := goImg.Pix
+	for y := 0; y < H; y++ {
+		for x := 0; x < W; x++ {
+			srcIdx := (y*W + x) * C
+			dstIdx := (y*W + x) * 4
+			pix[dstIdx+0] = uint8(clampF(data[srcIdx+0]*255+0.5, 0, 255))
+			pix[dstIdx+1] = uint8(clampF(data[srcIdx+1]*255+0.5, 0, 255))
+			pix[dstIdx+2] = uint8(clampF(data[srcIdx+2]*255+0.5, 0, 255))
+			pix[dstIdx+3] = 255
+		}
+	}
+
+	return goImg, nil
+}
+
+func clampF(v, min, max float32) float32 {
+	if v < min {
+		return min
+	}
+	if v > max {
+		return max
+	}
+	return v
+}
+
+// DecodeImage decodes image bytes with EXIF orientation applied.
+// Transparent images are composited onto a white background.
+func DecodeImage(data []byte) (image.Image, error) {
+	orientation := readJPEGOrientation(data)
+
+	img, _, err := image.Decode(bytes.NewReader(data))
+	if err != nil {
+		return nil, err
+	}
+
+	img = flattenAlpha(img)
+	return applyOrientation(img, orientation), nil
+}
+
+// flattenAlpha composites an image onto a white background,
+// removing any transparency. This is needed because image
+// generation models don't handle alpha channels well.
+func flattenAlpha(img image.Image) image.Image {
+	if _, ok := img.(*image.RGBA); !ok {
+		if _, ok := img.(*image.NRGBA); !ok {
+			// No alpha channel, return as-is
+			return img
+		}
+	}
+
+	bounds := img.Bounds()
+	dst := image.NewRGBA(bounds)
+
+	// Fill with white background
+	draw.Draw(dst, bounds, &image.Uniform{color.White}, image.Point{}, draw.Src)
+
+	// Composite the image on top
+	draw.Draw(dst, bounds, img, bounds.Min, draw.Over)
+
+	return dst
+}
+
+// readJPEGOrientation extracts EXIF orientation from JPEG bytes.
+// Returns 1 (normal) for non-JPEG or if orientation not found.
+func readJPEGOrientation(data []byte) int {
+	if len(data) < 2 || data[0] != 0xFF || data[1] != 0xD8 {
+		return 1 // Not JPEG
+	}
+
+	r := bytes.NewReader(data[2:])
+	for {
+		var marker [2]byte
+		if _, err := r.Read(marker[:]); err != nil || marker[0] != 0xFF {
+			return 1
+		}
+
+		if marker[1] == 0xE1 { // APP1 (EXIF)
+			var lenBytes [2]byte
+			if _, err := r.Read(lenBytes[:]); err != nil {
+				return 1
+			}
+			segLen := int(uint16(lenBytes[0])<<8|uint16(lenBytes[1])) - 2
+			if segLen < 14 {
+				r.Seek(int64(segLen), 1)
+				continue
+			}
+			seg := make([]byte, segLen)
+			if _, err := r.Read(seg); err != nil {
+				return 1
+			}
+			if string(seg[:4]) == "Exif" && seg[4] == 0 && seg[5] == 0 {
+				return parseTIFFOrientation(seg[6:])
+			}
+			continue
+		}
+
+		if marker[1] == 0xD9 || marker[1] == 0xDA {
+			return 1 // EOI or SOS
+		}
+		if marker[1] >= 0xD0 && marker[1] <= 0xD7 {
+			continue // RST markers
+		}
+
+		var lenBytes [2]byte
+		if _, err := r.Read(lenBytes[:]); err != nil {
+			return 1
+		}
+		segLen := int(uint16(lenBytes[0])<<8|uint16(lenBytes[1])) - 2
+		if segLen > 0 {
+			r.Seek(int64(segLen), 1)
+		}
+	}
+}
+
+func parseTIFFOrientation(tiff []byte) int {
+	if len(tiff) < 8 {
+		return 1
+	}
+
+	var big bool
+	switch string(tiff[:2]) {
+	case "MM":
+		big = true
+	case "II":
+		big = false
+	default:
+		return 1
+	}
+
+	u16 := func(b []byte) uint16 {
+		if big {
+			return uint16(b[0])<<8 | uint16(b[1])
+		}
+		return uint16(b[1])<<8 | uint16(b[0])
+	}
+	u32 := func(b []byte) uint32 {
+		if big {
+			return uint32(b[0])<<24 | uint32(b[1])<<16 | uint32(b[2])<<8 | uint32(b[3])
+		}
+		return uint32(b[3])<<24 | uint32(b[2])<<16 | uint32(b[1])<<8 | uint32(b[0])
+	}
+
+	if u16(tiff[2:4]) != 42 {
+		return 1
+	}
+
+	ifdOffset := u32(tiff[4:8])
+	if int(ifdOffset)+2 > len(tiff) {
+		return 1
+	}
+
+	numEntries := u16(tiff[ifdOffset : ifdOffset+2])
+	for i := range int(numEntries) {
+		offset := ifdOffset + 2 + uint32(i)*12
+		if int(offset)+12 > len(tiff) {
+			break
+		}
+		if u16(tiff[offset:offset+2]) == 0x0112 { // Orientation tag
+			o := int(u16(tiff[offset+8 : offset+10]))
+			if o >= 1 && o <= 8 {
+				return o
+			}
+			return 1
+		}
+	}
+	return 1
+}
+
+func applyOrientation(img image.Image, orientation int) image.Image {
+	if orientation <= 1 || orientation > 8 {
+		return img
+	}
+
+	bounds := img.Bounds()
+	w, h := bounds.Dx(), bounds.Dy()
+
+	outW, outH := w, h
+	if orientation >= 5 {
+		outW, outH = h, w
+	}
+
+	out := image.NewRGBA(image.Rect(0, 0, outW, outH))
+	for y := range h {
+		for x := range w {
+			var dx, dy int
+			switch orientation {
+			case 2:
+				dx, dy = w-1-x, y
+			case 3:
+				dx, dy = w-1-x, h-1-y
+			case 4:
+				dx, dy = x, h-1-y
+			case 5:
+				dx, dy = y, x
+			case 6:
+				dx, dy = h-1-y, x
+			case 7:
+				dx, dy = h-1-y, w-1-x
+			case 8:
+				dx, dy = y, w-1-x
+			}
+			out.Set(dx, dy, img.At(x+bounds.Min.X, y+bounds.Min.Y))
+		}
+	}
+	return out
+}

x/imagegen/image_processor.go

@@ -0,0 +1,56 @@
+package imagegen
+
+import (
+	"fmt"
+	"image"
+	_ "image/jpeg"
+	_ "image/png"
+	"os"
+
+	"github.com/ollama/ollama/x/imagegen/mlx"
+	"golang.org/x/image/draw"
+)
+
+// ProcessImage loads and preprocesses an image for multimodal vision towers.
+// Returns [1, H, W, C] tensor in NHWC format normalized for SigLIP.
+func ProcessImage(path string, imageSize int32) (*mlx.Array, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return nil, fmt.Errorf("open image: %w", err)
+	}
+	defer f.Close()
+
+	img, _, err := image.Decode(f)
+	if err != nil {
+		return nil, fmt.Errorf("decode image: %w", err)
+	}
+
+	return ProcessImageData(img, imageSize)
+}
+
+// ProcessImageData preprocesses an image.Image for multimodal vision towers.
+func ProcessImageData(img image.Image, imageSize int32) (*mlx.Array, error) {
+	// Resize to target size using bilinear interpolation.
+	resized := image.NewRGBA(image.Rect(0, 0, int(imageSize), int(imageSize)))
+	draw.BiLinear.Scale(resized, resized.Bounds(), img, img.Bounds(), draw.Over, nil)
+
+	// Convert to float32 array [H, W, C] and normalize.
+	// SigLIP normalization: (pixel / 255.0 - 0.5) / 0.5 = pixel / 127.5 - 1.0.
+	data := make([]float32, imageSize*imageSize*3)
+	idx := 0
+	for y := int32(0); y < imageSize; y++ {
+		for x := int32(0); x < imageSize; x++ {
+			r, g, b, _ := resized.At(int(x), int(y)).RGBA()
+			// RGBA returns 16-bit values, convert to 8-bit.
+			data[idx] = float32(r>>8)/127.5 - 1.0
+			data[idx+1] = float32(g>>8)/127.5 - 1.0
+			data[idx+2] = float32(b>>8)/127.5 - 1.0
+			idx += 3
+		}
+	}
+
+	// Create MLX array [1, H, W, C] for NHWC layout.
+	arr := mlx.NewArrayFloat32(data, []int32{1, imageSize, imageSize, 3})
+	mlx.Eval(arr) // Materialize to prevent use-after-free.
+	return arr, nil
+}

x/imagegen/imagegen.go

@@ -0,0 +1,132 @@
+package imagegen
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"log/slog"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/ollama/ollama/x/imagegen/manifest"
+	"github.com/ollama/ollama/x/imagegen/mlx"
+	"github.com/ollama/ollama/x/imagegen/models/flux2"
+	"github.com/ollama/ollama/x/imagegen/models/zimage"
+)
+
+// ImageModel is the interface for image generation models.
+type ImageModel interface {
+	GenerateImage(ctx context.Context, prompt string, width, height int32, steps int, seed int64, progress func(step, total int)) (*mlx.Array, error)
+}
+
+var imageGenMu sync.Mutex
+
+// loadImageModel loads an image generation model.
+func (s *server) loadImageModel() error {
+	// Check memory requirements before loading
+	var requiredMemory uint64
+	if modelManifest, err := manifest.LoadManifest(s.modelName); err == nil {
+		requiredMemory = uint64(modelManifest.TotalTensorSize())
+	}
+	availableMemory := mlx.GetMemoryLimit()
+	if availableMemory > 0 && requiredMemory > 0 && availableMemory < requiredMemory {
+		return fmt.Errorf("insufficient memory for image generation: need %d GB, have %d GB",
+			requiredMemory/(1024*1024*1024), availableMemory/(1024*1024*1024))
+	}
+
+	// Detect model type and load appropriate model
+	modelType := DetectModelType(s.modelName)
+	slog.Info("detected image model type", "type", modelType)
+
+	var model ImageModel
+	switch modelType {
+	case "Flux2KleinPipeline":
+		m := &flux2.Model{}
+		if err := m.Load(s.modelName); err != nil {
+			return fmt.Errorf("failed to load flux2 model: %w", err)
+		}
+		model = m
+	default:
+		// Default to Z-Image for ZImagePipeline, FluxPipeline, etc.
+		m := &zimage.Model{}
+		if err := m.Load(s.modelName); err != nil {
+			return fmt.Errorf("failed to load zimage model: %w", err)
+		}
+		model = m
+	}
+
+	s.imageModel = model
+	return nil
+}
+
+// handleImageCompletion handles image generation requests.
+func (s *server) handleImageCompletion(w http.ResponseWriter, r *http.Request, req Request) {
+	// Serialize generation requests - MLX model may not handle concurrent generation
+	imageGenMu.Lock()
+	defer imageGenMu.Unlock()
+
+	// Set seed if not provided
+	if req.Seed <= 0 {
+		req.Seed = time.Now().UnixNano()
+	}
+
+	// Set up streaming response
+	w.Header().Set("Content-Type", "application/x-ndjson")
+	w.Header().Set("Transfer-Encoding", "chunked")
+	flusher, ok := w.(http.Flusher)
+	if !ok {
+		http.Error(w, "streaming not supported", http.StatusInternalServerError)
+		return
+	}
+
+	ctx := r.Context()
+	enc := json.NewEncoder(w)
+
+	// Progress callback streams step updates
+	progress := func(step, total int) {
+		resp := Response{Step: step, Total: total}
+		enc.Encode(resp)
+		w.Write([]byte("\n"))
+		flusher.Flush()
+	}
+
+	// Generate image
+	img, err := s.imageModel.GenerateImage(ctx, req.Prompt, req.Width, req.Height, req.Steps, req.Seed, progress)
+	if err != nil {
+		// Don't send error for cancellation
+		if ctx.Err() != nil {
+			return
+		}
+		resp := Response{Content: fmt.Sprintf("error: %v", err), Done: true}
+		data, _ := json.Marshal(resp)
+		w.Write(data)
+		w.Write([]byte("\n"))
+		return
+	}
+
+	// Encode image as base64 PNG
+	imageData, err := EncodeImageBase64(img)
+	if err != nil {
+		resp := Response{Content: fmt.Sprintf("error encoding: %v", err), Done: true}
+		data, _ := json.Marshal(resp)
+		w.Write(data)
+		w.Write([]byte("\n"))
+		return
+	}
+
+	// Free the generated image array and clean up MLX state
+	img.Free()
+	mlx.ClearCache()
+	mlx.MetalResetPeakMemory()
+
+	// Send final response with image data
+	resp := Response{
+		Image: imageData,
+		Done:  true,
+	}
+	data, _ := json.Marshal(resp)
+	w.Write(data)
+	w.Write([]byte("\n"))
+	flusher.Flush()
+}

x/imagegen/manifest/manifest.go

@@ -0,0 +1,307 @@
+package manifest
+
+import (
+	"encoding/binary"
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"sort"
+	"strings"
+
+	"github.com/ollama/ollama/envconfig"
+)
+
+// ManifestLayer represents a layer in the manifest.
+type ManifestLayer struct {
+	MediaType string `json:"mediaType"`
+	Digest    string `json:"digest"`
+	Size      int64  `json:"size"`
+	Name      string `json:"name,omitempty"` // Path-style name: "component/tensor" or "path/to/config.json"
+}
+
+// Manifest represents the manifest JSON structure.
+type Manifest struct {
+	SchemaVersion int             `json:"schemaVersion"`
+	MediaType     string          `json:"mediaType"`
+	Config        ManifestLayer   `json:"config"`
+	Layers        []ManifestLayer `json:"layers"`
+}
+
+// ModelManifest holds a parsed manifest with helper methods.
+type ModelManifest struct {
+	Manifest *Manifest
+	BlobDir  string
+}
+
+func DefaultBlobDir() string {
+	return filepath.Join(envconfig.Models(), "blobs")
+}
+
+// DefaultManifestDir returns the manifest storage directory.
+// Respects OLLAMA_MODELS.
+
+func DefaultManifestDir() string {
+	return filepath.Join(envconfig.Models(), "manifests")
+}
+
+// LoadManifest loads a manifest for the given model name.
+// Model name format: "modelname" or "modelname:tag" or "host/namespace/name:tag"
+func LoadManifest(modelName string) (*ModelManifest, error) {
+	manifestPath := resolveManifestPath(modelName)
+
+	data, err := os.ReadFile(manifestPath)
+	if err != nil {
+		return nil, fmt.Errorf("read manifest: %w", err)
+	}
+
+	var manifest Manifest
+	if err := json.Unmarshal(data, &manifest); err != nil {
+		return nil, fmt.Errorf("parse manifest: %w", err)
+	}
+
+	return &ModelManifest{
+		Manifest: &manifest,
+		BlobDir:  DefaultBlobDir(),
+	}, nil
+}
+
+// resolveManifestPath converts a model name to a manifest file path.
+func resolveManifestPath(modelName string) string {
+	// Parse model name into components
+	// Default: registry.ollama.ai/library/<name>/<tag>
+	host := "registry.ollama.ai"
+	namespace := "library"
+	name := modelName
+	tag := "latest"
+
+	// Handle explicit tag
+	if idx := strings.LastIndex(name, ":"); idx != -1 {
+		tag = name[idx+1:]
+		name = name[:idx]
+	}
+
+	// Handle full path like "host/namespace/name"
+	parts := strings.Split(name, "/")
+	switch len(parts) {
+	case 3:
+		host = parts[0]
+		namespace = parts[1]
+		name = parts[2]
+	case 2:
+		namespace = parts[0]
+		name = parts[1]
+	}
+
+	return filepath.Join(DefaultManifestDir(), host, namespace, name, tag)
+}
+
+// BlobPath returns the full path to a blob given its digest.
+func (m *ModelManifest) BlobPath(digest string) string {
+	// Convert "sha256:abc123" to "sha256-abc123"
+	blobName := strings.Replace(digest, ":", "-", 1)
+	return filepath.Join(m.BlobDir, blobName)
+}
+
+// GetTensorLayers returns tensor layers, optionally filtered by component.
+// If component is empty, returns all tensor layers (for LLM models).
+// If component is specified (e.g., "text_encoder", "transformer", "vae"),
+// returns only layers with that prefix.
+func (m *ModelManifest) GetTensorLayers(component string) []ManifestLayer {
+	var layers []ManifestLayer
+	for _, layer := range m.Manifest.Layers {
+		if layer.MediaType != "application/vnd.ollama.image.tensor" {
+			continue
+		}
+		if component == "" || strings.HasPrefix(layer.Name, component+"/") {
+			layers = append(layers, layer)
+		}
+	}
+	return layers
+}
+
+// GetConfigLayer returns the config layer for a given path.
+func (m *ModelManifest) GetConfigLayer(configPath string) *ManifestLayer {
+	for _, layer := range m.Manifest.Layers {
+		if layer.MediaType == "application/vnd.ollama.image.json" && layer.Name == configPath {
+			return &layer
+		}
+	}
+	return nil
+}
+
+// ReadConfig reads and returns the content of a config file.
+func (m *ModelManifest) ReadConfig(configPath string) ([]byte, error) {
+	layer := m.GetConfigLayer(configPath)
+	if layer == nil {
+		return nil, fmt.Errorf("config %q not found in manifest", configPath)
+	}
+
+	blobPath := m.BlobPath(layer.Digest)
+	return os.ReadFile(blobPath)
+}
+
+// ReadConfigJSON reads and unmarshals a config file.
+func (m *ModelManifest) ReadConfigJSON(configPath string, v any) error {
+	data, err := m.ReadConfig(configPath)
+	if err != nil {
+		return err
+	}
+	return json.Unmarshal(data, v)
+}
+
+// OpenBlob opens a blob for reading.
+func (m *ModelManifest) OpenBlob(digest string) (io.ReadCloser, error) {
+	return os.Open(m.BlobPath(digest))
+}
+
+// HasTensorLayers returns true if the manifest has any tensor layers.
+func (m *ModelManifest) HasTensorLayers() bool {
+	for _, layer := range m.Manifest.Layers {
+		if layer.MediaType == "application/vnd.ollama.image.tensor" {
+			return true
+		}
+	}
+	return false
+}
+
+// TotalTensorSize returns the total size in bytes of all tensor layers.
+func (m *ModelManifest) TotalTensorSize() int64 {
+	var total int64
+	for _, layer := range m.Manifest.Layers {
+		if layer.MediaType == "application/vnd.ollama.image.tensor" {
+			total += layer.Size
+		}
+	}
+	return total
+}
+
+// ModelInfo contains metadata about an image generation model.
+type ModelInfo struct {
+	Architecture   string
+	ParameterCount int64
+	Quantization   string
+}
+
+// GetModelInfo returns metadata about an image generation model.
+func GetModelInfo(modelName string) (*ModelInfo, error) {
+	manifest, err := LoadManifest(modelName)
+	if err != nil {
+		return nil, fmt.Errorf("failed to load manifest: %w", err)
+	}
+
+	info := &ModelInfo{}
+
+	// Read model_index.json for architecture, parameter count, and quantization
+	if data, err := manifest.ReadConfig("model_index.json"); err == nil {
+		var index struct {
+			Architecture   string `json:"architecture"`
+			ParameterCount int64  `json:"parameter_count"`
+			Quantization   string `json:"quantization"`
+		}
+		if json.Unmarshal(data, &index) == nil {
+			info.Architecture = index.Architecture
+			info.ParameterCount = index.ParameterCount
+			info.Quantization = index.Quantization
+		}
+	}
+
+	// Fallback: detect quantization from first tensor blob's __metadata__
+	if info.Quantization == "" {
+		info.Quantization = detectQuantizationFromBlobs(manifest)
+	}
+	if info.Quantization == "" {
+		info.Quantization = "BF16"
+	}
+
+	// Fallback: estimate parameter count if not in config
+	if info.ParameterCount == 0 {
+		var totalSize int64
+		for _, layer := range manifest.Manifest.Layers {
+			if layer.MediaType == "application/vnd.ollama.image.tensor" {
+				totalSize += layer.Size
+			}
+		}
+		// Assume BF16 (2 bytes/param) as rough estimate
+		info.ParameterCount = totalSize / 2
+	}
+
+	return info, nil
+}
+
+// detectQuantizationFromBlobs reads __metadata__ from the first tensor blob
+// to detect quantization type.
+func detectQuantizationFromBlobs(manifest *ModelManifest) string {
+	for _, layer := range manifest.Manifest.Layers {
+		if layer.MediaType != "application/vnd.ollama.image.tensor" {
+			continue
+		}
+		data, err := readBlobHeader(manifest.BlobPath(layer.Digest))
+		if err != nil {
+			continue
+		}
+		var header map[string]json.RawMessage
+		if json.Unmarshal(data, &header) != nil {
+			continue
+		}
+		if metaRaw, ok := header["__metadata__"]; ok {
+			var meta map[string]string
+			if json.Unmarshal(metaRaw, &meta) == nil {
+				if qt, ok := meta["quant_type"]; ok && qt != "" {
+					return strings.ToUpper(qt)
+				}
+			}
+		}
+		// Only check the first tensor blob
+		break
+	}
+	return ""
+}
+
+// ParseBlobTensorNames reads a safetensors blob and returns all "main" tensor names.
+// Filters out __metadata__, .scale, and .bias entries to return only primary weight tensors.
+func ParseBlobTensorNames(path string) ([]string, error) {
+	data, err := readBlobHeader(path)
+	if err != nil {
+		return nil, err
+	}
+
+	var header map[string]json.RawMessage
+	if err := json.Unmarshal(data, &header); err != nil {
+		return nil, err
+	}
+
+	var names []string
+	for k := range header {
+		if k == "__metadata__" || strings.HasSuffix(k, ".scale") || strings.HasSuffix(k, ".bias") {
+			continue
+		}
+		names = append(names, k)
+	}
+
+	sort.Strings(names)
+	return names, nil
+}
+
+// readBlobHeader reads the JSON header bytes from a safetensors blob file.
+func readBlobHeader(path string) ([]byte, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+
+	var headerSize uint64
+	if err := binary.Read(f, binary.LittleEndian, &headerSize); err != nil {
+		return nil, err
+	}
+	if headerSize > 1024*1024 {
+		return nil, fmt.Errorf("header too large: %d", headerSize)
+	}
+	data := make([]byte, headerSize)
+	if _, err := io.ReadFull(f, data); err != nil {
+		return nil, err
+	}
+	return data, nil
+}

x/imagegen/manifest/manifest_test.go

@@ -0,0 +1,57 @@
+package manifest
+
+import (
+	"path/filepath"
+	"testing"
+)
+
+func TestTotalTensorSize(t *testing.T) {
+	m := &ModelManifest{
+		Manifest: &Manifest{
+			Layers: []ManifestLayer{
+				{MediaType: "application/vnd.ollama.image.tensor", Size: 1000},
+				{MediaType: "application/vnd.ollama.image.tensor", Size: 2000},
+				{MediaType: "application/vnd.ollama.image.json", Size: 500}, // not a tensor
+				{MediaType: "application/vnd.ollama.image.tensor", Size: 3000},
+			},
+		},
+	}
+
+	got := m.TotalTensorSize()
+	want := int64(6000)
+	if got != want {
+		t.Errorf("TotalTensorSize() = %d, want %d", got, want)
+	}
+}
+
+func TestTotalTensorSizeEmpty(t *testing.T) {
+	m := &ModelManifest{
+		Manifest: &Manifest{
+			Layers: []ManifestLayer{},
+		},
+	}
+
+	if got := m.TotalTensorSize(); got != 0 {
+		t.Errorf("TotalTensorSize() = %d, want 0", got)
+	}
+}
+
+func TestManifestAndBlobDirsRespectOLLAMAModels(t *testing.T) {
+	modelsDir := filepath.Join(t.TempDir(), "models")
+
+	// Simulate packaged/systemd environment
+	t.Setenv("OLLAMA_MODELS", modelsDir)
+	t.Setenv("HOME", "/usr/share/ollama")
+
+	// Manifest dir must respect OLLAMA_MODELS
+	wantManifest := filepath.Join(modelsDir, "manifests")
+	if got := DefaultManifestDir(); got != wantManifest {
+		t.Fatalf("DefaultManifestDir() = %q, want %q", got, wantManifest)
+	}
+
+	// Blob dir must respect OLLAMA_MODELS
+	wantBlobs := filepath.Join(modelsDir, "blobs")
+	if got := DefaultBlobDir(); got != wantBlobs {
+		t.Fatalf("DefaultBlobDir() = %q, want %q", got, wantBlobs)
+	}
+}

x/imagegen/manifest/weights.go

@@ -0,0 +1,296 @@
+package manifest
+
+import (
+	"fmt"
+	"sort"
+	"strconv"
+	"strings"
+
+	"github.com/ollama/ollama/x/imagegen/mlx"
+)
+
+// ManifestWeights provides fast weight loading from tensor blobs.
+// Uses native mmap loading with synthetic safetensors headers for zero-copy.
+type ManifestWeights struct {
+	manifest    *ModelManifest
+	component   string
+	tensors     map[string]ManifestLayer // name -> layer
+	cache       map[string]*mlx.Array    // name -> loaded array
+	nativeCache []*mlx.SafetensorsFile   // keep native handles alive
+	quantType   string                   // quantization type from blob metadata (e.g., "int4", "int8")
+	groupSize   int                      // quantization group size from blob metadata
+}
+
+// LoadWeightsFromManifest creates a weight loader from manifest storage.
+// If component is empty, loads all tensors (for LLM models).
+// If component is specified, loads only tensors for that component and strips the prefix.
+func LoadWeightsFromManifest(manifest *ModelManifest, component string) (*ManifestWeights, error) {
+	layers := manifest.GetTensorLayers(component)
+	if len(layers) == 0 {
+		if component == "" {
+			return nil, fmt.Errorf("no tensor layers found in manifest")
+		}
+		return nil, fmt.Errorf("no tensor layers found for component %q", component)
+	}
+
+	// Strip component prefix from tensor names for model loading
+	// e.g., "text_encoder/model.embed_tokens.weight" -> "model.embed_tokens.weight"
+	tensors := make(map[string]ManifestLayer, len(layers))
+	for _, layer := range layers {
+		if component == "" {
+			tensors[layer.Name] = layer
+		} else {
+			tensorName := strings.TrimPrefix(layer.Name, component+"/")
+			tensors[tensorName] = layer
+		}
+	}
+
+	return &ManifestWeights{
+		manifest:  manifest,
+		component: component,
+		tensors:   tensors,
+		cache:     make(map[string]*mlx.Array),
+	}, nil
+}
+
+// Load loads all tensor blobs using native mmap (zero-copy).
+// Blobs are stored in safetensors format for native mlx_load_safetensors mmap.
+// Combined quantized blobs contain tensors keyed by name, name+".scale", and optional name+".bias"
+// with quantization metadata. Scale and bias are stored in cache as name+"_scale"
+// and name+"_qbias" for compatibility with downstream loading code.
+// Packed blobs (e.g., for expert groups) contain multiple tensors; the manifest name
+// is a group prefix and individual tensors are loaded by their actual names from the blob.
+// If dtype is non-zero, non-quantized tensors are converted to the specified dtype.
+func (mw *ManifestWeights) Load(dtype mlx.Dtype) error {
+	// Track native handles to free after batch eval
+	nativeHandles := make([]*mlx.SafetensorsFile, 0, len(mw.tensors))
+	arrays := make([]*mlx.Array, 0, len(mw.tensors))
+
+	// Group tensors by digest to avoid loading the same blob multiple times
+	type blobEntry struct {
+		name  string
+		layer ManifestLayer
+	}
+	blobGroups := make(map[string][]blobEntry)
+	for name, layer := range mw.tensors {
+		blobGroups[layer.Digest] = append(blobGroups[layer.Digest], blobEntry{name, layer})
+	}
+
+	for digest, entries := range blobGroups {
+		path := mw.manifest.BlobPath(digest)
+
+		// Load blob as safetensors (native mmap, zero-copy)
+		sf, err := mlx.LoadSafetensorsNative(path)
+		if err != nil {
+			for _, h := range nativeHandles {
+				h.Free()
+			}
+			return fmt.Errorf("load %s: %w", entries[0].name, err)
+		}
+		nativeHandles = append(nativeHandles, sf)
+
+		// Read quantization metadata from blob
+		if qt := sf.GetMetadata("quant_type"); qt != "" && mw.quantType == "" {
+			mw.quantType = qt
+			if gs := sf.GetMetadata("group_size"); gs != "" {
+				mw.groupSize, _ = strconv.Atoi(gs)
+			}
+		}
+
+		for _, entry := range entries {
+			name := entry.name
+
+			// Try to get tensor by stripped name first, then with component prefix,
+			// then fall back to "data" for legacy blobs created by older versions
+			// that stored all tensors with the generic key "data".
+			lookupName := name
+			arr := sf.Get(lookupName)
+			if arr == nil && mw.component != "" {
+				lookupName = mw.component + "/" + name
+				arr = sf.Get(lookupName)
+			}
+			if arr == nil {
+				// Legacy blob format: tensor stored as "data"
+				lookupName = "data"
+				arr = sf.Get(lookupName)
+			}
+			if arr != nil {
+				// Single-tensor blob or tensor found by name
+				if dtype != 0 && arr.Dtype() != dtype {
+					arr = mlx.AsType(arr, dtype)
+				}
+				arr = mlx.Contiguous(arr)
+				mw.cache[name] = arr
+				arrays = append(arrays, arr)
+
+				// Check for scale tensor
+				if scale := sf.Get(lookupName + ".scale"); scale != nil {
+					scale = mlx.Contiguous(scale)
+					mw.cache[name+"_scale"] = scale
+					arrays = append(arrays, scale)
+				}
+
+				// Check for bias tensor
+				if bias := sf.Get(lookupName + ".bias"); bias != nil {
+					bias = mlx.Contiguous(bias)
+					mw.cache[name+"_qbias"] = bias
+					arrays = append(arrays, bias)
+				}
+			} else {
+				// Packed blob: manifest name is a group prefix, not a tensor name.
+				// Load all individual tensors from the blob.
+				tensorNames, err := ParseBlobTensorNames(path)
+				if err != nil {
+					for _, h := range nativeHandles {
+						h.Free()
+					}
+					return fmt.Errorf("parse packed blob for %s: %w", name, err)
+				}
+
+				for _, tensorName := range tensorNames {
+					tArr := sf.Get(tensorName)
+					if tArr == nil {
+						continue
+					}
+
+					if dtype != 0 && tArr.Dtype() != dtype {
+						tArr = mlx.AsType(tArr, dtype)
+					}
+					tArr = mlx.Contiguous(tArr)
+
+					// Strip component prefix from blob-internal names so cache keys
+					// match the stripped names used by LoadModule.
+					cacheName := tensorName
+					if mw.component != "" {
+						cacheName = strings.TrimPrefix(tensorName, mw.component+"/")
+					}
+					mw.cache[cacheName] = tArr
+					arrays = append(arrays, tArr)
+
+					// Check for scale tensor
+					if scale := sf.Get(tensorName + ".scale"); scale != nil {
+						scale = mlx.Contiguous(scale)
+						mw.cache[cacheName+"_scale"] = scale
+						arrays = append(arrays, scale)
+					}
+
+					// Check for bias tensor
+					if bias := sf.Get(tensorName + ".bias"); bias != nil {
+						bias = mlx.Contiguous(bias)
+						mw.cache[cacheName+"_qbias"] = bias
+						arrays = append(arrays, bias)
+					}
+				}
+			}
+		}
+	}
+
+	// Batch evaluate all tensors at once (much faster than one at a time)
+	mlx.Eval(arrays...)
+
+	// Now safe to free all native handles
+	for _, sf := range nativeHandles {
+		sf.Free()
+	}
+
+	return nil
+}
+
+// GetTensor returns a tensor from cache. Call Load() first.
+func (mw *ManifestWeights) GetTensor(name string) (*mlx.Array, error) {
+	if mw.cache == nil {
+		return nil, fmt.Errorf("cache not initialized: call Load() first")
+	}
+	arr, ok := mw.cache[name]
+	if !ok {
+		return nil, fmt.Errorf("tensor %q not found", name)
+	}
+	return arr, nil
+}
+
+// ListTensors returns all tensor names in sorted order.
+// Includes both manifest tensor names and scale/bias entries from combined blobs.
+func (mw *ManifestWeights) ListTensors() []string {
+	seen := make(map[string]bool, len(mw.tensors)+len(mw.cache))
+	for name := range mw.tensors {
+		seen[name] = true
+	}
+	// Also include cache entries (scale/bias from combined blobs)
+	for name := range mw.cache {
+		seen[name] = true
+	}
+	names := make([]string, 0, len(seen))
+	for name := range seen {
+		names = append(names, name)
+	}
+	sort.Strings(names)
+	return names
+}
+
+// HasTensor checks if a tensor exists in the manifest or cache.
+func (mw *ManifestWeights) HasTensor(name string) bool {
+	if _, ok := mw.tensors[name]; ok {
+		return true
+	}
+	// Also check cache for scale/bias entries from combined blobs
+	if _, ok := mw.cache[name]; ok {
+		return true
+	}
+	return false
+}
+
+// Quantization returns the model's quantization type.
+// Returns the quant_type from blob metadata (e.g., "int4", "int8", "nvfp4", "mxfp8").
+// Returns empty string if not quantized.
+// Falls back to model_index.json for image gen models.
+func (mw *ManifestWeights) Quantization() string {
+	if mw.quantType != "" {
+		return strings.ToUpper(mw.quantType)
+	}
+
+	if mw.manifest == nil {
+		return ""
+	}
+
+	// Fallback: read from model_index.json (for image gen models)
+	var index struct {
+		Quantization string `json:"quantization"`
+	}
+	if err := mw.manifest.ReadConfigJSON("model_index.json", &index); err == nil && index.Quantization != "" {
+		return index.Quantization
+	}
+
+	return ""
+}
+
+// GroupSize returns the quantization group size.
+// Returns the group_size from blob metadata.
+// Returns 0 if not specified (caller should use default based on quantization type).
+func (mw *ManifestWeights) GroupSize() int {
+	if mw.groupSize > 0 {
+		return mw.groupSize
+	}
+
+	if mw.manifest == nil {
+		return 0
+	}
+
+	// Fallback: read from model_index.json (for image gen models)
+	var index struct {
+		GroupSize int `json:"group_size"`
+	}
+	if err := mw.manifest.ReadConfigJSON("model_index.json", &index); err == nil && index.GroupSize > 0 {
+		return index.GroupSize
+	}
+
+	return 0
+}
+
+// ReleaseAll frees all native handles and clears the tensor cache.
+func (mw *ManifestWeights) ReleaseAll() {
+	for _, sf := range mw.nativeCache {
+		sf.Free()
+	}
+	mw.nativeCache = nil
+	mw.cache = nil
+}

x/imagegen/memory.go

@@ -0,0 +1,80 @@
+// Package imagegen provides experimental image generation capabilities for Ollama.
+//
+// This package is in x/ because the tensor model storage format is under development.
+// The goal is to integrate these capabilities into the main Ollama packages once
+// the format is stable.
+//
+// TODO (jmorganca): Integrate into main packages when stable:
+//   - CLI commands → cmd/
+//   - API endpoints → api/
+//   - Model creation → server/
+package imagegen
+
+import (
+	"encoding/json"
+	"fmt"
+	"runtime"
+
+	"github.com/ollama/ollama/x/imagegen/manifest"
+)
+
+// SupportedBackends lists the backends that support image generation.
+var SupportedBackends = []string{"metal", "cuda", "cpu"}
+
+// CheckPlatformSupport validates that image generation is supported on the current platform.
+// Returns nil if supported, or an error describing why it's not supported.
+func CheckPlatformSupport() error {
+	switch runtime.GOOS {
+	case "darwin":
+		// macOS: Metal is supported via MLX
+		if runtime.GOARCH != "arm64" {
+			return fmt.Errorf("image generation on macOS requires Apple Silicon (arm64), got %s", runtime.GOARCH)
+		}
+		return nil
+	case "linux", "windows":
+		// Linux/Windows: CUDA support (requires mlx or cuda build)
+		// The actual backend availability is checked at runtime
+		return nil
+	default:
+		return fmt.Errorf("image generation is not supported on %s", runtime.GOOS)
+	}
+}
+
+// ResolveModelName checks if a model name is a known image generation model.
+// Returns the normalized model name if found, empty string otherwise.
+func ResolveModelName(modelName string) string {
+	modelManifest, err := manifest.LoadManifest(modelName)
+	if err == nil && modelManifest.HasTensorLayers() {
+		return modelName
+	}
+	return ""
+}
+
+// DetectModelType reads model_index.json and returns the model type.
+// Checks both "architecture" (Ollama format) and "_class_name" (diffusers format).
+// Returns empty string if detection fails.
+func DetectModelType(modelName string) string {
+	modelManifest, err := manifest.LoadManifest(modelName)
+	if err != nil {
+		return ""
+	}
+
+	data, err := modelManifest.ReadConfig("model_index.json")
+	if err != nil {
+		return ""
+	}
+
+	var index struct {
+		Architecture string `json:"architecture"`
+		ClassName    string `json:"_class_name"`
+	}
+	if err := json.Unmarshal(data, &index); err != nil {
+		return ""
+	}
+
+	// Prefer architecture (Ollama format), fall back to _class_name (diffusers)
+	if index.Architecture != "" {
+		return index.Architecture
+	}
+	return index.ClassName
+}

x/imagegen/memory_test.go

@@ -0,0 +1,39 @@
+package imagegen
+
+import (
+	"runtime"
+	"testing"
+)
+
+func TestCheckPlatformSupport(t *testing.T) {
+	err := CheckPlatformSupport()
+
+	switch runtime.GOOS {
+	case "darwin":
+		if runtime.GOARCH == "arm64" {
+			if err != nil {
+				t.Errorf("Expected nil error on darwin/arm64, got: %v", err)
+			}
+		} else {
+			if err == nil {
+				t.Error("Expected error on darwin/non-arm64")
+			}
+		}
+	case "linux", "windows":
+		if err != nil {
+			t.Errorf("Expected nil error on %s, got: %v", runtime.GOOS, err)
+		}
+	default:
+		if err == nil {
+			t.Errorf("Expected error on unsupported platform %s", runtime.GOOS)
+		}
+	}
+}
+
+func TestResolveModelName(t *testing.T) {
+	// Non-existent model should return empty string
+	result := ResolveModelName("nonexistent-model")
+	if result != "" {
+		t.Errorf("ResolveModelName() = %q, want empty string", result)
+	}
+}

x/imagegen/mlx/CMakeLists.txt

@@ -0,0 +1,196 @@
+include(FetchContent)
+
+# Read MLX-C version from top-level file (shared with Dockerfile)
+file(READ "${CMAKE_SOURCE_DIR}/MLX_C_VERSION" MLX_C_GIT_TAG)
+string(STRIP "${MLX_C_GIT_TAG}" MLX_C_GIT_TAG)
+
+# Read MLX version from top-level file
+file(READ "${CMAKE_SOURCE_DIR}/MLX_VERSION" MLX_GIT_TAG)
+string(STRIP "${MLX_GIT_TAG}" MLX_GIT_TAG)
+
+set(MLX_C_BUILD_EXAMPLES OFF)
+
+set(MLX_BUILD_GGUF OFF)
+set(MLX_BUILD_SAFETENSORS ON)
+
+function(set_target_output_directory _target)
+    if(TARGET ${_target})
+        set_target_properties(${_target} PROPERTIES
+            RUNTIME_OUTPUT_DIRECTORY ${OLLAMA_BUILD_DIR}
+            LIBRARY_OUTPUT_DIRECTORY ${OLLAMA_BUILD_DIR}
+            ARCHIVE_OUTPUT_DIRECTORY ${OLLAMA_BUILD_DIR}
+        )
+    endif()
+endfunction()
+
+# Check for Metal support (macOS only)
+if(CMAKE_SYSTEM_NAME MATCHES "Darwin")
+    execute_process(
+      COMMAND
+        zsh "-c"
+        "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal ${XCRUN_FLAGS} -E -x metal -P - | tail -1 | tr -d '\n'"
+      OUTPUT_VARIABLE MLX_METAL_VERSION COMMAND_ERROR_IS_FATAL ANY)
+
+    if(NOT MLX_METAL_VERSION)
+        message(STATUS "`xcrun metal` error. Setting MLX_BUILD_METAL=OFF")
+        set(MLX_BUILD_METAL OFF)
+    endif()
+else()
+    # On Linux, disable Metal backend
+    message(STATUS "Non-macOS platform detected. Setting MLX_BUILD_METAL=OFF")
+    set(MLX_BUILD_METAL OFF)
+endif()
+
+# Map CMAKE_CUDA_ARCHITECTURES to MLX_CUDA_ARCHITECTURES if not explicitly set
+if(NOT MLX_CUDA_ARCHITECTURES AND CMAKE_CUDA_ARCHITECTURES)
+    set(MLX_CUDA_ARCHITECTURES ${CMAKE_CUDA_ARCHITECTURES})
+    message(STATUS "Using CMAKE_CUDA_ARCHITECTURES for MLX: ${MLX_CUDA_ARCHITECTURES}")
+endif()
+
+# Forward cuDNN environment variables to cmake variables so MLX's FindCUDNN.cmake
+# can find them via HINTS ${CUDNN_INCLUDE_PATH} / ${CUDNN_LIBRARY_PATH}.
+if(DEFINED ENV{CUDNN_INCLUDE_PATH} AND NOT CUDNN_INCLUDE_PATH)
+    set(CUDNN_INCLUDE_PATH "$ENV{CUDNN_INCLUDE_PATH}" CACHE PATH "cuDNN include path")
+    message(STATUS "Using CUDNN_INCLUDE_PATH from environment: ${CUDNN_INCLUDE_PATH}")
+endif()
+if(DEFINED ENV{CUDNN_LIBRARY_PATH} AND NOT CUDNN_LIBRARY_PATH)
+    set(CUDNN_LIBRARY_PATH "$ENV{CUDNN_LIBRARY_PATH}" CACHE PATH "cuDNN library path")
+    message(STATUS "Using CUDNN_LIBRARY_PATH from environment: ${CUDNN_LIBRARY_PATH}")
+endif()
+
+# Enable CUDA backend if CUDA architectures are specified and CUDA compiler is available
+if(MLX_CUDA_ARCHITECTURES AND CMAKE_CUDA_COMPILER)
+    set(MLX_BUILD_CUDA ON CACHE BOOL "Build CUDA backend for MLX" FORCE)
+    message(STATUS "Enabling MLX CUDA backend with architectures: ${MLX_CUDA_ARCHITECTURES}")
+elseif(MLX_CUDA_ARCHITECTURES)
+    message(WARNING "MLX_CUDA_ARCHITECTURES specified but CUDA compiler not found, CUDA backend will be disabled")
+endif()
+
+# Allow local source overrides via environment variables
+# Resolve to absolute paths so FetchContent doesn't break on relative dirs.
+if(DEFINED ENV{OLLAMA_MLX_SOURCE})
+    get_filename_component(_mlx_src "$ENV{OLLAMA_MLX_SOURCE}" ABSOLUTE BASE_DIR ${CMAKE_SOURCE_DIR})
+    set(FETCHCONTENT_SOURCE_DIR_MLX "${_mlx_src}" CACHE PATH "" FORCE)
+    message(STATUS "Using local MLX source: ${_mlx_src}")
+endif()
+if(DEFINED ENV{OLLAMA_MLX_C_SOURCE})
+    get_filename_component(_mlx_c_src "$ENV{OLLAMA_MLX_C_SOURCE}" ABSOLUTE BASE_DIR ${CMAKE_SOURCE_DIR})
+    set(FETCHCONTENT_SOURCE_DIR_MLX-C "${_mlx_c_src}" CACHE PATH "" FORCE)
+    message(STATUS "Using local MLX-C source: ${_mlx_c_src}")
+endif()
+
+# Pre-declare mlx so our pinned version takes precedence over the one
+# hardcoded in mlx-c's CMakeLists.txt (first FetchContent_Declare wins).
+FetchContent_Declare(
+    mlx
+    GIT_REPOSITORY "https://github.com/ml-explore/mlx.git"
+    GIT_TAG ${MLX_GIT_TAG}
+)
+
+FetchContent_Declare(
+    mlx-c
+    GIT_REPOSITORY "https://github.com/ml-explore/mlx-c.git"
+    GIT_TAG ${MLX_C_GIT_TAG}
+)
+FetchContent_MakeAvailable(mlx-c)
+
+# To avoid a "long tail" when building MLX with a large set of GPU
+# architectures, utilize a higher --threads (-t) setting. At high -t
+# every .cu spawns concurrent cicc instances; each cicc can consume several GB
+# compiling MLX's CUTLASS-using kernels.  This in turn can cause OOMs.
+#
+# We use a pool to cover all MLX CUDA sources. Pool size is derived from total 
+# host RAM via a per-file memory budget.
+#
+# This was calibrated with `-t 6`.  Higher -t may require overriding
+# MLX_CUDA_RAM_MB
+if(CMAKE_GENERATOR STREQUAL "Ninja")
+    file(GLOB_RECURSE _mlx_cu
+        "${mlx_SOURCE_DIR}/mlx/backend/cuda/*.cu"
+        "${mlx_BINARY_DIR}/mlx/backend/cuda/*.cu"
+    )
+    if(_mlx_cu)
+        set(MLX_CUDA_RAM_MB 22000 CACHE STRING
+            "Per-file memory budget (MB) for the cuda_compile JOB_POOL.  Override for higher -t.")
+        cmake_host_system_information(RESULT _ram_mb QUERY TOTAL_PHYSICAL_MEMORY)
+        math(EXPR _cuda_pool "${_ram_mb} / ${MLX_CUDA_RAM_MB}")
+        if(_cuda_pool LESS 2)
+            set(_cuda_pool 2)
+        endif()
+        set_property(GLOBAL APPEND PROPERTY JOB_POOLS cuda_compile=${_cuda_pool})
+        list(LENGTH _mlx_cu _cu_count)
+        # SOURCE properties default to directory-scoped, which means a plain
+        # set_property(SOURCE ...) here would NOT affect the build rules
+        # generated for the mlx target (defined in mlx_SOURCE_DIR after
+        # FetchContent).  TARGET_DIRECTORY mlx puts the property in the
+        # directory where mlx was defined, so it actually applies.
+        foreach(f ${_mlx_cu})
+            set_property(SOURCE "${f}"
+                TARGET_DIRECTORY mlx
+                PROPERTY JOB_POOL_COMPILE cuda_compile)
+        endforeach()
+        message(STATUS "MLX cuda_compile JOB_POOL: ${_cu_count} files, pool size ${_cuda_pool} (host RAM ${_ram_mb} MB / ${MLX_CUDA_RAM_MB} MB per file)")
+    else()
+        message(WARNING "MLX cuda_compile JOB_POOL: no .cu files found under mlx/backend/cuda/ - check MLX layout")
+    endif()
+endif()
+
+# Sync vendored headers with fetched version
+file(GLOB _mlx_c_hdrs "${mlx-c_SOURCE_DIR}/mlx/c/*.h")
+file(COPY ${_mlx_c_hdrs} DESTINATION "${CMAKE_SOURCE_DIR}/x/mlxrunner/mlx/include/mlx/c/")
+
+# Regenerate Go/C shim wrappers from the (possibly updated) headers.
+# Skip during cross-compilation — the generated files are arch-independent.
+if(CMAKE_SYSTEM_PROCESSOR STREQUAL CMAKE_HOST_SYSTEM_PROCESSOR OR NOT APPLE)
+    find_program(GO_EXECUTABLE go REQUIRED)
+    message(STATUS "Regenerating MLX Go wrappers")
+
+    # CGo's probe compilation is sensitive to CGO_CFLAGS/CGO_CXXFLAGS and CC.
+    # Clear them so go generate uses default compiler settings:
+    # - On Windows, CC may contain spaces (e.g., "C:/Program Files/.../cl.exe")
+    #   which breaks CGo's CC parsing.
+    # - On macOS, CGO_CFLAGS with -mmacosx-version-min breaks header search
+    #   when cmake also sets CMAKE_OSX_DEPLOYMENT_TARGET.
+    set(_SAVE_CC "$ENV{CC}")
+    set(_SAVE_CGO_CFLAGS "$ENV{CGO_CFLAGS}")
+    set(_SAVE_CGO_CXXFLAGS "$ENV{CGO_CXXFLAGS}")
+    set(ENV{CC} "")
+    set(ENV{CGO_CFLAGS} "")
+    set(ENV{CGO_CXXFLAGS} "")
+
+    execute_process(
+        COMMAND ${GO_EXECUTABLE} generate ./x/...
+        WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}
+        COMMAND_ERROR_IS_FATAL ANY
+    )
+
+    set(ENV{CC} "${_SAVE_CC}")
+    set(ENV{CGO_CFLAGS} "${_SAVE_CGO_CFLAGS}")
+    set(ENV{CGO_CXXFLAGS} "${_SAVE_CGO_CXXFLAGS}")
+else()
+    message(STATUS "Skipping MLX Go wrapper generation (cross-compiling)")
+endif()
+
+# For local dev builds, override MLX_VERSION with git describe output
+if(TARGET mlx_version AND DEFINED FETCHCONTENT_SOURCE_DIR_MLX)
+    execute_process(
+        COMMAND git describe --tags --first-parent --abbrev=7 --long --dirty --always
+        WORKING_DIRECTORY ${mlx_SOURCE_DIR}
+        OUTPUT_VARIABLE _mlx_git_version
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+        ERROR_QUIET
+        RESULT_VARIABLE _mlx_git_result
+    )
+    if(_mlx_git_result EQUAL 0 AND _mlx_git_version)
+        # Strip leading "v" prefix for consistency
+        string(REGEX REPLACE "^v" "" _mlx_git_version "${_mlx_git_version}")
+        get_target_property(_mlx_defs mlx_version COMPILE_DEFINITIONS)
+        list(FILTER _mlx_defs EXCLUDE REGEX "^MLX_VERSION=")
+        set_target_properties(mlx_version PROPERTIES COMPILE_DEFINITIONS "${_mlx_defs}")
+        target_compile_definitions(mlx_version PRIVATE "MLX_VERSION=\"${_mlx_git_version}\"")
+        message(STATUS "MLX version (local dev): ${_mlx_git_version}")
+    endif()
+endif()
+
+set_target_output_directory(mlx)
+set_target_output_directory(mlxc)

x/imagegen/mlx/README.md

@@ -0,0 +1,46 @@
+# MLX Memory Management
+
+| This package will get consolidated with `x/ml/backend/mlx` in the future.
+
+## Automatic Tracking
+
+All arrays are automatically tracked when created. On `Eval()`, non-kept arrays are freed.
+
+### API
+
+```go
+result := mlx.Matmul(x, w) // arrays automatically tracked
+mlx.Eval(result)           // free non-kept, eval result (auto-kept)
+```
+
+### Key Functions
+
+- `mlx.Eval(outputs...)` - free non-kept arrays, then evaluate (outputs auto-kept)
+- `mlx.AsyncEval(outputs...)` - async version of Eval (outputs auto-kept)
+- `mlx.Keep(arrays...)` - mark arrays to survive cleanup (for weights, caches)
+- `array.Free()` - mark array for cleanup on next Eval
+
+### Loop Pattern
+
+```go
+for step := 0; step < maxTokens; step++ {
+    logits := model.Forward(token, caches)
+    oldToken := token
+    token = sample(logits)
+
+    // Keep cache state across iterations
+    for _, c := range caches {
+        mlx.Keep(c.State()...)
+    }
+
+    oldToken.Free()       // mark for cleanup
+    mlx.AsyncEval(token)  // frees old, evals new
+}
+```
+
+### Notes
+
+- `Eval()` and `AsyncEval()` auto-keep their outputs
+- `Free()` marks for cleanup - actual free happens during next Eval
+- Use `Keep()` for weights and cache state that must survive multiple Eval cycles
+- Arrays created inside compiled closures are managed by MLX, not tracked

x/imagegen/mlx/compile.go

@@ -0,0 +1,171 @@
+package mlx
+
+/*
+#include "mlx.h"
+#include <stdlib.h>
+
+// Forward declaration for Go callback
+extern int goClosureCallback(mlx_vector_array* res, mlx_vector_array input, void* payload);
+
+// Destructor for payload (Go handle)
+extern void goClosureDestructor(void* payload);
+*/
+import "C"
+import (
+	"runtime/cgo"
+	"sync"
+	"unsafe"
+)
+
+// inClosureCallback is set to true during closure callback execution.
+var inClosureCallback bool
+var closureCallbackMu sync.Mutex
+
+// InClosureCallback returns true if we're currently executing inside a closure callback.
+func InClosureCallback() bool {
+	closureCallbackMu.Lock()
+	defer closureCallbackMu.Unlock()
+	return inClosureCallback
+}
+
+// CompiledFunc is a compiled MLX function that can be called efficiently.
+// All intermediate arrays during execution stay inside MLX - only inputs
+// and outputs cross the Go boundary.
+type CompiledFunc struct {
+	closure  C.mlx_closure
+	compiled C.mlx_closure
+}
+
+// ClosureFunc is the signature for functions that can be compiled.
+// It takes a slice of input arrays and returns a slice of output arrays.
+type ClosureFunc func(inputs []*Array) []*Array
+
+// Compile compiles a Go function into an optimized MLX closure.
+// The function is traced once during compilation, then subsequent calls
+// run the optimized graph without creating Go intermediate arrays.
+//
+// Example:
+//
+//	compiled := mlx.Compile(func(inputs []*mlx.Array) []*mlx.Array {
+//	    a, b := inputs[0], inputs[1]
+//	    c := mlx.Add(a, b)
+//	    d := mlx.Mul(c, c)
+//	    return []*mlx.Array{d}
+//	})
+//	defer compiled.Free()
+//
+//	result := compiled.Call(x, y)[0]
+func Compile(fn ClosureFunc) *CompiledFunc {
+	return CompileShapeless(fn, false)
+}
+
+// CompileShapeless compiles with optional shapeless mode.
+// If shapeless=true, the function works for any input shape after tracing.
+func CompileShapeless(fn ClosureFunc, shapeless bool) *CompiledFunc {
+	// Create a cgo.Handle to prevent the Go function from being GC'd
+	handle := cgo.NewHandle(fn)
+
+	// Create the closure from the Go callback
+	closure := C.mlx_closure_new_func_payload(
+		(*[0]byte)(C.goClosureCallback),
+		unsafe.Pointer(handle),
+		(*[0]byte)(C.goClosureDestructor),
+	)
+
+	// Compile the closure
+	compiled := C.mlx_closure_new()
+	C.mlx_compile(&compiled, closure, C.bool(shapeless))
+
+	return &CompiledFunc{
+		closure:  closure,
+		compiled: compiled,
+	}
+}
+
+// Call invokes the compiled function with the given inputs.
+func (cf *CompiledFunc) Call(inputs ...*Array) []*Array {
+	// Pack inputs into vector
+	inputVec := C.mlx_vector_array_new()
+	for _, arr := range inputs {
+		C.mlx_vector_array_append_value(inputVec, arr.c)
+	}
+
+	// Apply compiled closure
+	outputVec := C.mlx_vector_array_new()
+	C.mlx_closure_apply(&outputVec, cf.compiled, inputVec)
+	C.mlx_vector_array_free(inputVec)
+
+	// Unpack outputs
+	numOutputs := int(C.mlx_vector_array_size(outputVec))
+	outputs := make([]*Array, numOutputs)
+	for i := 0; i < numOutputs; i++ {
+		var arr C.mlx_array
+		C.mlx_vector_array_get(&arr, outputVec, C.size_t(i))
+		outputs[i] = newArray(arr)
+	}
+	C.mlx_vector_array_free(outputVec)
+
+	return outputs
+}
+
+// CallEval invokes the compiled function and evaluates the results.
+func (cf *CompiledFunc) CallEval(inputs ...*Array) []*Array {
+	outputs := cf.Call(inputs...)
+	Eval(outputs...)
+	return outputs
+}
+
+// Free releases the compiled function resources.
+func (cf *CompiledFunc) Free() {
+	C.mlx_closure_free(cf.compiled)
+	C.mlx_closure_free(cf.closure)
+}
+
+// borrowArray wraps a C array WITHOUT setting up GC cleanup.
+// Use this for arrays we don't own (e.g., borrowed references in callbacks).
+func borrowArray(array C.mlx_array) *Array {
+	return &Array{c: array}
+}
+
+//export goClosureCallback
+func goClosureCallback(res *C.mlx_vector_array, input C.mlx_vector_array, payload unsafe.Pointer) C.int {
+	// Set flag to disable AddCleanup during callback
+	closureCallbackMu.Lock()
+	inClosureCallback = true
+	closureCallbackMu.Unlock()
+	defer func() {
+		closureCallbackMu.Lock()
+		inClosureCallback = false
+		closureCallbackMu.Unlock()
+	}()
+
+	// Recover the Go function from the handle
+	handle := cgo.Handle(payload)
+	fn := handle.Value().(ClosureFunc)
+
+	// Convert input vector to Go slice - use borrowArray since MLX owns these
+	numInputs := int(C.mlx_vector_array_size(input))
+	inputs := make([]*Array, numInputs)
+	for i := 0; i < numInputs; i++ {
+		var arr C.mlx_array
+		C.mlx_vector_array_get(&arr, input, C.size_t(i))
+		inputs[i] = borrowArray(arr) // Don't set up cleanup - MLX owns these
+	}
+
+	// Call the Go function
+	outputs := fn(inputs)
+
+	// Build output vector
+	*res = C.mlx_vector_array_new()
+	for _, arr := range outputs {
+		C.mlx_vector_array_append_value(*res, arr.c)
+	}
+
+	return 0
+}
+
+//export goClosureDestructor
+func goClosureDestructor(payload unsafe.Pointer) {
+	handle := cgo.Handle(payload)
+	handle.Delete()
+}

x/imagegen/mlx/doc.go

@@ -0,0 +1,4 @@
+// Package mlx provides Go bindings for the MLX-C library with dynamic loading support.
+//
+//go:generate go run generate_wrappers.go ../../mlxrunner/mlx/include/mlx/c mlx.h mlx.c
+package mlx

x/imagegen/mlx/generate_wrappers.go

@@ -0,0 +1,451 @@
+//go:build ignore
+
+// This tool generates MLX-C dynamic loading wrappers.
+// Usage: go run generate_wrappers.go <mlx-c-include-dir> <output-header> [output-impl]
+package main
+
+import (
+	"bytes"
+	"flag"
+	"fmt"
+	"io/fs"
+	"os"
+	"path/filepath"
+	"regexp"
+	"strings"
+)
+
+type Function struct {
+	Name            string
+	ReturnType      string
+	Params          string
+	ParamNames      []string
+	NeedsARM64Guard bool
+}
+
+func findHeaders(directory string) ([]string, error) {
+	var headers []string
+	err := filepath.WalkDir(directory, func(path string, d fs.DirEntry, err error) error {
+		if err != nil {
+			return err
+		}
+		// Private headers contain C++ implementation helpers and are not part of
+		// the C API surface; parsing them can produce invalid wrapper signatures.
+		if d.IsDir() && d.Name() == "private" {
+			return fs.SkipDir
+		}
+		if !d.IsDir() && strings.HasSuffix(path, ".h") {
+			headers = append(headers, path)
+		}
+		return nil
+	})
+	return headers, err
+}
+
+func cleanContent(content string) string {
+	// Remove single-line comments
+	re := regexp.MustCompile(`//.*?\n`)
+	content = re.ReplaceAllString(content, "\n")
+
+	// Remove multi-line comments
+	re = regexp.MustCompile(`/\*.*?\*/`)
+	content = re.ReplaceAllString(content, "")
+
+	// Remove preprocessor directives (lines starting with #) - use multiline mode
+	re = regexp.MustCompile(`(?m)^\s*#.*?$`)
+	content = re.ReplaceAllString(content, "")
+
+	// Remove extern "C" { and } blocks more conservatively
+	// Only remove the extern "C" { line, not the content inside
+	re = regexp.MustCompile(`extern\s+"C"\s*\{\s*?\n`)
+	content = re.ReplaceAllString(content, "\n")
+	// Remove standalone closing braces that are not part of function declarations
+	re = regexp.MustCompile(`\n\s*\}\s*\n`)
+	content = re.ReplaceAllString(content, "\n")
+
+	// Collapse whitespace and newlines
+	re = regexp.MustCompile(`\s+`)
+	content = re.ReplaceAllString(content, " ")
+
+	return content
+}
+
+func extractParamNames(params string) []string {
+	if params == "" || strings.TrimSpace(params) == "void" {
+		return []string{}
+	}
+
+	var names []string
+
+	// Split by comma, but respect parentheses (for function pointers)
+	parts := splitParams(params)
+
+	// Remove array brackets
+	arrayBrackets := regexp.MustCompile(`\[.*?\]`)
+
+	// Function pointer pattern
+	funcPtrPattern := regexp.MustCompile(`\(\s*\*\s*(\w+)\s*\)`)
+
+	// Type keywords to skip
+	typeKeywords := map[string]bool{
+		"const":     true,
+		"struct":    true,
+		"unsigned":  true,
+		"signed":    true,
+		"long":      true,
+		"short":     true,
+		"int":       true,
+		"char":      true,
+		"float":     true,
+		"double":    true,
+		"void":      true,
+		"size_t":    true,
+		"uint8_t":   true,
+		"uint16_t":  true,
+		"uint32_t":  true,
+		"uint64_t":  true,
+		"int8_t":    true,
+		"int16_t":   true,
+		"int32_t":   true,
+		"int64_t":   true,
+		"intptr_t":  true,
+		"uintptr_t": true,
+	}
+
+	for _, part := range parts {
+		if part == "" {
+			continue
+		}
+
+		// Remove array brackets
+		part = arrayBrackets.ReplaceAllString(part, "")
+
+		// For function pointers like "void (*callback)(int)"
+		if matches := funcPtrPattern.FindStringSubmatch(part); len(matches) > 1 {
+			names = append(names, matches[1])
+			continue
+		}
+
+		// Regular parameter: last identifier
+		tokens := regexp.MustCompile(`\w+`).FindAllString(part, -1)
+		if len(tokens) > 0 {
+			// The last token is usually the parameter name
+			// Skip type keywords
+			for i := len(tokens) - 1; i >= 0; i-- {
+				if !typeKeywords[tokens[i]] {
+					names = append(names, tokens[i])
+					break
+				}
+			}
+		}
+	}
+
+	return names
+}
+
+func splitParams(params string) []string {
+	var parts []string
+	var current bytes.Buffer
+	depth := 0
+
+	for _, char := range params + "," {
+		switch char {
+		case '(':
+			depth++
+			current.WriteRune(char)
+		case ')':
+			depth--
+			current.WriteRune(char)
+		case ',':
+			if depth == 0 {
+				parts = append(parts, strings.TrimSpace(current.String()))
+				current.Reset()
+			} else {
+				current.WriteRune(char)
+			}
+		default:
+			current.WriteRune(char)
+		}
+	}
+
+	return parts
+}
+
+func parseFunctions(content string) []Function {
+	var functions []Function
+
+	// Match function declarations: return_type function_name(params);
+	// Matches both mlx_* and _mlx_* functions
+	pattern := regexp.MustCompile(`\b((?:const\s+)?(?:struct\s+)?[\w\s]+?[\*\s]*)\s+(_?mlx_\w+)\s*\(([^)]*(?:\([^)]*\)[^)]*)*)\)\s*;`)
+
+	matches := pattern.FindAllStringSubmatch(content, -1)
+	for _, match := range matches {
+		returnType := strings.TrimSpace(match[1])
+		funcName := strings.TrimSpace(match[2])
+		params := strings.TrimSpace(match[3])
+
+		// Skip if this looks like a variable declaration
+		if params == "" || strings.Contains(params, "{") {
+			continue
+		}
+
+		// Clean up return type
+		returnType = strings.Join(strings.Fields(returnType), " ")
+
+		// Extract parameter names
+		paramNames := extractParamNames(params)
+
+		// Check if ARM64 guard is needed
+		needsGuard := needsARM64Guard(funcName, returnType, params)
+
+		functions = append(functions, Function{
+			Name:            funcName,
+			ReturnType:      returnType,
+			Params:          params,
+			ParamNames:      paramNames,
+			NeedsARM64Guard: needsGuard,
+		})
+	}
+
+	return functions
+}
+
+func needsARM64Guard(name, retType, params string) bool {
+	return strings.Contains(name, "float16") ||
+		strings.Contains(name, "bfloat16") ||
+		strings.Contains(retType, "float16_t") ||
+		strings.Contains(retType, "bfloat16_t") ||
+		strings.Contains(params, "float16_t") ||
+		strings.Contains(params, "bfloat16_t")
+}
+
+func generateWrapperFiles(functions []Function, headerPath, implPath string) error {
+	// Generate header file
+	var headerBuf bytes.Buffer
+
+	headerBuf.WriteString("// AUTO-GENERATED by generate_wrappers.go - DO NOT EDIT\n")
+	headerBuf.WriteString("// This file provides wrapper declarations for MLX-C functions that use dlopen/dlsym\n")
+	headerBuf.WriteString("//\n")
+	headerBuf.WriteString("// Strategy: Include MLX-C headers for type definitions, then provide wrapper\n")
+	headerBuf.WriteString("// functions that shadow the originals, allowing Go code to call them directly (e.g., C.mlx_add).\n")
+	headerBuf.WriteString("// Function pointers are defined in mlx.c (single compilation unit).\n\n")
+	headerBuf.WriteString("#ifndef MLX_WRAPPERS_H\n")
+	headerBuf.WriteString("#define MLX_WRAPPERS_H\n\n")
+
+	headerBuf.WriteString("// Include MLX headers for type definitions and original declarations\n")
+	headerBuf.WriteString("#include \"mlx/c/mlx.h\"\n")
+	headerBuf.WriteString("#include \"mlx_dynamic.h\"\n")
+	headerBuf.WriteString("#include <stdio.h>\n\n")
+
+	// Undef all MLX functions to avoid conflicts
+	headerBuf.WriteString("// Undefine any existing MLX function macros\n")
+	for _, fn := range functions {
+		headerBuf.WriteString(fmt.Sprintf("#undef %s\n", fn.Name))
+	}
+	headerBuf.WriteString("\n")
+
+	// Function pointer extern declarations
+	headerBuf.WriteString("// Function pointer declarations (defined in mlx.c, loaded via dlsym)\n")
+	for _, fn := range functions {
+		if fn.NeedsARM64Guard {
+			headerBuf.WriteString("#if defined(__aarch64__) || defined(_M_ARM64)\n")
+		}
+		headerBuf.WriteString(fmt.Sprintf("extern %s (*%s_ptr)(%s);\n", fn.ReturnType, fn.Name, fn.Params))
+		if fn.NeedsARM64Guard {
+			headerBuf.WriteString("#endif\n")
+		}
+	}
+	headerBuf.WriteString("\n")
+
+	// Initialization function declaration
+	headerBuf.WriteString("// Initialize all function pointers via dlsym (defined in mlx.c)\n")
+	headerBuf.WriteString("int mlx_load_functions(void* handle);\n\n")
+
+	// Wrapper function declarations
+	headerBuf.WriteString("// Wrapper function declarations that call through function pointers\n")
+	headerBuf.WriteString("// Go code calls these directly as C.mlx_* (no #define redirection needed)\n")
+	for _, fn := range functions {
+		if fn.NeedsARM64Guard {
+			headerBuf.WriteString("#if defined(__aarch64__) || defined(_M_ARM64)\n")
+		}
+		headerBuf.WriteString(fmt.Sprintf("%s %s(%s);\n", fn.ReturnType, fn.Name, fn.Params))
+		if fn.NeedsARM64Guard {
+			headerBuf.WriteString("#endif\n")
+		}
+		headerBuf.WriteString("\n")
+	}
+
+	headerBuf.WriteString("#endif // MLX_WRAPPERS_H\n")
+
+	// Write header file
+	if err := os.WriteFile(headerPath, headerBuf.Bytes(), 0644); err != nil {
+		return fmt.Errorf("failed to write header file: %w", err)
+	}
+
+	// Generate implementation file
+	var implBuf bytes.Buffer
+
+	implBuf.WriteString("// AUTO-GENERATED by generate_wrappers.go - DO NOT EDIT\n")
+	implBuf.WriteString("// This file contains the function pointer definitions and initialization\n")
+	implBuf.WriteString("// All function pointers are in a single compilation unit to avoid duplication\n\n")
+
+	implBuf.WriteString("#include \"mlx/c/mlx.h\"\n")
+	implBuf.WriteString("#include \"mlx_dynamic.h\"\n")
+	implBuf.WriteString("#include <stdio.h>\n\n")
+	implBuf.WriteString("// Platform-specific dynamic loading\n")
+	implBuf.WriteString("#ifdef _WIN32\n")
+	implBuf.WriteString("#include <windows.h>\n")
+	implBuf.WriteString("#define GET_SYM(handle, name) (void*)GetProcAddress((HMODULE)(handle), name)\n")
+	implBuf.WriteString("#else\n")
+	implBuf.WriteString("#include <dlfcn.h>\n")
+	implBuf.WriteString("#define GET_SYM(handle, name) dlsym(handle, name)\n")
+	implBuf.WriteString("#endif\n\n")
+
+	// Function pointer definitions
+	implBuf.WriteString("// Function pointer definitions\n")
+	for _, fn := range functions {
+		if fn.NeedsARM64Guard {
+			implBuf.WriteString("#if defined(__aarch64__) || defined(_M_ARM64)\n")
+		}
+		implBuf.WriteString(fmt.Sprintf("%s (*%s_ptr)(%s) = NULL;\n", fn.ReturnType, fn.Name, fn.Params))
+		if fn.NeedsARM64Guard {
+			implBuf.WriteString("#endif\n")
+		}
+	}
+	implBuf.WriteString("\n")
+
+	// Initialization function
+	implBuf.WriteString("// Initialize all function pointers\n")
+	implBuf.WriteString("int mlx_load_functions(void* handle) {\n")
+	implBuf.WriteString("    if (handle == NULL) {\n")
+	implBuf.WriteString("        fprintf(stderr, \"MLX: Invalid library handle\\n\");\n")
+	implBuf.WriteString("        return -1;\n")
+	implBuf.WriteString("    }\n\n")
+
+	for _, fn := range functions {
+		if fn.NeedsARM64Guard {
+			implBuf.WriteString("#if defined(__aarch64__) || defined(_M_ARM64)\n")
+		}
+		implBuf.WriteString(fmt.Sprintf("    %s_ptr = GET_SYM(handle, \"%s\");\n", fn.Name, fn.Name))
+		implBuf.WriteString(fmt.Sprintf("    if (%s_ptr == NULL) {\n", fn.Name))
+		implBuf.WriteString(fmt.Sprintf("        fprintf(stderr, \"MLX: Failed to load symbol: %s\\n\");\n", fn.Name))
+		implBuf.WriteString("        return -1;\n")
+		implBuf.WriteString("    }\n")
+		if fn.NeedsARM64Guard {
+			implBuf.WriteString("#endif\n")
+		}
+	}
+
+	implBuf.WriteString("    return 0;\n")
+	implBuf.WriteString("}\n\n")
+
+	// Wrapper function implementations
+	implBuf.WriteString("// Wrapper function implementations that call through function pointers\n")
+	for _, fn := range functions {
+		if fn.NeedsARM64Guard {
+			implBuf.WriteString("#if defined(__aarch64__) || defined(_M_ARM64)\n")
+		}
+		implBuf.WriteString(fmt.Sprintf("%s %s(%s) {\n", fn.ReturnType, fn.Name, fn.Params))
+
+		// Call through function pointer
+		if fn.ReturnType != "void" {
+			implBuf.WriteString(fmt.Sprintf("    return %s_ptr(", fn.Name))
+		} else {
+			implBuf.WriteString(fmt.Sprintf("    %s_ptr(", fn.Name))
+		}
+
+		// Pass parameters
+		implBuf.WriteString(strings.Join(fn.ParamNames, ", "))
+		implBuf.WriteString(");\n")
+		implBuf.WriteString("}\n")
+		if fn.NeedsARM64Guard {
+			implBuf.WriteString("#endif\n")
+		}
+		implBuf.WriteString("\n")
+	}
+
+	// Write implementation file
+	if err := os.WriteFile(implPath, implBuf.Bytes(), 0644); err != nil {
+		return fmt.Errorf("failed to write implementation file: %w", err)
+	}
+
+	return nil
+}
+
+func main() {
+	flag.Usage = func() {
+		fmt.Fprintf(flag.CommandLine.Output(), "Usage: go run generate_wrappers.go <mlx-c-include-dir> <output-header> [output-impl]\n")
+		fmt.Fprintf(flag.CommandLine.Output(), "Generate MLX-C dynamic loading wrappers.\n\n")
+		flag.PrintDefaults()
+	}
+	flag.Parse()
+
+	args := flag.Args()
+	if len(args) < 2 {
+		fmt.Fprintf(flag.CommandLine.Output(), "ERROR: Missing required arguments\n\n")
+		flag.Usage()
+		os.Exit(1)
+	}
+
+	headerDir := args[0]
+	outputHeader := args[1]
+	// Default implementation file is same name with .c extension
+	outputImpl := outputHeader
+	if len(args) > 2 {
+		outputImpl = args[2]
+	} else if strings.HasSuffix(outputHeader, ".h") {
+		outputImpl = outputHeader[:len(outputHeader)-2] + ".c"
+	}
+
+	// Check if header directory exists
+	if _, err := os.Stat(headerDir); os.IsNotExist(err) {
+		fmt.Fprintf(os.Stderr, "ERROR: MLX-C headers directory not found at: %s\n\n", headerDir)
+		fmt.Fprintf(os.Stderr, "Please run CMake first to download MLX-C dependencies:\n")
+		fmt.Fprintf(os.Stderr, "  cmake -B build\n\n")
+		fmt.Fprintf(os.Stderr, "The CMake build will download and extract MLX-C headers needed for wrapper generation.\n")
+		os.Exit(1)
+	}
+
+	fmt.Fprintf(os.Stderr, "Parsing MLX-C headers from: %s\n", headerDir)
+
+	// Find all headers
+	headers, err := findHeaders(headerDir)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "ERROR: Failed to find header files: %v\n", err)
+		os.Exit(1)
+	}
+	fmt.Fprintf(os.Stderr, "Found %d header files\n", len(headers))
+
+	// Parse all headers
+	var allFunctions []Function
+	seen := make(map[string]bool)
+
+	for _, header := range headers {
+		content, err := os.ReadFile(header)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "Error reading %s: %v\n", header, err)
+			continue
+		}
+
+		cleaned := cleanContent(string(content))
+		functions := parseFunctions(cleaned)
+
+		// Deduplicate
+		for _, fn := range functions {
+			if !seen[fn.Name] {
+				seen[fn.Name] = true
+				allFunctions = append(allFunctions, fn)
+			}
+		}
+	}
+
+	fmt.Fprintf(os.Stderr, "Found %d unique function declarations\n", len(allFunctions))
+
+	// Generate wrapper files
+	if err := generateWrapperFiles(allFunctions, outputHeader, outputImpl); err != nil {
+		fmt.Fprintf(os.Stderr, "ERROR: Failed to generate wrapper files: %v\n", err)
+		os.Exit(1)
+	}
+
+	fmt.Fprintf(os.Stderr, "Generated %s and %s successfully\n", outputHeader, outputImpl)
+}

x/imagegen/mlx/mlx_dynamic.c

@@ -0,0 +1,93 @@
+// mlx_dynamic.c - Dynamic loading wrapper for MLX-C library
+// This file provides runtime dynamic loading of libmlxc instead of link-time binding
+
+#include "mlx_dynamic.h"
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#ifdef _WIN32
+#include <windows.h>
+typedef HMODULE lib_handle_t;
+static char win_error_buffer[256] = {0};
+static const char* get_win_error(void) {
+    DWORD err = GetLastError();
+    snprintf(win_error_buffer, sizeof(win_error_buffer), "error code %lu", err);
+    return win_error_buffer;
+}
+#define LIB_ERROR() get_win_error()
+#else
+#include <dlfcn.h>
+typedef void* lib_handle_t;
+#define LIB_ERROR() dlerror()
+#endif
+
+static lib_handle_t mlx_handle = NULL;
+static int mlx_initialized = 0;
+static char mlx_error_buffer[512] = {0};
+
+#ifdef _WIN32
+// Windows: Load library from a path with dependency resolution.
+// Temporarily adds the library's directory to the DLL search path
+// so that dependencies (like mlx.dll) in the same directory are found.
+static int try_load_win(const char* path) {
+    if (!path) return 0;
+
+    // Extract directory and add to DLL search path for dependency resolution
+    char dir_path[MAX_PATH];
+    strncpy(dir_path, path, MAX_PATH - 1);
+    dir_path[MAX_PATH - 1] = '\0';
+    char* last_slash = strrchr(dir_path, '\\');
+    if (!last_slash) last_slash = strrchr(dir_path, '/');
+    if (last_slash) {
+        *last_slash = '\0';
+        SetDllDirectoryA(dir_path);
+    }
+
+    mlx_handle = LoadLibraryA(path);
+    SetDllDirectoryA(NULL);
+    return mlx_handle != NULL;
+}
+#endif
+
+// Try to load library from a specific path
+static int try_load_lib(const char* path) {
+    if (!path) return 0;
+#ifdef _WIN32
+    return try_load_win(path);
+#else
+    mlx_handle = dlopen(path, RTLD_LAZY | RTLD_GLOBAL);
+    return mlx_handle != NULL;
+#endif
+}
+
+// Initialize the MLX dynamic library from a specific path.
+// Returns 0 on success, -1 on failure.
+int mlx_dynamic_init_path(const char* path) {
+    if (mlx_initialized) {
+        return 0;
+    }
+
+    if (try_load_lib(path)) {
+        mlx_initialized = 1;
+        snprintf(mlx_error_buffer, sizeof(mlx_error_buffer),
+                 "MLX: Successfully loaded %s", path ? path : "library");
+        return 0;
+    }
+
+    const char* err = LIB_ERROR();
+    snprintf(mlx_error_buffer, sizeof(mlx_error_buffer),
+             "MLX: Failed to load %s: %s", path ? path : "(null)", err ? err : "unknown error");
+    return -1;
+}
+
+// Get the last error message
+const char* mlx_dynamic_error(void) {
+    return mlx_error_buffer;
+}
+
+// Get the library handle (for use by generated wrappers)
+void* mlx_get_handle(void) {
+    return mlx_handle;
+}
+

x/imagegen/mlx/mlx_dynamic.h

@@ -0,0 +1,23 @@
+// mlx_dynamic.h - Dynamic loading interface for MLX-C library
+#ifndef MLX_DYNAMIC_H
+#define MLX_DYNAMIC_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Initialize the MLX dynamic library from a specific path
+// Returns 0 on success, -1 on failure
+int mlx_dynamic_init_path(const char* path);
+
+// Get the last error message from dynamic loading
+const char* mlx_dynamic_error(void);
+
+// Get the library handle (for use by generated wrappers)
+void* mlx_get_handle(void);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLX_DYNAMIC_H

x/imagegen/mlx/mlx_error_handler.c

@@ -0,0 +1,38 @@
+// mlx_error_handler.c - Safe error handling for MLX initialization
+// Provides a non-fatal error handler for use during init(), so that
+// GPU failures are captured instead of calling exit(-1).
+
+#include "mlx_error_handler.h"
+#include "mlx.h"
+#include <string.h>
+
+static char mlx_init_error_msg[1024] = {0};
+static int  mlx_init_error_flag = 0;
+
+// Error handler that silently stores the error message.
+// The error is surfaced on the Go side via mlxInitError / GetMLXInitError()
+// only when MLX is actually needed.
+static void mlx_silent_error_handler(const char* msg, void* data) {
+    (void)data;
+    strncpy(mlx_init_error_msg, msg, sizeof(mlx_init_error_msg) - 1);
+    mlx_init_error_msg[sizeof(mlx_init_error_msg) - 1] = '\0';
+    mlx_init_error_flag = 1;
+}
+
+void mlx_set_safe_init_mode(void) {
+    mlx_init_error_flag = 0;
+    mlx_init_error_msg[0] = '\0';
+    mlx_set_error_handler(mlx_silent_error_handler, NULL, NULL);
+}
+
+void mlx_set_default_error_mode(void) {
+    mlx_set_error_handler(NULL, NULL, NULL);
+}
+
+int mlx_had_init_error(void) {
+    return mlx_init_error_flag;
+}
+
+const char* mlx_get_init_error(void) {
+    return mlx_init_error_flag ? mlx_init_error_msg : NULL;
+}

x/imagegen/mlx/mlx_error_handler.h

@@ -0,0 +1,22 @@
+// mlx_error_handler.h - Safe error handling for MLX initialization
+// This replaces the default exit(-1) MLX error handler during init()
+// so that GPU failures don't kill the process.
+
+#ifndef MLX_ERROR_HANDLER_H
+#define MLX_ERROR_HANDLER_H
+
+// Enter safe mode before any MLX compute calls during init().
+// Replaces the default exit(-1) handler with one that silently stores errors.
+void mlx_set_safe_init_mode(void);
+
+// Restore the default MLX error handler (exit on error).
+// Call from runner entry points after confirming MLX is available.
+void mlx_set_default_error_mode(void);
+
+// Check whether an error occurred while in safe init mode.
+int mlx_had_init_error(void);
+
+// Get the error message from the last init error, or NULL if none.
+const char* mlx_get_init_error(void);
+
+#endif // MLX_ERROR_HANDLER_H

x/imagegen/models/flux2/flux2.go

@@ -0,0 +1,551 @@
+// Package flux2 implements the FLUX.2 Klein diffusion transformer model.
+// Klein is a 4B parameter distilled model that supports sub-second inference.
+package flux2
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"image"
+	"math"
+	"time"
+
+	"github.com/ollama/ollama/x/imagegen/manifest"
+	"github.com/ollama/ollama/x/imagegen/mlx"
+	"github.com/ollama/ollama/x/imagegen/models/qwen3"
+	"github.com/ollama/ollama/x/imagegen/tokenizer"
+	"golang.org/x/image/draw"
+)
+
+// GenerateConfig holds all options for image generation.
+type GenerateConfig struct {
+	Prompt        string
+	Width         int32                 // Image width (default: 1024)
+	Height        int32                 // Image height (default: 1024)
+	Steps         int                   // Denoising steps (default: 4 for Klein)
+	GuidanceScale float32               // Guidance scale (default: 1.0, Klein doesn't need CFG)
+	Seed          int64                 // Random seed
+	Progress      func(step, totalSteps int) // Optional progress callback
+	CapturePath   string                // GPU capture path (debug)
+	InputImages   []image.Image         // Reference images for image conditioning (already loaded)
+}
+
+// Model represents a FLUX.2 Klein model.
+type Model struct {
+	ModelName       string
+	Tokenizer       *tokenizer.Tokenizer
+	TextEncoder     *qwen3.TextEncoder
+	Transformer     *Flux2Transformer2DModel
+	VAE             *AutoencoderKLFlux2
+	SchedulerConfig *SchedulerConfig
+}
+
+// TextEncoderLayerIndices are the layers from which to extract text embeddings.
+// Diffusers uses hidden_states[9, 18, 27]. In Python, hidden_states[0] is the embedding
+// output before any layers, so hidden_states[9] = after layer 8 (0-indexed).
+// Go's ForwardWithLayerOutputs captures after layer i runs, so we use [8, 17, 26].
+var TextEncoderLayerIndices = []int{8, 17, 26}
+
+// Load loads the FLUX.2 Klein model from ollama blob storage.
+func (m *Model) Load(modelName string) error {
+	fmt.Printf("Loading FLUX.2 Klein model from manifest: %s...\n", modelName)
+	start := time.Now()
+
+	if mlx.GPUIsAvailable() {
+		mlx.SetDefaultDeviceGPU()
+		mlx.EnableCompile()
+	}
+
+	m.ModelName = modelName
+
+	// Load manifest
+	manifest, err := manifest.LoadManifest(modelName)
+	if err != nil {
+		return fmt.Errorf("load manifest: %w", err)
+	}
+
+	// Load tokenizer
+	fmt.Print("  Loading tokenizer... ")
+	tokData, err := manifest.ReadConfig("tokenizer/tokenizer.json")
+	if err != nil {
+		return fmt.Errorf("tokenizer: %w", err)
+	}
+
+	tokConfig := &tokenizer.TokenizerConfig{}
+	if data, err := manifest.ReadConfig("tokenizer/tokenizer_config.json"); err == nil {
+		tokConfig.TokenizerConfigJSON = data
+	}
+	if data, err := manifest.ReadConfig("tokenizer/generation_config.json"); err == nil {
+		tokConfig.GenerationConfigJSON = data
+	}
+	if data, err := manifest.ReadConfig("tokenizer/special_tokens_map.json"); err == nil {
+		tokConfig.SpecialTokensMapJSON = data
+	}
+
+	tok, err := tokenizer.LoadFromBytesWithConfig(tokData, tokConfig)
+	if err != nil {
+		return fmt.Errorf("tokenizer: %w", err)
+	}
+	m.Tokenizer = tok
+	fmt.Println("✓")
+
+	// Load text encoder
+	m.TextEncoder = &qwen3.TextEncoder{}
+	if err := m.TextEncoder.Load(manifest, "text_encoder/config.json"); err != nil {
+		return fmt.Errorf("text encoder: %w", err)
+	}
+
+	// Load transformer
+	m.Transformer = &Flux2Transformer2DModel{}
+	if err := m.Transformer.Load(manifest); err != nil {
+		return fmt.Errorf("transformer: %w", err)
+	}
+
+	// Load VAE
+	m.VAE = &AutoencoderKLFlux2{}
+	if err := m.VAE.Load(manifest); err != nil {
+		return fmt.Errorf("VAE: %w", err)
+	}
+
+	// Evaluate all weights in a single batch (reduces GPU sync overhead)
+	fmt.Print("  Evaluating weights... ")
+	allWeights := mlx.Collect(m.TextEncoder)
+	allWeights = append(allWeights, mlx.Collect(m.Transformer)...)
+	allWeights = append(allWeights, mlx.Collect(m.VAE)...)
+	mlx.Eval(allWeights...)
+	fmt.Println("✓")
+
+	// Load scheduler config
+	m.SchedulerConfig = DefaultSchedulerConfig()
+	if schedData, err := manifest.ReadConfig("scheduler/scheduler_config.json"); err == nil {
+		if err := json.Unmarshal(schedData, m.SchedulerConfig); err != nil {
+			fmt.Printf("  Warning: failed to parse scheduler config: %v\n", err)
+		}
+	}
+
+	mem := mlx.MetalGetActiveMemory()
+	fmt.Printf("  Loaded in %.2fs (%.1f GB VRAM)\n", time.Since(start).Seconds(), float64(mem)/(1024*1024*1024))
+
+	return nil
+}
+
+// Generate creates an image from a prompt.
+func (m *Model) Generate(prompt string, width, height int32, steps int, seed int64) (*mlx.Array, error) {
+	return m.GenerateFromConfig(context.Background(), &GenerateConfig{
+		Prompt: prompt,
+		Width:  width,
+		Height: height,
+		Steps:  steps,
+		Seed:   seed,
+	})
+}
+
+// GenerateWithProgress creates an image with progress callback.
+func (m *Model) GenerateWithProgress(prompt string, width, height int32, steps int, seed int64, progress func(step, totalSteps int)) (*mlx.Array, error) {
+	return m.GenerateFromConfig(context.Background(), &GenerateConfig{
+		Prompt:   prompt,
+		Width:    width,
+		Height:   height,
+		Steps:    steps,
+		Seed:     seed,
+		Progress: progress,
+	})
+}
+
+// GenerateFromConfig generates an image using the unified config struct.
+func (m *Model) GenerateFromConfig(ctx context.Context, cfg *GenerateConfig) (*mlx.Array, error) {
+	start := time.Now()
+	result, err := m.generate(ctx, cfg)
+	if err != nil {
+		return nil, err
+	}
+	fmt.Printf("Generated in %.2fs (%d steps)\n", time.Since(start).Seconds(), cfg.Steps)
+	return result, nil
+}
+
+// GenerateImage implements runner.ImageModel interface.
+func (m *Model) GenerateImage(ctx context.Context, prompt string, width, height int32, steps int, seed int64, progress func(step, total int)) (*mlx.Array, error) {
+	return m.GenerateFromConfig(ctx, &GenerateConfig{
+		Prompt:   prompt,
+		Width:    width,
+		Height:   height,
+		Steps:    steps,
+		Seed:     seed,
+		Progress: progress,
+	})
+}
+
+// GenerateImageWithInputs implements runner.ImageEditModel interface.
+// It generates an image conditioned on the provided input images for image editing.
+func (m *Model) GenerateImageWithInputs(ctx context.Context, prompt string, width, height int32, steps int, seed int64, inputImages []image.Image, progress func(step, total int)) (*mlx.Array, error) {
+	return m.GenerateFromConfig(ctx, &GenerateConfig{
+		Prompt:      prompt,
+		Width:       width,
+		Height:      height,
+		Steps:       steps,
+		Seed:        seed,
+		InputImages: inputImages,
+		Progress:    progress,
+	})
+}
+
+// MaxOutputPixels is the maximum output resolution (4 megapixels, ~2048x2048)
+const MaxOutputPixels = 2048 * 2048
+
+// MaxRefPixels is the maximum resolution for reference images (smaller to reduce attention memory)
+const MaxRefPixels = 728 * 728
+
+// generate is the internal denoising pipeline.
+func (m *Model) generate(ctx context.Context, cfg *GenerateConfig) (*mlx.Array, error) {
+	// Enable MLX compilation for fused kernels
+	mlx.EnableCompile()
+
+	// Apply defaults
+	if cfg.Steps <= 0 {
+		cfg.Steps = 4 // Klein default: 4 steps for distilled model
+	}
+	if cfg.GuidanceScale <= 0 {
+		cfg.GuidanceScale = 1.0 // Klein doesn't need guidance
+	}
+
+	// Determine output dimensions
+	if len(cfg.InputImages) > 0 {
+		// With input images, compute missing dimension from aspect ratio
+		// Images are already EXIF-rotated by the caller
+		bounds := cfg.InputImages[0].Bounds()
+		imgW, imgH := bounds.Dx(), bounds.Dy()
+		aspectRatio := float64(imgH) / float64(imgW)
+		if cfg.Width > 0 && cfg.Height <= 0 {
+			// Width specified, compute height
+			cfg.Height = int32(math.Round(float64(cfg.Width)*aspectRatio/16) * 16)
+		} else if cfg.Height > 0 && cfg.Width <= 0 {
+			// Height specified, compute width
+			cfg.Width = int32(math.Round(float64(cfg.Height)/aspectRatio/16) * 16)
+		} else if cfg.Width <= 0 && cfg.Height <= 0 {
+			// Neither specified, use input dimensions
+			cfg.Width = int32(imgW)
+			cfg.Height = int32(imgH)
+		}
+	}
+	if cfg.Width <= 0 {
+		cfg.Width = 1024
+	}
+	if cfg.Height <= 0 {
+		cfg.Height = 1024
+	}
+
+	// Cap to max pixels, preserve aspect ratio, round to multiple of 16
+	pixels := int(cfg.Width) * int(cfg.Height)
+	if pixels > MaxOutputPixels {
+		scale := math.Sqrt(float64(MaxOutputPixels) / float64(pixels))
+		cfg.Width = int32(math.Round(float64(cfg.Width) * scale / 16) * 16)
+		cfg.Height = int32(math.Round(float64(cfg.Height) * scale / 16) * 16)
+	}
+	cfg.Height = int32((cfg.Height + 8) / 16 * 16) // round to nearest 16
+	cfg.Width = int32((cfg.Width + 8) / 16 * 16)
+	fmt.Printf("  Output: %dx%d\n", cfg.Width, cfg.Height)
+
+	tcfg := m.Transformer.TransformerConfig
+	patchSize := m.VAE.Config.PatchSize
+
+	// Latent dimensions: image / 8 (VAE downscale) / patch_size
+	latentH := cfg.Height / 8
+	latentW := cfg.Width / 8
+	patchH := latentH / patchSize[0]
+	patchW := latentW / patchSize[1]
+	imgSeqLen := patchH * patchW
+
+	// Text encoding with multi-layer extraction (no padding, use true sequence length)
+	fmt.Print("  Encoding prompt... ")
+	promptEmbeds, textLen := m.TextEncoder.EncodePromptWithLayers(m.Tokenizer, cfg.Prompt, 512, TextEncoderLayerIndices, false)
+	fmt.Println("✓")
+
+	// Encode reference images if provided
+	var refTokens *ImageCondTokens
+	var refHeights, refWidths []int32
+	if len(cfg.InputImages) > 0 {
+		fmt.Printf("  Encoding %d reference image(s):\n", len(cfg.InputImages))
+
+		var err error
+		refTokens, err = m.EncodeImageRefs(cfg.InputImages)
+		if err != nil {
+			return nil, fmt.Errorf("encode reference images: %w", err)
+		}
+
+		// Extract heights/widths for RoPE computation (same limits as EncodeImageRefs)
+		limitPixels := MaxRefPixels
+		if len(cfg.InputImages) > 1 {
+			limitPixels = MaxRefPixels / 2
+		}
+		for _, img := range cfg.InputImages {
+			_, w, h := PrepareImage(img, limitPixels)
+			refHeights = append(refHeights, int32(h/16))
+			refWidths = append(refWidths, int32(w/16))
+		}
+	}
+
+	// Scheduler
+	scheduler := NewFlowMatchScheduler(m.SchedulerConfig)
+	scheduler.SetTimestepsWithMu(cfg.Steps, CalculateShift(imgSeqLen, cfg.Steps))
+
+	// Init latents in packed form [B, C*4, H/2, W/2] like diffusers
+	// diffusers creates noise in [B, 128, 64, 64] and packs to [B, 4096, 128]
+	latentChannels := m.VAE.Config.LatentChannels
+	packedChannels := latentChannels * 4 // 32 * 4 = 128
+	latents := scheduler.InitNoise([]int32{1, packedChannels, patchH, patchW}, cfg.Seed)
+
+	// Pack latents (transpose): [B, C, H, W] -> [B, H*W, C]
+	// This matches diffusers' _pack_latents
+	patches := packLatents(latents)
+	noiseSeqLen := patches.Shape()[1]
+
+	// RoPE cache - includes reference images if present
+	rope := PrepareRoPECache(textLen, patchH, patchW, tcfg.AxesDimsRoPE, tcfg.RopeTheta, refHeights, refWidths, ImageRefScale)
+
+	// Cleanup setup arrays when done
+	defer func() {
+		rope.Cos.Free()
+		rope.Sin.Free()
+		promptEmbeds.Free()
+		if refTokens != nil {
+			refTokens.Tokens.Free()
+		}
+	}()
+
+	// Pre-compute all timesteps before the loop to avoid per-step tensor creation
+	timesteps := make([]*mlx.Array, cfg.Steps)
+	for i := 0; i < cfg.Steps; i++ {
+		tCurr := scheduler.Timesteps[i] / float32(m.SchedulerConfig.NumTrainTimesteps)
+		timesteps[i] = mlx.ToBFloat16(mlx.NewArray([]float32{tCurr}, []int32{1}))
+	}
+
+	// Evaluate setup arrays
+	fmt.Print("  Evaluating setup... ")
+	setupStart := time.Now()
+	toEval := []*mlx.Array{promptEmbeds, patches, rope.Cos, rope.Sin}
+	toEval = append(toEval, timesteps...)
+	if refTokens != nil {
+		toEval = append(toEval, refTokens.Tokens)
+	}
+	mlx.Eval(toEval...)
+	mlx.MetalResetPeakMemory() // Reset peak to measure generation separately
+	fmt.Printf("✓ (%.2fs, %.1f GB)\n", time.Since(setupStart).Seconds(),
+		float64(mlx.MetalGetActiveMemory())/(1024*1024*1024))
+
+	if cfg.Progress != nil {
+		cfg.Progress(0, cfg.Steps)
+	}
+
+	loopStart := time.Now()
+	stepStart := time.Now()
+
+	// Denoising loop
+	for i := 0; i < cfg.Steps; i++ {
+		// Check for cancellation
+		if ctx != nil {
+			select {
+			case <-ctx.Done():
+				return nil, ctx.Err()
+			default:
+			}
+		}
+
+		// GPU capture on step 2 if requested
+		if cfg.CapturePath != "" && i == 1 {
+			mlx.MetalStartCapture(cfg.CapturePath)
+		}
+
+		timestep := timesteps[i]
+
+		// Prepare input - concatenate noise patches with reference tokens if present
+		imgInput := patches
+		if refTokens != nil {
+			imgInput = mlx.Concatenate([]*mlx.Array{patches, refTokens.Tokens}, 1)
+		}
+
+		// Transformer forward pass
+		output := m.Transformer.Forward(imgInput, promptEmbeds, timestep, rope)
+
+		// If we concatenated reference tokens, slice to only get noise portion
+		if refTokens != nil {
+			output = mlx.Slice(output, []int32{0, 0, 0}, []int32{1, noiseSeqLen, output.Shape()[2]})
+		}
+
+		// Scheduler step (keep reference to old patches for the computation graph)
+		newPatches := scheduler.Step(output, patches, i)
+
+		if cfg.CapturePath != "" && i == 1 {
+			mlx.MetalStopCapture()
+		}
+
+		mlx.Eval(newPatches)
+		patches = newPatches
+
+		elapsed := time.Since(stepStart).Seconds()
+		peakGB := float64(mlx.MetalGetPeakMemory()) / (1024 * 1024 * 1024)
+		if i == 0 {
+			fmt.Printf("    step %d: %.2fs (JIT warmup), peak %.1f GB\n", i+1, elapsed, peakGB)
+		} else {
+			fmt.Printf("    step %d: %.2fs, peak %.1f GB\n", i+1, elapsed, peakGB)
+		}
+		stepStart = time.Now()
+		if cfg.Progress != nil {
+			cfg.Progress(i+1, cfg.Steps)
+		}
+	}
+
+	loopTime := time.Since(loopStart).Seconds()
+	peakMem := float64(mlx.MetalGetPeakMemory()) / (1024 * 1024 * 1024)
+	fmt.Printf("  Denoised %d steps in %.2fs (%.2fs/step), peak %.1f GB\n",
+		cfg.Steps, loopTime, loopTime/float64(cfg.Steps), peakMem)
+
+	// Free timesteps now that denoising is done
+	for _, ts := range timesteps {
+		ts.Free()
+	}
+
+	// VAE decode with tiling for larger images
+	fmt.Print("  Decoding VAE... ")
+	vaeStart := time.Now()
+	// Enable tiling for images > 512x512 (latent > 64x64)
+	// VAE attention is O(n²) on latent pixels, tiling reduces memory significantly
+	if patchH*2 > 64 || patchW*2 > 64 {
+		m.VAE.Tiling = DefaultTilingConfig()
+	}
+	decoded := m.VAE.Decode(patches, patchH, patchW)
+	mlx.Eval(decoded)
+
+	// Free patches now that decode is done
+	patches.Free()
+
+	fmt.Printf("✓ (%.2fs, peak %.1f GB)\n", time.Since(vaeStart).Seconds(),
+		float64(mlx.MetalGetPeakMemory())/(1024*1024*1024))
+
+	return decoded, nil
+}
+
+// packLatents converts [B, C, H, W] to [B, H*W, C] (matches diffusers _pack_latents)
+func packLatents(x *mlx.Array) *mlx.Array {
+	shape := x.Shape()
+	B := shape[0]
+	C := shape[1]
+	H := shape[2]
+	W := shape[3]
+	// [B, C, H, W] -> [B, C, H*W] -> [B, H*W, C]
+	x = mlx.Reshape(x, B, C, H*W)
+	return mlx.Transpose(x, 0, 2, 1)
+}
+
+// LoadPersistent loads the model and keeps it in memory for repeated use.
+func LoadPersistent(modelName string) (*Model, error) {
+	m := &Model{}
+	if err := m.Load(modelName); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
+// ImageRefScale is the time coordinate offset between reference images (matches diffusers scale=10)
+const ImageRefScale = 10
+
+// PrepareImage resizes and crops an image to be a multiple of 16, with optional pixel limit.
+// Returns the processed image and its dimensions.
+func PrepareImage(img image.Image, limitPixels int) (image.Image, int, int) {
+	bounds := img.Bounds()
+	w, h := bounds.Dx(), bounds.Dy()
+
+	// Cap pixels if needed (like diffusers cap_pixels)
+	if limitPixels > 0 && w*h > limitPixels {
+		scale := math.Sqrt(float64(limitPixels) / float64(w*h))
+		w = int(float64(w) * scale)
+		h = int(float64(h) * scale)
+	}
+
+	// Round down to multiple of 16
+	w = (w / 16) * 16
+	h = (h / 16) * 16
+
+	if w < 16 {
+		w = 16
+	}
+	if h < 16 {
+		h = 16
+	}
+
+	// Resize using high-quality bicubic interpolation (matches diffusers' default lanczos)
+	resized := image.NewRGBA(image.Rect(0, 0, w, h))
+	draw.CatmullRom.Scale(resized, resized.Bounds(), img, img.Bounds(), draw.Over, nil)
+
+	return resized, w, h
+}
+
+// ImageToTensor converts an image to a tensor in [-1, 1] range with shape [1, C, H, W].
+func ImageToTensor(img image.Image) *mlx.Array {
+	bounds := img.Bounds()
+	w, h := bounds.Dx(), bounds.Dy()
+
+	// Convert to float32 array in NCHW format [1, 3, H, W] with values in [-1, 1]
+	data := make([]float32, 3*h*w)
+
+	for y := 0; y < h; y++ {
+		for x := 0; x < w; x++ {
+			r, g, b, _ := img.At(x+bounds.Min.X, y+bounds.Min.Y).RGBA()
+			// RGBA returns 16-bit values, convert to [-1, 1]
+			data[0*h*w+y*w+x] = float32(r>>8)/127.5 - 1.0
+			data[1*h*w+y*w+x] = float32(g>>8)/127.5 - 1.0
+			data[2*h*w+y*w+x] = float32(b>>8)/127.5 - 1.0
+		}
+	}
+
+	arr := mlx.NewArrayFloat32(data, []int32{1, 3, int32(h), int32(w)})
+	return arr
+}
+
+// ImageCondTokens holds encoded reference image tokens.
+type ImageCondTokens struct {
+	Tokens *mlx.Array // [1, total_tokens, C] - concatenated reference tokens
+}
+
+// EncodeImageRefs encodes reference images using the VAE.
+func (m *Model) EncodeImageRefs(images []image.Image) (*ImageCondTokens, error) {
+	if len(images) == 0 {
+		return nil, nil
+	}
+
+	// Limit reference images to reduce attention memory
+	limitPixels := MaxRefPixels
+	if len(images) > 1 {
+		limitPixels = MaxRefPixels / 2
+	}
+
+	var allTokens []*mlx.Array
+
+	for _, img := range images {
+		// Prepare image (resize, crop to multiple of 16)
+		prepared, prepW, prepH := PrepareImage(img, limitPixels)
+		fmt.Printf("    Encoding %dx%d image... ", prepW, prepH)
+
+		// Convert to tensor [-1, 1]
+		tensor := ImageToTensor(prepared)
+
+		// Encode with VAE - returns [1, L, 128]
+		encoded := m.VAE.EncodeImage(tensor)
+		squeezed := mlx.Squeeze(encoded, 0) // [L, C]
+
+		// Defer eval - will be done with other setup arrays
+		allTokens = append(allTokens, squeezed)
+		fmt.Println("✓")
+	}
+
+	// For single image, just add batch dimension directly
+	// For multiple images, concatenate first
+	var tokens *mlx.Array
+	if len(allTokens) == 1 {
+		tokens = mlx.ExpandDims(allTokens[0], 0) // [1, L, C]
+	} else {
+		tokens = mlx.Concatenate(allTokens, 0) // [total_L, C]
+		tokens = mlx.ExpandDims(tokens, 0)     // [1, total_L, C]
+	}
+
+	return &ImageCondTokens{Tokens: tokens}, nil
+}

x/imagegen/models/flux2/rope.go

@@ -0,0 +1,222 @@
+package flux2
+
+import (
+	"math"
+
+	"github.com/ollama/ollama/x/imagegen/mlx"
+)
+
+// RoPEConfig holds 4D RoPE configuration for Flux2
+type RoPEConfig struct {
+	Theta    int32   // 2000 for Klein
+	AxesDims []int32 // [32, 32, 32, 32] - dimensions for T, H, W, L axes
+}
+
+// RoPECache holds precomputed RoPE cos/sin values
+type RoPECache struct {
+	Cos      *mlx.Array // [1, TotalSeqLen, 1, head_dim/2]
+	Sin      *mlx.Array // [1, TotalSeqLen, 1, head_dim/2]
+	TextLen  int32      // Length of text sequence
+	ImageLen int32      // Length of image sequence
+}
+
+// PrepareTextIDs creates position IDs for text tokens.
+// Text tokens use: T=0, H=0, W=0, L=0..seqLen-1
+// Returns: [seqLen, 4]
+func PrepareTextIDs(seqLen int32) *mlx.Array {
+	ids := make([]float32, seqLen*4)
+	for i := int32(0); i < seqLen; i++ {
+		idx := i * 4
+		ids[idx+0] = 0             // T = 0
+		ids[idx+1] = 0             // H = 0
+		ids[idx+2] = 0             // W = 0
+		ids[idx+3] = float32(i)    // L = sequence position
+	}
+	return mlx.NewArray(ids, []int32{seqLen, 4})
+}
+
+// PrepareLatentIDs creates position IDs for image latent tokens.
+// Latent tokens use: T=0, H=0..height-1, W=0..width-1, L=0
+// The latents are in row-major order (H then W).
+// Returns: [height*width, 4]
+func PrepareLatentIDs(height, width int32) *mlx.Array {
+	seqLen := height * width
+	ids := make([]float32, seqLen*4)
+	idx := 0
+	for h := int32(0); h < height; h++ {
+		for w := int32(0); w < width; w++ {
+			ids[idx*4+0] = 0           // T = 0
+			ids[idx*4+1] = float32(h)  // H = row
+			ids[idx*4+2] = float32(w)  // W = column
+			ids[idx*4+3] = 0           // L = 0
+			idx++
+		}
+	}
+	return mlx.NewArray(ids, []int32{seqLen, 4})
+}
+
+// PrepareImageIDs creates position IDs for reference image tokens (used in editing).
+// Reference images use: T=scale*(i+1), H=0..h-1, W=0..w-1, L=0
+// where i is the image index (0, 1, 2, ...) and scale separates images in T dimension.
+// Returns: [total_tokens, 4]
+func PrepareImageIDs(imageHeights, imageWidths []int32, scale int32) *mlx.Array {
+	// Calculate total tokens
+	totalTokens := int32(0)
+	for i := range imageHeights {
+		totalTokens += imageHeights[i] * imageWidths[i]
+	}
+
+	ids := make([]float32, totalTokens*4)
+	idx := int32(0)
+	for imgIdx, h := range imageHeights {
+		w := imageWidths[imgIdx]
+		tValue := float32(scale * int32(imgIdx+1))
+		for hi := int32(0); hi < h; hi++ {
+			for wi := int32(0); wi < w; wi++ {
+				ids[idx*4+0] = tValue       // T = scale * (imgIdx + 1)
+				ids[idx*4+1] = float32(hi)  // H = row
+				ids[idx*4+2] = float32(wi)  // W = column
+				ids[idx*4+3] = 0            // L = 0
+				idx++
+			}
+		}
+	}
+	return mlx.NewArray(ids, []int32{totalTokens, 4})
+}
+
+// ComputeRoPE computes cos and sin for 4D rotary position embeddings.
+// ids: [L, 4] with (T, H, W, L) coordinates
+// axesDims: [32, 32, 32, 32] - each axis has this many dimensions (total = head_dim = 128)
+// theta: base frequency (2000 for Klein)
+// Returns: cos, sin each [1, L, 1, head_dim] with repeat_interleave applied
+func ComputeRoPE(ids *mlx.Array, axesDims []int32, theta int32) (*mlx.Array, *mlx.Array) {
+	shape := ids.Shape()
+	seqLen := shape[0]
+
+	// Compute total head dim (sum of all axes dims)
+	headDim := int32(0)
+	for _, d := range axesDims {
+		headDim += d
+	}
+
+	// Extract each coordinate dimension
+	// ids[:, 0] = T, ids[:, 1] = H, ids[:, 2] = W, ids[:, 3] = L
+	posT := mlx.Slice(ids, []int32{0, 0}, []int32{seqLen, 1}) // [L, 1]
+	posH := mlx.Slice(ids, []int32{0, 1}, []int32{seqLen, 2}) // [L, 1]
+	posW := mlx.Slice(ids, []int32{0, 2}, []int32{seqLen, 3}) // [L, 1]
+	posL := mlx.Slice(ids, []int32{0, 3}, []int32{seqLen, 4}) // [L, 1]
+
+	// Compute frequencies for each axis
+	logTheta := float32(math.Log(float64(theta)))
+	cosArrs := make([]*mlx.Array, 4)
+	sinArrs := make([]*mlx.Array, 4)
+	positions := []*mlx.Array{posT, posH, posW, posL}
+
+	for i, axisDim := range axesDims {
+		half := axisDim / 2
+
+		// Create frequency array for this axis: theta^(-2j/dim) for j=0..half-1
+		// This matches diffusers: 1.0 / (theta ** (torch.arange(0, dim, 2) / dim))
+		freqs := make([]float32, half)
+		for j := int32(0); j < half; j++ {
+			freqs[j] = float32(math.Exp(float64(-logTheta * float32(2*j) / float32(axisDim))))
+		}
+		freqArr := mlx.NewArray(freqs, []int32{1, half})
+
+		// Compute pos * freq -> [L, half]
+		posExpanded := positions[i] // [L, 1]
+		args := mlx.Mul(posExpanded, freqArr) // [L, half]
+
+		// Compute cos and sin for this axis
+		cosAxis := mlx.Cos(args) // [L, half]
+		sinAxis := mlx.Sin(args) // [L, half]
+
+		// repeat_interleave(2): [c0, c1, ...] -> [c0, c0, c1, c1, ...]
+		// Reshape [L, half] -> [L, half, 1], tile to [L, half, 2], reshape to [L, axisDim]
+		cosAxis = mlx.ExpandDims(cosAxis, 2)                        // [L, half, 1]
+		cosAxis = mlx.Tile(cosAxis, []int32{1, 1, 2})               // [L, half, 2]
+		cosAxis = mlx.Reshape(cosAxis, seqLen, axisDim)             // [L, axisDim]
+
+		sinAxis = mlx.ExpandDims(sinAxis, 2)
+		sinAxis = mlx.Tile(sinAxis, []int32{1, 1, 2})
+		sinAxis = mlx.Reshape(sinAxis, seqLen, axisDim)
+
+		cosArrs[i] = cosAxis
+		sinArrs[i] = sinAxis
+	}
+
+	// Concatenate all axes: [L, headDim]
+	cos := mlx.Concatenate(cosArrs, 1)
+	sin := mlx.Concatenate(sinArrs, 1)
+
+	// Reshape to [1, L, 1, headDim] for broadcasting with attention
+	cos = mlx.Reshape(cos, 1, seqLen, 1, headDim)
+	sin = mlx.Reshape(sin, 1, seqLen, 1, headDim)
+
+	return cos, sin
+}
+
+// ApplyRoPE4D applies 4D rotary position embeddings to queries and keys.
+// x: [B, L, nheads, head_dim]
+// cos, sin: [1, L, 1, head_dim] (with repeat_interleave applied)
+// Returns: x with RoPE applied
+// Matches diffusers apply_rotary_emb with use_real=True, use_real_unbind_dim=-1
+func ApplyRoPE4D(x *mlx.Array, cos, sin *mlx.Array) *mlx.Array {
+	shape := x.Shape()
+	B := shape[0]
+	L := shape[1]
+	nheads := shape[2]
+	headDim := shape[3]
+	half := headDim / 2
+
+	// Reshape x to [B, L, nheads, half, 2] and split into real/imag
+	xReshaped := mlx.Reshape(x, B, L, nheads, half, 2)
+
+	// Extract real (index 0) and imag (index 1) parts
+	xReal := mlx.Slice(xReshaped, []int32{0, 0, 0, 0, 0}, []int32{B, L, nheads, half, 1})
+	xImag := mlx.Slice(xReshaped, []int32{0, 0, 0, 0, 1}, []int32{B, L, nheads, half, 2})
+	xReal = mlx.Squeeze(xReal, 4) // [B, L, nheads, half]
+	xImag = mlx.Squeeze(xImag, 4) // [B, L, nheads, half]
+
+	// x_rotated = stack([-x_imag, x_real], dim=-1).flatten(-2)
+	// This creates [-x_imag[0], x_real[0], -x_imag[1], x_real[1], ...]
+	negXImag := mlx.Neg(xImag)
+	negXImag = mlx.ExpandDims(negXImag, 4) // [B, L, nheads, half, 1]
+	xReal = mlx.ExpandDims(xReal, 4)       // [B, L, nheads, half, 1]
+	xRotated := mlx.Concatenate([]*mlx.Array{negXImag, xReal}, 4) // [B, L, nheads, half, 2]
+	xRotated = mlx.Reshape(xRotated, B, L, nheads, headDim)       // [B, L, nheads, headDim]
+
+	// out = x * cos + x_rotated * sin
+	return mlx.Add(mlx.Mul(x, cos), mlx.Mul(xRotated, sin))
+}
+
+// PrepareRoPECache creates RoPE cache for text + noise, optionally with reference images.
+// textLen: number of text tokens
+// noiseH, noiseW: dimensions of the noise latent in patch tokens
+// axesDims: [32, 32, 32, 32]
+// theta: 2000
+// refHeights, refWidths: optional reference image dimensions (pass nil/empty for no images)
+// scale: time coordinate offset between reference images (e.g., 10)
+func PrepareRoPECache(textLen, noiseH, noiseW int32, axesDims []int32, theta int32, refHeights, refWidths []int32, scale int32) *RoPECache {
+	textIDs := PrepareTextIDs(textLen)
+	noiseIDs := PrepareLatentIDs(noiseH, noiseW)
+
+	var allIDs *mlx.Array
+	imageLen := noiseH * noiseW
+
+	if len(refHeights) > 0 {
+		refIDs := PrepareImageIDs(refHeights, refWidths, scale)
+		allIDs = mlx.Concatenate([]*mlx.Array{textIDs, noiseIDs, refIDs}, 0)
+		for i := range refHeights {
+			imageLen += refHeights[i] * refWidths[i]
+		}
+	} else {
+		allIDs = mlx.Concatenate([]*mlx.Array{textIDs, noiseIDs}, 0)
+	}
+
+	cos, sin := ComputeRoPE(allIDs, axesDims, theta)
+	cos = mlx.ToBFloat16(cos)
+	sin = mlx.ToBFloat16(sin)
+
+	return &RoPECache{Cos: cos, Sin: sin, TextLen: textLen, ImageLen: imageLen}
+}

x/imagegen/models/flux2/scheduler.go

@@ -0,0 +1,147 @@
+package flux2
+
+import (
+	"math"
+
+	"github.com/ollama/ollama/x/imagegen/mlx"
+)
+
+// SchedulerConfig holds Flow-Match scheduler configuration
+type SchedulerConfig struct {
+	NumTrainTimesteps  int32   `json:"num_train_timesteps"`  // 1000
+	Shift              float32 `json:"shift"`                // 3.0 for Klein
+	UseDynamicShifting bool    `json:"use_dynamic_shifting"` // true
+	TimeShiftType      string  `json:"time_shift_type"`      // "exponential" or "linear"
+}
+
+// DefaultSchedulerConfig returns default config for Klein
+func DefaultSchedulerConfig() *SchedulerConfig {
+	return &SchedulerConfig{
+		NumTrainTimesteps:  1000,
+		Shift:              3.0, // Klein uses 3.0
+		UseDynamicShifting: true,
+		TimeShiftType:      "exponential",
+	}
+}
+
+// FlowMatchScheduler implements the Flow-Match Euler discrete scheduler
+type FlowMatchScheduler struct {
+	Config    *SchedulerConfig
+	Timesteps []float32 // Discretized timesteps (t from 1 to 0)
+	Sigmas    []float32 // Noise levels at each timestep
+	NumSteps  int       // Number of inference steps
+}
+
+// NewFlowMatchScheduler creates a new scheduler
+func NewFlowMatchScheduler(cfg *SchedulerConfig) *FlowMatchScheduler {
+	return &FlowMatchScheduler{
+		Config: cfg,
+	}
+}
+
+// SetTimesteps sets up the scheduler for the given number of inference steps
+func (s *FlowMatchScheduler) SetTimesteps(numSteps int) {
+	s.SetTimestepsWithMu(numSteps, 0)
+}
+
+// SetTimestepsWithMu sets up scheduler matching diffusers set_timesteps(sigmas=..., mu=...)
+func (s *FlowMatchScheduler) SetTimestepsWithMu(numSteps int, mu float32) {
+	s.NumSteps = numSteps
+
+	// diffusers: sigmas = linspace(1, 1/num_steps, num_steps)
+	// Then applies time shift, appends 0.0 at end
+	s.Sigmas = make([]float32, numSteps+1)
+
+	for i := 0; i < numSteps; i++ {
+		// linspace(1, 1/num_steps, num_steps)
+		var sigma float32
+		if numSteps == 1 {
+			sigma = 1.0
+		} else {
+			sigma = 1.0 - float32(i)/float32(numSteps-1)*(1.0-1.0/float32(numSteps))
+		}
+
+		// Apply time shift if using dynamic shifting
+		if s.Config.UseDynamicShifting && mu != 0 {
+			sigma = s.timeShift(mu, sigma)
+		} else {
+			// If not dynamic shifting, apply fixed shift scaling like diffusers
+			shift := s.Config.Shift
+			sigma = shift * sigma / (1 + (shift-1)*sigma)
+		}
+		s.Sigmas[i] = sigma
+	}
+	// Append terminal zero
+	s.Sigmas[numSteps] = 0.0
+
+	// Timesteps scaled to training range (matches diffusers: timesteps = sigmas * num_train_timesteps)
+	s.Timesteps = make([]float32, numSteps+1)
+	for i, v := range s.Sigmas {
+		s.Timesteps[i] = v * float32(s.Config.NumTrainTimesteps)
+	}
+}
+
+// timeShift applies the dynamic time shift
+func (s *FlowMatchScheduler) timeShift(mu float32, t float32) float32 {
+	if t <= 0 {
+		return 0
+	}
+	if s.Config.TimeShiftType == "linear" {
+		return mu / (mu + (1.0/t-1.0))
+	}
+	// Default: exponential
+	expMu := float32(math.Exp(float64(mu)))
+	return expMu / (expMu + (1.0/t - 1.0))
+}
+
+// Step performs one denoising step
+func (s *FlowMatchScheduler) Step(modelOutput, sample *mlx.Array, timestepIdx int) *mlx.Array {
+	sigma := s.Sigmas[timestepIdx]
+	sigmaNext := s.Sigmas[timestepIdx+1]
+
+	// Euler step: x_{t-dt} = x_t + (sigma_next - sigma) * v_t
+	dt := sigmaNext - sigma
+
+	// Upcast to float32 for precision (matches diffusers)
+	sampleF32 := mlx.AsType(sample, mlx.DtypeFloat32)
+	outputF32 := mlx.AsType(modelOutput, mlx.DtypeFloat32)
+
+	scaledOutput := mlx.MulScalar(outputF32, dt)
+	result := mlx.Add(sampleF32, scaledOutput)
+
+	// Cast back to bfloat16
+	return mlx.ToBFloat16(result)
+}
+
+// GetTimestep returns the timestep value at the given index
+func (s *FlowMatchScheduler) GetTimestep(idx int) float32 {
+	if idx < len(s.Timesteps) {
+		return s.Timesteps[idx]
+	}
+	return 0.0
+}
+
+// InitNoise creates initial noise for sampling
+func (s *FlowMatchScheduler) InitNoise(shape []int32, seed int64) *mlx.Array {
+	return mlx.RandomNormalWithDtype(shape, uint64(seed), mlx.DtypeBFloat16)
+}
+
+// CalculateShift computes the mu shift value for dynamic scheduling
+// Matches diffusers compute_empirical_mu function
+func CalculateShift(imgSeqLen int32, numSteps int) float32 {
+	a1, b1 := float32(8.73809524e-05), float32(1.89833333)
+	a2, b2 := float32(0.00016927), float32(0.45666666)
+
+	seqLen := float32(imgSeqLen)
+
+	if imgSeqLen > 4300 {
+		return a2*seqLen + b2
+	}
+
+	m200 := a2*seqLen + b2
+	m10 := a1*seqLen + b1
+
+	a := (m200 - m10) / 190.0
+	b := m200 - 200.0*a
+	return a*float32(numSteps) + b
+}

x/imagegen/models/flux2/transformer.go

@@ -0,0 +1,560 @@
+package flux2
+
+import (
+	"fmt"
+	"math"
+
+	"github.com/ollama/ollama/x/imagegen/manifest"
+	"github.com/ollama/ollama/x/imagegen/mlx"
+	"github.com/ollama/ollama/x/imagegen/nn"
+	"github.com/ollama/ollama/x/imagegen/safetensors"
+)
+
+// TransformerConfig holds Flux2 transformer configuration
+type TransformerConfig struct {
+	AttentionHeadDim         int32   `json:"attention_head_dim"`         // 128
+	AxesDimsRoPE             []int32 `json:"axes_dims_rope"`             // [32, 32, 32, 32]
+	Eps                      float32 `json:"eps"`                        // 1e-6
+	GuidanceEmbeds           bool    `json:"guidance_embeds"`            // false for Klein
+	InChannels               int32   `json:"in_channels"`                // 128
+	JointAttentionDim        int32   `json:"joint_attention_dim"`        // 7680
+	MLPRatio                 float32 `json:"mlp_ratio"`                  // 3.0
+	NumAttentionHeads        int32   `json:"num_attention_heads"`        // 24
+	NumLayers                int32   `json:"num_layers"`                 // 5
+	NumSingleLayers          int32   `json:"num_single_layers"`          // 20
+	PatchSize                int32   `json:"patch_size"`                 // 1
+	RopeTheta                int32   `json:"rope_theta"`                 // 2000
+	TimestepGuidanceChannels int32   `json:"timestep_guidance_channels"` // 256
+}
+
+// Computed dimensions
+func (c *TransformerConfig) InnerDim() int32 {
+	return c.NumAttentionHeads * c.AttentionHeadDim // 24 * 128 = 3072
+}
+
+func (c *TransformerConfig) MLPHiddenDim() int32 {
+	return int32(float32(c.InnerDim()) * c.MLPRatio) // 3072 * 3.0 = 9216
+}
+
+// TimestepEmbedder creates timestep embeddings
+// Weight names: time_guidance_embed.timestep_embedder.linear_1.weight, linear_2.weight
+type TimestepEmbedder struct {
+	Linear1  nn.LinearLayer `weight:"linear_1"`
+	Linear2  nn.LinearLayer `weight:"linear_2"`
+	EmbedDim int32          // 256
+}
+
+// Forward creates sinusoidal embeddings and projects them
+func (t *TimestepEmbedder) Forward(timesteps *mlx.Array) *mlx.Array {
+	half := t.EmbedDim / 2
+	freqs := make([]float32, half)
+	for i := int32(0); i < half; i++ {
+		freqs[i] = float32(math.Exp(-math.Log(10000.0) * float64(i) / float64(half)))
+	}
+	freqsArr := mlx.NewArray(freqs, []int32{1, half})
+
+	// timesteps: [B] -> [B, 1]
+	tExpanded := mlx.ExpandDims(timesteps, 1)
+	// args: [B, half]
+	args := mlx.Mul(tExpanded, freqsArr)
+
+	// [cos(args), sin(args)] -> [B, embed_dim]
+	sinEmbed := mlx.Concatenate([]*mlx.Array{mlx.Cos(args), mlx.Sin(args)}, 1)
+
+	// MLP: linear_1 -> silu -> linear_2
+	h := t.Linear1.Forward(sinEmbed)
+	h = mlx.SiLU(h)
+	return t.Linear2.Forward(h)
+}
+
+// TimeGuidanceEmbed wraps the timestep embedder
+// Weight names: time_guidance_embed.timestep_embedder.*
+type TimeGuidanceEmbed struct {
+	TimestepEmbedder *TimestepEmbedder `weight:"timestep_embedder"`
+}
+
+// Forward computes timestep embeddings
+func (t *TimeGuidanceEmbed) Forward(timesteps *mlx.Array) *mlx.Array {
+	return t.TimestepEmbedder.Forward(timesteps)
+}
+
+// Modulation computes adaptive modulation parameters
+// Weight names: double_stream_modulation_img.linear.weight, etc.
+type Modulation struct {
+	Linear nn.LinearLayer `weight:"linear"`
+}
+
+// Forward computes modulation parameters
+func (m *Modulation) Forward(temb *mlx.Array) *mlx.Array {
+	h := mlx.SiLU(temb)
+	return m.Linear.Forward(h)
+}
+
+// TransformerBlockAttn implements dual-stream attention
+// Weight names: transformer_blocks.N.attn.*
+type TransformerBlockAttn struct {
+	// Image stream (separate Q, K, V projections)
+	ToQ nn.LinearLayer `weight:"to_q"`
+	ToK nn.LinearLayer `weight:"to_k"`
+	ToV nn.LinearLayer `weight:"to_v"`
+	// Note: to_out has .0 suffix in weights, handled specially
+	ToOut0 nn.LinearLayer `weight:"to_out.0"`
+
+	// Text stream (add_ projections)
+	AddQProj nn.LinearLayer `weight:"add_q_proj"`
+	AddKProj nn.LinearLayer `weight:"add_k_proj"`
+	AddVProj nn.LinearLayer `weight:"add_v_proj"`
+	ToAddOut nn.LinearLayer `weight:"to_add_out"`
+
+	// QK norms for image stream
+	NormQ *mlx.Array `weight:"norm_q.weight"`
+	NormK *mlx.Array `weight:"norm_k.weight"`
+
+	// QK norms for text stream (added)
+	NormAddedQ *mlx.Array `weight:"norm_added_q.weight"`
+	NormAddedK *mlx.Array `weight:"norm_added_k.weight"`
+}
+
+// FeedForward implements SwiGLU MLP
+// Weight names: transformer_blocks.N.ff.linear_in.weight, linear_out.weight
+type FeedForward struct {
+	LinearIn  nn.LinearLayer `weight:"linear_in"`
+	LinearOut nn.LinearLayer `weight:"linear_out"`
+}
+
+// Forward applies SwiGLU MLP
+func (ff *FeedForward) Forward(x *mlx.Array) *mlx.Array {
+	// LinearIn outputs 2x hidden dim for SwiGLU
+	h := ff.LinearIn.Forward(x)
+	shape := h.Shape()
+	half := shape[len(shape)-1] / 2
+
+	// Split into gate and up
+	gate := mlx.Slice(h, []int32{0, 0, 0}, []int32{shape[0], shape[1], half})
+	up := mlx.Slice(h, []int32{0, 0, half}, []int32{shape[0], shape[1], shape[2]})
+
+	// SwiGLU: silu(gate) * up
+	h = mlx.Mul(mlx.SiLU(gate), up)
+	return ff.LinearOut.Forward(h)
+}
+
+// TransformerBlock implements a dual-stream transformer block
+// Weight names: transformer_blocks.N.*
+type TransformerBlock struct {
+	Attn      *TransformerBlockAttn `weight:"attn"`
+	FF        *FeedForward          `weight:"ff"`
+	FFContext *FeedForward          `weight:"ff_context"`
+
+	// Config (set after loading)
+	NHeads  int32
+	HeadDim int32
+	Scale   float32
+}
+
+// Forward applies the dual-stream block
+// imgHidden: [B, imgLen, dim]
+// txtHidden: [B, txtLen, dim]
+// imgMod, txtMod: modulation params [B, 6*dim] each
+// cos, sin: RoPE values
+func (block *TransformerBlock) Forward(imgHidden, txtHidden *mlx.Array, imgMod, txtMod *mlx.Array, cos, sin *mlx.Array) (*mlx.Array, *mlx.Array) {
+	imgShape := imgHidden.Shape()
+	B := imgShape[0]
+	imgLen := imgShape[1]
+	dim := imgShape[2]
+	txtLen := txtHidden.Shape()[1]
+
+	// Parse modulation: 6 params each (shift1, scale1, gate1, shift2, scale2, gate2)
+	imgShift1, imgScale1, imgGate1 := parseModulation3(imgMod, dim, 0)
+	imgShift2, imgScale2, imgGate2 := parseModulation3(imgMod, dim, 3)
+	txtShift1, txtScale1, txtGate1 := parseModulation3(txtMod, dim, 0)
+	txtShift2, txtScale2, txtGate2 := parseModulation3(txtMod, dim, 3)
+
+	// === Attention branch ===
+	// Modulate inputs
+	imgNorm := modulateLayerNorm(imgHidden, imgShift1, imgScale1)
+	txtNorm := modulateLayerNorm(txtHidden, txtShift1, txtScale1)
+
+	// Compute Q, K, V for image stream (separate projections)
+	imgQ := block.Attn.ToQ.Forward(imgNorm)
+	imgK := block.Attn.ToK.Forward(imgNorm)
+	imgV := block.Attn.ToV.Forward(imgNorm)
+
+	// Compute Q, K, V for text stream (add_ projections)
+	txtQ := block.Attn.AddQProj.Forward(txtNorm)
+	txtK := block.Attn.AddKProj.Forward(txtNorm)
+	txtV := block.Attn.AddVProj.Forward(txtNorm)
+
+	// Reshape for attention: [B, L, dim] -> [B, L, nheads, headDim]
+	imgQ = mlx.Reshape(imgQ, B, imgLen, block.NHeads, block.HeadDim)
+	imgK = mlx.Reshape(imgK, B, imgLen, block.NHeads, block.HeadDim)
+	imgV = mlx.Reshape(imgV, B, imgLen, block.NHeads, block.HeadDim)
+	txtQ = mlx.Reshape(txtQ, B, txtLen, block.NHeads, block.HeadDim)
+	txtK = mlx.Reshape(txtK, B, txtLen, block.NHeads, block.HeadDim)
+	txtV = mlx.Reshape(txtV, B, txtLen, block.NHeads, block.HeadDim)
+
+	// Apply QK norm (RMSNorm with learned scale)
+	imgQ = applyQKNorm(imgQ, block.Attn.NormQ)
+	imgK = applyQKNorm(imgK, block.Attn.NormK)
+	txtQ = applyQKNorm(txtQ, block.Attn.NormAddedQ)
+	txtK = applyQKNorm(txtK, block.Attn.NormAddedK)
+
+	// Concatenate for joint attention: text first, then image
+	q := mlx.Concatenate([]*mlx.Array{txtQ, imgQ}, 1)
+	k := mlx.Concatenate([]*mlx.Array{txtK, imgK}, 1)
+	v := mlx.Concatenate([]*mlx.Array{txtV, imgV}, 1)
+
+	// Apply RoPE
+	q = ApplyRoPE4D(q, cos, sin)
+	k = ApplyRoPE4D(k, cos, sin)
+
+	// Transpose for SDPA: [B, nheads, L, headDim]
+	q = mlx.Transpose(q, 0, 2, 1, 3)
+	k = mlx.Transpose(k, 0, 2, 1, 3)
+	v = mlx.Transpose(v, 0, 2, 1, 3)
+
+	// Scaled dot-product attention
+	out := mlx.ScaledDotProductAttention(q, k, v, block.Scale, false)
+
+	// Transpose back: [B, L, nheads, headDim]
+	out = mlx.Transpose(out, 0, 2, 1, 3)
+
+	// Split back into txt and img
+	totalLen := txtLen + imgLen
+	txtOut := mlx.Slice(out, []int32{0, 0, 0, 0}, []int32{B, txtLen, block.NHeads, block.HeadDim})
+	imgOut := mlx.Slice(out, []int32{0, txtLen, 0, 0}, []int32{B, totalLen, block.NHeads, block.HeadDim})
+
+	// Reshape and project
+	txtOut = mlx.Reshape(txtOut, B, txtLen, dim)
+	imgOut = mlx.Reshape(imgOut, B, imgLen, dim)
+	txtOut = block.Attn.ToAddOut.Forward(txtOut)
+	imgOut = block.Attn.ToOut0.Forward(imgOut)
+
+	// Apply gates and residual
+	imgHidden = mlx.Add(imgHidden, mlx.Mul(imgGate1, imgOut))
+	txtHidden = mlx.Add(txtHidden, mlx.Mul(txtGate1, txtOut))
+
+	// === MLP branch ===
+	imgNorm = modulateLayerNorm(imgHidden, imgShift2, imgScale2)
+	txtNorm = modulateLayerNorm(txtHidden, txtShift2, txtScale2)
+
+	imgFFOut := block.FF.Forward(imgNorm)
+	txtFFOut := block.FFContext.Forward(txtNorm)
+
+	imgHidden = mlx.Add(imgHidden, mlx.Mul(imgGate2, imgFFOut))
+	txtHidden = mlx.Add(txtHidden, mlx.Mul(txtGate2, txtFFOut))
+
+	return imgHidden, txtHidden
+}
+
+// SingleTransformerBlockAttn implements attention for single-stream blocks
+// Weight names: single_transformer_blocks.N.attn.*
+type SingleTransformerBlockAttn struct {
+	ToQKVMlpProj nn.LinearLayer `weight:"to_qkv_mlp_proj"` // Fused QKV + MLP input
+	ToOut        nn.LinearLayer `weight:"to_out"`          // Fused attn_out + MLP out
+	NormQ        *mlx.Array     `weight:"norm_q.weight"`
+	NormK        *mlx.Array     `weight:"norm_k.weight"`
+}
+
+// SingleTransformerBlock implements a single-stream transformer block
+// Weight names: single_transformer_blocks.N.*
+type SingleTransformerBlock struct {
+	Attn *SingleTransformerBlockAttn `weight:"attn"`
+
+	// Config
+	NHeads    int32
+	HeadDim   int32
+	InnerDim  int32
+	MLPHidDim int32
+	Scale     float32
+}
+
+// Forward applies the single-stream block
+// x: [B, L, dim] concatenated text+image
+// mod: modulation [B, 3*dim]
+func (block *SingleTransformerBlock) Forward(x *mlx.Array, mod *mlx.Array, cos, sin *mlx.Array) *mlx.Array {
+	shape := x.Shape()
+	B := shape[0]
+	L := shape[1]
+	dim := shape[2]
+
+	// Parse modulation: (shift, scale, gate)
+	shift, scale, gate := parseModulation3(mod, dim, 0)
+
+	// Modulate input
+	h := modulateLayerNorm(x, shift, scale)
+
+	// Fused projection: QKV + MLP gate/up
+	// linear1 outputs: [q, k, v, mlp_gate, mlp_up] = [dim, dim, dim, mlpHid, mlpHid]
+	qkvMlp := block.Attn.ToQKVMlpProj.Forward(h)
+
+	// Split: first 3*dim is QKV, rest is MLP
+	qkvDim := 3 * block.InnerDim
+	qkv := mlx.Slice(qkvMlp, []int32{0, 0, 0}, []int32{B, L, qkvDim})
+	mlpIn := mlx.Slice(qkvMlp, []int32{0, 0, qkvDim}, []int32{B, L, qkvMlp.Shape()[2]})
+
+	// Split QKV
+	q, k, v := splitQKV(qkv, B, L, block.InnerDim)
+
+	// Reshape for attention
+	q = mlx.Reshape(q, B, L, block.NHeads, block.HeadDim)
+	k = mlx.Reshape(k, B, L, block.NHeads, block.HeadDim)
+	v = mlx.Reshape(v, B, L, block.NHeads, block.HeadDim)
+
+	// QK norm
+	q = applyQKNorm(q, block.Attn.NormQ)
+	k = applyQKNorm(k, block.Attn.NormK)
+
+	// Apply RoPE
+	q = ApplyRoPE4D(q, cos, sin)
+	k = ApplyRoPE4D(k, cos, sin)
+
+	// Transpose for SDPA
+	q = mlx.Transpose(q, 0, 2, 1, 3)
+	k = mlx.Transpose(k, 0, 2, 1, 3)
+	v = mlx.Transpose(v, 0, 2, 1, 3)
+
+	// SDPA
+	attnOut := mlx.ScaledDotProductAttention(q, k, v, block.Scale, false)
+
+	// Transpose back and reshape
+	attnOut = mlx.Transpose(attnOut, 0, 2, 1, 3)
+	attnOut = mlx.Reshape(attnOut, B, L, block.InnerDim)
+
+	// MLP: SwiGLU
+	mlpShape := mlpIn.Shape()
+	half := mlpShape[2] / 2
+	mlpGate := mlx.Slice(mlpIn, []int32{0, 0, 0}, []int32{B, L, half})
+	mlpUp := mlx.Slice(mlpIn, []int32{0, 0, half}, []int32{B, L, mlpShape[2]})
+	mlpOut := mlx.Mul(mlx.SiLU(mlpGate), mlpUp)
+
+	// Concatenate attention and MLP for fused output
+	combined := mlx.Concatenate([]*mlx.Array{attnOut, mlpOut}, 2)
+
+	// Output projection
+	out := block.Attn.ToOut.Forward(combined)
+
+	// Apply gate and residual
+	return mlx.Add(x, mlx.Mul(gate, out))
+}
+
+// NormOut implements the output normalization with modulation
+// Weight names: norm_out.linear.weight
+type NormOut struct {
+	Linear nn.LinearLayer `weight:"linear"`
+}
+
+// Forward computes final modulated output
+func (n *NormOut) Forward(x *mlx.Array, temb *mlx.Array) *mlx.Array {
+	shape := x.Shape()
+	B := shape[0]
+	dim := shape[2]
+
+	// Modulation: temb -> silu -> linear -> [shift, scale]
+	mod := mlx.SiLU(temb)
+	mod = n.Linear.Forward(mod)
+
+	// Split into scale and shift (diffusers order: scale first, shift second)
+	scale := mlx.Slice(mod, []int32{0, 0}, []int32{B, dim})
+	shift := mlx.Slice(mod, []int32{0, dim}, []int32{B, 2 * dim})
+	shift = mlx.ExpandDims(shift, 1)
+	scale = mlx.ExpandDims(scale, 1)
+
+	// Modulate with RMSNorm
+	return modulateLayerNorm(x, shift, scale)
+}
+
+// Flux2Transformer2DModel is the main Flux2 transformer
+// Weight names at top level: time_guidance_embed.*, double_stream_modulation_*.*, etc.
+type Flux2Transformer2DModel struct {
+	// Timestep embedding
+	TimeGuidanceEmbed *TimeGuidanceEmbed `weight:"time_guidance_embed"`
+
+	// Shared modulation
+	DoubleStreamModulationImg *Modulation `weight:"double_stream_modulation_img"`
+	DoubleStreamModulationTxt *Modulation `weight:"double_stream_modulation_txt"`
+	SingleStreamModulation    *Modulation `weight:"single_stream_modulation"`
+
+	// Embedders
+	XEmbedder       nn.LinearLayer `weight:"x_embedder"`
+	ContextEmbedder nn.LinearLayer `weight:"context_embedder"`
+
+	// Transformer blocks
+	TransformerBlocks       []*TransformerBlock       `weight:"transformer_blocks"`
+	SingleTransformerBlocks []*SingleTransformerBlock `weight:"single_transformer_blocks"`
+
+	// Output
+	NormOut *NormOut       `weight:"norm_out"`
+	ProjOut nn.LinearLayer `weight:"proj_out"`
+
+	*TransformerConfig
+}
+
+// Load loads the Flux2 transformer from ollama blob storage.
+func (m *Flux2Transformer2DModel) Load(modelManifest *manifest.ModelManifest) error {
+	fmt.Print("  Loading transformer... ")
+
+	// Load config from blob
+	var cfg TransformerConfig
+	if err := modelManifest.ReadConfigJSON("transformer/config.json", &cfg); err != nil {
+		return fmt.Errorf("config: %w", err)
+	}
+	m.TransformerConfig = &cfg
+
+	// Initialize slices
+	m.TransformerBlocks = make([]*TransformerBlock, cfg.NumLayers)
+	m.SingleTransformerBlocks = make([]*SingleTransformerBlock, cfg.NumSingleLayers)
+
+	// Initialize TimeGuidanceEmbed with embed dim
+	m.TimeGuidanceEmbed = &TimeGuidanceEmbed{
+		TimestepEmbedder: &TimestepEmbedder{EmbedDim: cfg.TimestepGuidanceChannels},
+	}
+
+	// Load weights from tensor blobs
+	weights, err := manifest.LoadWeightsFromManifest(modelManifest, "transformer")
+	if err != nil {
+		return fmt.Errorf("weights: %w", err)
+	}
+	if err := weights.Load(0); err != nil {
+		return fmt.Errorf("load weights: %w", err)
+	}
+	defer weights.ReleaseAll()
+
+	return m.loadWeights(weights)
+}
+
+// loadWeights loads weights from any WeightSource into the model
+func (m *Flux2Transformer2DModel) loadWeights(weights safetensors.WeightSource) error {
+	if err := safetensors.LoadModule(m, weights, ""); err != nil {
+		return fmt.Errorf("load module: %w", err)
+	}
+	m.initComputedFields()
+	fmt.Println("✓")
+	return nil
+}
+
+// initComputedFields initializes computed fields after loading weights
+func (m *Flux2Transformer2DModel) initComputedFields() {
+	cfg := m.TransformerConfig
+	innerDim := cfg.InnerDim()
+	scale := float32(1.0 / math.Sqrt(float64(cfg.AttentionHeadDim)))
+
+	// Initialize transformer blocks
+	for _, block := range m.TransformerBlocks {
+		block.NHeads = cfg.NumAttentionHeads
+		block.HeadDim = cfg.AttentionHeadDim
+		block.Scale = scale
+	}
+
+	// Initialize single transformer blocks
+	for _, block := range m.SingleTransformerBlocks {
+		block.NHeads = cfg.NumAttentionHeads
+		block.HeadDim = cfg.AttentionHeadDim
+		block.InnerDim = innerDim
+		block.MLPHidDim = cfg.MLPHiddenDim()
+		block.Scale = scale
+	}
+}
+
+// Forward runs the Flux2 transformer
+func (m *Flux2Transformer2DModel) Forward(patches, txtEmbeds *mlx.Array, timesteps *mlx.Array, rope *RoPECache) *mlx.Array {
+	patchShape := patches.Shape()
+	B := patchShape[0]
+	imgLen := patchShape[1]
+	txtLen := txtEmbeds.Shape()[1]
+
+	// Scale timestep to 0-1000 range (diffusers multiplies by 1000)
+	scaledTimesteps := mlx.MulScalar(timesteps, 1000.0)
+
+	// Compute timestep embedding
+	temb := m.TimeGuidanceEmbed.Forward(scaledTimesteps)
+
+	// Embed patches and text
+	imgHidden := m.XEmbedder.Forward(patches)
+	txtHidden := m.ContextEmbedder.Forward(txtEmbeds)
+
+	// Compute shared modulation
+	imgMod := m.DoubleStreamModulationImg.Forward(temb)
+	txtMod := m.DoubleStreamModulationTxt.Forward(temb)
+	singleMod := m.SingleStreamModulation.Forward(temb)
+
+	// Double (dual-stream) blocks
+	for _, block := range m.TransformerBlocks {
+		imgHidden, txtHidden = block.Forward(imgHidden, txtHidden, imgMod, txtMod, rope.Cos, rope.Sin)
+	}
+
+	// Concatenate for single-stream: text first, then image
+	hidden := mlx.Concatenate([]*mlx.Array{txtHidden, imgHidden}, 1)
+
+	// Single-stream blocks
+	for _, block := range m.SingleTransformerBlocks {
+		hidden = block.Forward(hidden, singleMod, rope.Cos, rope.Sin)
+	}
+
+	// Extract image portion
+	totalLen := txtLen + imgLen
+	imgOut := mlx.Slice(hidden, []int32{0, txtLen, 0}, []int32{B, totalLen, hidden.Shape()[2]})
+
+	// Final norm and projection
+	imgOut = m.NormOut.Forward(imgOut, temb)
+	return m.ProjOut.Forward(imgOut)
+}
+
+// Note: QK normalization uses mlx.RMSNorm (the fast version) directly
+// See applyQKNorm function below
+
+// compiledSwiGLU fuses: silu(gate) * up
+// Called 30x per step (10 in dual-stream + 20 in single-stream blocks)
+var compiledSwiGLU *mlx.CompiledFunc
+
+func getCompiledSwiGLU() *mlx.CompiledFunc {
+	if compiledSwiGLU == nil {
+		compiledSwiGLU = mlx.CompileShapeless(func(inputs []*mlx.Array) []*mlx.Array {
+			gate, up := inputs[0], inputs[1]
+			return []*mlx.Array{mlx.Mul(mlx.SiLU(gate), up)}
+		}, true)
+	}
+	return compiledSwiGLU
+}
+
+// Helper functions
+
+// parseModulation3 extracts 3 modulation params (shift, scale, gate) starting at offset
+func parseModulation3(mod *mlx.Array, dim int32, offset int32) (*mlx.Array, *mlx.Array, *mlx.Array) {
+	B := mod.Shape()[0]
+	start := offset * dim
+	shift := mlx.Slice(mod, []int32{0, start}, []int32{B, start + dim})
+	scale := mlx.Slice(mod, []int32{0, start + dim}, []int32{B, start + 2*dim})
+	gate := mlx.Slice(mod, []int32{0, start + 2*dim}, []int32{B, start + 3*dim})
+
+	// Expand for broadcasting [B, dim] -> [B, 1, dim]
+	shift = mlx.ExpandDims(shift, 1)
+	scale = mlx.ExpandDims(scale, 1)
+	gate = mlx.ExpandDims(gate, 1)
+
+	return shift, scale, gate
+}
+
+// modulateLayerNorm applies LayerNorm then shift/scale modulation
+// Diffusers uses LayerNorm(elementwise_affine=False) which centers the data
+func modulateLayerNorm(x *mlx.Array, shift, scale *mlx.Array) *mlx.Array {
+	// Fast LayerNorm without learnable params
+	x = mlx.LayerNorm(x, 1e-6)
+
+	// Modulate: x * (1 + scale) + shift
+	x = mlx.Mul(x, mlx.AddScalar(scale, 1.0))
+	return mlx.Add(x, shift)
+}
+
+// splitQKV splits a fused QKV tensor into Q, K, V
+func splitQKV(qkv *mlx.Array, B, L, dim int32) (*mlx.Array, *mlx.Array, *mlx.Array) {
+	q := mlx.Slice(qkv, []int32{0, 0, 0}, []int32{B, L, dim})
+	k := mlx.Slice(qkv, []int32{0, 0, dim}, []int32{B, L, 2 * dim})
+	v := mlx.Slice(qkv, []int32{0, 0, 2 * dim}, []int32{B, L, 3 * dim})
+	return q, k, v
+}
+
+// applyQKNorm applies RMSNorm with learned scale (no bias)
+// Uses the optimized mlx_fast_rms_norm
+func applyQKNorm(x *mlx.Array, scale *mlx.Array) *mlx.Array {
+	return mlx.RMSNorm(x, scale, 1e-6)
+}

x/imagegen/models/flux2/vae.go

@@ -0,0 +1,802 @@
+package flux2
+
+import (
+	"fmt"
+	"math"
+
+	"github.com/ollama/ollama/x/imagegen/manifest"
+	"github.com/ollama/ollama/x/imagegen/mlx"
+	"github.com/ollama/ollama/x/imagegen/nn"
+	"github.com/ollama/ollama/x/imagegen/safetensors"
+	"github.com/ollama/ollama/x/imagegen/vae"
+)
+
+// VAEConfig holds AutoencoderKLFlux2 configuration
+type VAEConfig struct {
+	ActFn             string  `json:"act_fn"`                  // "silu"
+	BatchNormEps      float32 `json:"batch_norm_eps"`          // 0.0001
+	BatchNormMomentum float32 `json:"batch_norm_momentum"`     // 0.1
+	BlockOutChannels  []int32 `json:"block_out_channels"`      // [128, 256, 512, 512]
+	ForceUpcast       bool    `json:"force_upcast"`            // true
+	InChannels        int32   `json:"in_channels"`             // 3
+	LatentChannels    int32   `json:"latent_channels"`         // 32
+	LayersPerBlock    int32   `json:"layers_per_block"`        // 2
+	MidBlockAddAttn   bool    `json:"mid_block_add_attention"` // true
+	NormNumGroups     int32   `json:"norm_num_groups"`         // 32
+	OutChannels       int32   `json:"out_channels"`            // 3
+	PatchSize         []int32 `json:"patch_size"`              // [2, 2]
+	SampleSize        int32   `json:"sample_size"`             // 1024
+	UsePostQuantConv  bool    `json:"use_post_quant_conv"`     // true
+	UseQuantConv      bool    `json:"use_quant_conv"`          // true
+}
+
+// BatchNorm2D implements 2D batch normalization with running statistics
+type BatchNorm2D struct {
+	RunningMean *mlx.Array // [C]
+	RunningVar  *mlx.Array // [C]
+	Weight      *mlx.Array // [C] gamma
+	Bias        *mlx.Array // [C] beta
+	Eps         float32
+	Momentum    float32
+}
+
+// Forward applies batch normalization (inference mode - uses running stats)
+// Input and output are in NHWC format [B, H, W, C]
+func (bn *BatchNorm2D) Forward(x *mlx.Array) *mlx.Array {
+	shape := x.Shape()
+	C := shape[3]
+
+	// Reshape stats for broadcasting [1, 1, 1, C]
+	mean := mlx.Reshape(bn.RunningMean, 1, 1, 1, C)
+	variance := mlx.Reshape(bn.RunningVar, 1, 1, 1, C)
+
+	// Normalize: (x - mean) / sqrt(var + eps)
+	xNorm := mlx.Sub(x, mean)
+	xNorm = mlx.Div(xNorm, mlx.Sqrt(mlx.AddScalar(variance, bn.Eps)))
+
+	// Scale and shift (only if affine=True)
+	if bn.Weight != nil {
+		weight := mlx.Reshape(bn.Weight, 1, 1, 1, C)
+		xNorm = mlx.Mul(xNorm, weight)
+	}
+	if bn.Bias != nil {
+		bias := mlx.Reshape(bn.Bias, 1, 1, 1, C)
+		xNorm = mlx.Add(xNorm, bias)
+	}
+
+	return xNorm
+}
+
+// Denormalize inverts the batch normalization
+// Used when decoding latents
+func (bn *BatchNorm2D) Denormalize(x *mlx.Array) *mlx.Array {
+	shape := x.Shape()
+	C := shape[3]
+
+	// Reshape stats for broadcasting [1, 1, 1, C]
+	mean := mlx.Reshape(bn.RunningMean, 1, 1, 1, C)
+	variance := mlx.Reshape(bn.RunningVar, 1, 1, 1, C)
+
+	// Inverse: first undo affine, then undo normalization
+	// For affine=False: x_denorm = x * sqrt(var + eps) + mean
+	if bn.Bias != nil {
+		bias := mlx.Reshape(bn.Bias, 1, 1, 1, C)
+		x = mlx.Sub(x, bias)
+	}
+	if bn.Weight != nil {
+		weight := mlx.Reshape(bn.Weight, 1, 1, 1, C)
+		x = mlx.Div(x, weight)
+	}
+	x = mlx.Mul(x, mlx.Sqrt(mlx.AddScalar(variance, bn.Eps)))
+	x = mlx.Add(x, mean)
+
+	return x
+}
+
+// GroupNormLayer implements group normalization
+// Reused from zimage package pattern
+type GroupNormLayer struct {
+	Weight    *mlx.Array `weight:"weight"`
+	Bias      *mlx.Array `weight:"bias"`
+	NumGroups int32
+	Eps       float32
+}
+
+// Forward applies group normalization
+// Input and output are in NHWC format [B, H, W, C]
+func (gn *GroupNormLayer) Forward(x *mlx.Array) *mlx.Array {
+	shape := x.Shape()
+	B := shape[0]
+	H := shape[1]
+	W := shape[2]
+	C := shape[3]
+
+	// Reshape to [B, H, W, groups, C/groups]
+	groupSize := C / gn.NumGroups
+	x = mlx.Reshape(x, B, H, W, gn.NumGroups, groupSize)
+
+	// Compute mean and variance per group
+	mean := mlx.Mean(x, 1, true)
+	mean = mlx.Mean(mean, 2, true)
+	mean = mlx.Mean(mean, 4, true)
+
+	xCentered := mlx.Sub(x, mean)
+
+	sq := mlx.Square(xCentered)
+	variance := mlx.Mean(sq, 1, true)
+	variance = mlx.Mean(variance, 2, true)
+	variance = mlx.Mean(variance, 4, true)
+
+	// Normalize
+	xNorm := mlx.Div(xCentered, mlx.Sqrt(mlx.AddScalar(variance, gn.Eps)))
+
+	// Reshape back to [B, H, W, C]
+	xNorm = mlx.Reshape(xNorm, B, H, W, C)
+
+	// Scale and shift
+	if gn.Weight != nil {
+		weight := mlx.Reshape(gn.Weight, 1, 1, 1, C)
+		xNorm = mlx.Mul(xNorm, weight)
+	}
+	if gn.Bias != nil {
+		bias := mlx.Reshape(gn.Bias, 1, 1, 1, C)
+		xNorm = mlx.Add(xNorm, bias)
+	}
+
+	return xNorm
+}
+
+// Conv2D represents a 2D convolution layer (reused pattern)
+type Conv2D struct {
+	Weight  *mlx.Array `weight:"weight"`
+	Bias    *mlx.Array `weight:"bias,optional"`
+	Stride  int32
+	Padding int32
+}
+
+// Transform implements safetensors.Transformer to transpose weights from PyTorch's OIHW to MLX's OHWI.
+func (conv *Conv2D) Transform(field string, arr *mlx.Array) *mlx.Array {
+	if field == "Weight" {
+		return mlx.Transpose(arr, 0, 2, 3, 1)
+	}
+	return arr
+}
+
+// Forward applies convolution (NHWC format)
+func (conv *Conv2D) Forward(x *mlx.Array) *mlx.Array {
+	out := mlx.Conv2d(x, conv.Weight, conv.Stride, conv.Padding)
+
+	if conv.Bias != nil {
+		bias := mlx.Reshape(conv.Bias, 1, 1, 1, conv.Bias.Dim(0))
+		out = mlx.Add(out, bias)
+	}
+
+	return out
+}
+
+// ResnetBlock2D implements a ResNet block for VAE
+type ResnetBlock2D struct {
+	Norm1        *GroupNormLayer `weight:"norm1"`
+	Conv1        *Conv2D         `weight:"conv1"`
+	Norm2        *GroupNormLayer `weight:"norm2"`
+	Conv2        *Conv2D         `weight:"conv2"`
+	ConvShortcut *Conv2D         `weight:"conv_shortcut,optional"`
+}
+
+// Forward applies the ResNet block
+func (rb *ResnetBlock2D) Forward(x *mlx.Array) *mlx.Array {
+	h := rb.Norm1.Forward(x)
+	h = mlx.SiLU(h)
+	h = rb.Conv1.Forward(h)
+
+	h = rb.Norm2.Forward(h)
+	h = mlx.SiLU(h)
+	h = rb.Conv2.Forward(h)
+
+	if rb.ConvShortcut != nil {
+		x = rb.ConvShortcut.Forward(x)
+	}
+
+	return mlx.Add(h, x)
+}
+
+// VAEAttentionBlock implements self-attention for VAE
+type VAEAttentionBlock struct {
+	GroupNorm *GroupNormLayer `weight:"group_norm"`
+	ToQ       nn.LinearLayer  `weight:"to_q"`
+	ToK       nn.LinearLayer  `weight:"to_k"`
+	ToV       nn.LinearLayer  `weight:"to_v"`
+	ToOut     nn.LinearLayer  `weight:"to_out.0"`
+}
+
+// Forward applies attention (NHWC format)
+func (ab *VAEAttentionBlock) Forward(x *mlx.Array) *mlx.Array {
+	residual := x
+	shape := x.Shape()
+	B := shape[0]
+	H := shape[1]
+	W := shape[2]
+	C := shape[3]
+
+	h := ab.GroupNorm.Forward(x)
+	h = mlx.Reshape(h, B, H*W, C)
+
+	q := ab.ToQ.Forward(h)
+	k := ab.ToK.Forward(h)
+	v := ab.ToV.Forward(h)
+
+	q = mlx.ExpandDims(q, 1)
+	k = mlx.ExpandDims(k, 1)
+	v = mlx.ExpandDims(v, 1)
+
+	scale := float32(1.0 / math.Sqrt(float64(C)))
+	out := mlx.ScaledDotProductAttention(q, k, v, scale, false)
+	out = mlx.Squeeze(out, 1)
+
+	out = ab.ToOut.Forward(out)
+	out = mlx.Reshape(out, B, H, W, C)
+	out = mlx.Add(out, residual)
+
+	return out
+}
+
+// UpDecoderBlock2D implements an upsampling decoder block
+type UpDecoderBlock2D struct {
+	ResnetBlocks []*ResnetBlock2D
+	Upsample     *Conv2D
+}
+
+// Forward applies the up decoder block
+func (ub *UpDecoderBlock2D) Forward(x *mlx.Array) *mlx.Array {
+	for _, resnet := range ub.ResnetBlocks {
+		x = resnet.Forward(x)
+	}
+
+	if ub.Upsample != nil {
+		x = upsample2x(x)
+		x = ub.Upsample.Forward(x)
+	}
+
+	return x
+}
+
+// upsample2x performs 2x nearest neighbor upsampling
+func upsample2x(x *mlx.Array) *mlx.Array {
+	shape := x.Shape()
+	H := shape[1]
+	W := shape[2]
+
+	hIdx := mlx.ArangeInt(0, H, 1, mlx.DtypeInt32)
+	hIdx = mlx.Reshape(hIdx, H, 1)
+	hIdx = mlx.BroadcastTo(hIdx, []int32{H, 2})
+	hIdx = mlx.Reshape(hIdx, H*2)
+
+	wIdx := mlx.ArangeInt(0, W, 1, mlx.DtypeInt32)
+	wIdx = mlx.Reshape(wIdx, W, 1)
+	wIdx = mlx.BroadcastTo(wIdx, []int32{W, 2})
+	wIdx = mlx.Reshape(wIdx, W*2)
+
+	x = mlx.Take(x, hIdx, 1)
+	x = mlx.Take(x, wIdx, 2)
+
+	return x
+}
+
+// VAEMidBlock is the middle block with attention
+type VAEMidBlock struct {
+	Resnet1   *ResnetBlock2D
+	Attention *VAEAttentionBlock
+	Resnet2   *ResnetBlock2D
+}
+
+// Forward applies the mid block
+func (mb *VAEMidBlock) Forward(x *mlx.Array) *mlx.Array {
+	x = mb.Resnet1.Forward(x)
+	x = mb.Attention.Forward(x)
+	x = mb.Resnet2.Forward(x)
+	return x
+}
+
+// DefaultTilingConfig returns reasonable defaults for tiled decoding
+// Matches diffusers: tile_latent_min_size=64, tile_overlap_factor=0.25
+func DefaultTilingConfig() *vae.TilingConfig {
+	return vae.DefaultTilingConfig()
+}
+
+// AutoencoderKLFlux2 is the Flux2 VAE with BatchNorm
+type AutoencoderKLFlux2 struct {
+	Config *VAEConfig
+
+	// Encoder components (for image editing)
+	EncoderConvIn  *Conv2D
+	EncoderMid     *VAEMidBlock
+	EncoderDown    []*DownEncoderBlock2D
+	EncoderNormOut *GroupNormLayer
+	EncoderConvOut *Conv2D
+
+	// Decoder components
+	DecoderConvIn  *Conv2D
+	DecoderMid     *VAEMidBlock
+	DecoderUp      []*UpDecoderBlock2D
+	DecoderNormOut *GroupNormLayer
+	DecoderConvOut *Conv2D
+
+	// Quant conv layers
+	QuantConv     *Conv2D
+	PostQuantConv *Conv2D
+
+	// BatchNorm for latent normalization
+	LatentBN *BatchNorm2D
+
+	// Tiling configuration (nil = no tiling)
+	Tiling *vae.TilingConfig
+}
+
+// DownEncoderBlock2D implements a downsampling encoder block
+type DownEncoderBlock2D struct {
+	ResnetBlocks []*ResnetBlock2D
+	Downsample   *Conv2D
+}
+
+// Forward applies the down encoder block
+func (db *DownEncoderBlock2D) Forward(x *mlx.Array) *mlx.Array {
+	for _, resnet := range db.ResnetBlocks {
+		x = resnet.Forward(x)
+	}
+
+	if db.Downsample != nil {
+		// Pad then conv with stride 2
+		x = mlx.Pad(x, []int32{0, 0, 0, 1, 0, 1, 0, 0})
+		x = db.Downsample.Forward(x)
+	}
+
+	return x
+}
+
+// Load loads the Flux2 VAE from ollama blob storage.
+func (m *AutoencoderKLFlux2) Load(modelManifest *manifest.ModelManifest) error {
+	fmt.Print("  Loading VAE... ")
+
+	// Load config from blob
+	var cfg VAEConfig
+	if err := modelManifest.ReadConfigJSON("vae/config.json", &cfg); err != nil {
+		return fmt.Errorf("config: %w", err)
+	}
+	m.Config = &cfg
+
+	// Load weights from tensor blobs
+	weights, err := manifest.LoadWeightsFromManifest(modelManifest, "vae")
+	if err != nil {
+		return fmt.Errorf("weights: %w", err)
+	}
+	if err := weights.Load(0); err != nil {
+		return fmt.Errorf("load weights: %w", err)
+	}
+	defer weights.ReleaseAll()
+
+	return m.loadWeights(weights, &cfg)
+}
+
+// loadWeights loads VAE weights from any WeightSource
+func (m *AutoencoderKLFlux2) loadWeights(weights safetensors.WeightSource, cfg *VAEConfig) error {
+	var err error
+
+	// Load encoder components (for image conditioning)
+	if err := m.loadEncoderWeights(weights, cfg); err != nil {
+		return fmt.Errorf("encoder: %w", err)
+	}
+
+	// Load decoder conv_in
+	m.DecoderConvIn = &Conv2D{Stride: 1, Padding: 1}
+	if err := safetensors.LoadModule(m.DecoderConvIn, weights, "decoder.conv_in"); err != nil {
+		return fmt.Errorf("decoder.conv_in: %w", err)
+	}
+
+	// Load mid block
+	m.DecoderMid, err = loadVAEMidBlock(weights, "decoder.mid_block", cfg.NormNumGroups)
+	if err != nil {
+		return fmt.Errorf("decoder.mid_block: %w", err)
+	}
+
+	// Load up blocks
+	numBlocks := len(cfg.BlockOutChannels)
+	m.DecoderUp = make([]*UpDecoderBlock2D, numBlocks)
+	for i := 0; i < numBlocks; i++ {
+		prefix := fmt.Sprintf("decoder.up_blocks.%d", i)
+		hasUpsample := i < numBlocks-1
+		m.DecoderUp[i], err = loadUpDecoderBlock2D(weights, prefix, cfg.LayersPerBlock+1, cfg.NormNumGroups, hasUpsample)
+		if err != nil {
+			return fmt.Errorf("%s: %w", prefix, err)
+		}
+	}
+
+	// Load decoder conv_norm_out and conv_out
+	m.DecoderNormOut = &GroupNormLayer{NumGroups: cfg.NormNumGroups, Eps: 1e-5}
+	if err := safetensors.LoadModule(m.DecoderNormOut, weights, "decoder.conv_norm_out"); err != nil {
+		return fmt.Errorf("decoder.conv_norm_out: %w", err)
+	}
+
+	m.DecoderConvOut = &Conv2D{Stride: 1, Padding: 1}
+	if err := safetensors.LoadModule(m.DecoderConvOut, weights, "decoder.conv_out"); err != nil {
+		return fmt.Errorf("decoder.conv_out: %w", err)
+	}
+
+	// Load post_quant_conv
+	if cfg.UsePostQuantConv {
+		m.PostQuantConv = &Conv2D{Stride: 1, Padding: 0}
+		if err := safetensors.LoadModule(m.PostQuantConv, weights, "post_quant_conv"); err != nil {
+			return fmt.Errorf("post_quant_conv: %w", err)
+		}
+	}
+
+	// Load latent BatchNorm (affine=False, so no weight/bias)
+	bnMean, err := weights.GetTensor("bn.running_mean")
+	if err != nil {
+		return fmt.Errorf("bn.running_mean: %w", err)
+	}
+	bnVar, err := weights.GetTensor("bn.running_var")
+	if err != nil {
+		return fmt.Errorf("bn.running_var: %w", err)
+	}
+	m.LatentBN = &BatchNorm2D{
+		RunningMean: bnMean,
+		RunningVar:  bnVar,
+		Weight:      nil, // affine=False
+		Bias:        nil, // affine=False
+		Eps:         cfg.BatchNormEps,
+		Momentum:    cfg.BatchNormMomentum,
+	}
+
+	fmt.Println("✓")
+	return nil
+}
+
+// loadVAEMidBlock loads the mid block.
+func loadVAEMidBlock(weights safetensors.WeightSource, prefix string, numGroups int32) (*VAEMidBlock, error) {
+	resnet1, err := loadResnetBlock2D(weights, prefix+".resnets.0", numGroups)
+	if err != nil {
+		return nil, err
+	}
+
+	attention, err := loadVAEAttentionBlock(weights, prefix+".attentions.0", numGroups)
+	if err != nil {
+		return nil, err
+	}
+
+	resnet2, err := loadResnetBlock2D(weights, prefix+".resnets.1", numGroups)
+	if err != nil {
+		return nil, err
+	}
+
+	return &VAEMidBlock{
+		Resnet1:   resnet1,
+		Attention: attention,
+		Resnet2:   resnet2,
+	}, nil
+}
+
+// loadResnetBlock2D loads a ResNet block.
+func loadResnetBlock2D(weights safetensors.WeightSource, prefix string, numGroups int32) (*ResnetBlock2D, error) {
+	block := &ResnetBlock2D{
+		Norm1:        &GroupNormLayer{NumGroups: numGroups, Eps: 1e-5},
+		Conv1:        &Conv2D{Stride: 1, Padding: 1},
+		Norm2:        &GroupNormLayer{NumGroups: numGroups, Eps: 1e-5},
+		Conv2:        &Conv2D{Stride: 1, Padding: 1},
+		ConvShortcut: &Conv2D{Stride: 1, Padding: 0}, // Pre-allocate for optional loading
+	}
+	if err := safetensors.LoadModule(block, weights, prefix); err != nil {
+		return nil, err
+	}
+	// If ConvShortcut wasn't loaded (no weights found), nil it out
+	if block.ConvShortcut.Weight == nil {
+		block.ConvShortcut = nil
+	}
+	return block, nil
+}
+
+// loadVAEAttentionBlock loads an attention block using LoadModule.
+func loadVAEAttentionBlock(weights safetensors.WeightSource, prefix string, numGroups int32) (*VAEAttentionBlock, error) {
+	ab := &VAEAttentionBlock{
+		GroupNorm: &GroupNormLayer{NumGroups: numGroups, Eps: 1e-5},
+	}
+	if err := safetensors.LoadModule(ab, weights, prefix); err != nil {
+		return nil, err
+	}
+	return ab, nil
+}
+
+// loadUpDecoderBlock2D loads an up decoder block.
+func loadUpDecoderBlock2D(weights safetensors.WeightSource, prefix string, numLayers, numGroups int32, hasUpsample bool) (*UpDecoderBlock2D, error) {
+	resnets := make([]*ResnetBlock2D, numLayers)
+	for i := int32(0); i < numLayers; i++ {
+		resPrefix := fmt.Sprintf("%s.resnets.%d", prefix, i)
+		resnet, err := loadResnetBlock2D(weights, resPrefix, numGroups)
+		if err != nil {
+			return nil, err
+		}
+		resnets[i] = resnet
+	}
+
+	var upsample *Conv2D
+	if hasUpsample {
+		upsample = &Conv2D{Stride: 1, Padding: 1}
+		if err := safetensors.LoadModule(upsample, weights, prefix+".upsamplers.0.conv"); err != nil {
+			return nil, err
+		}
+	}
+
+	return &UpDecoderBlock2D{
+		ResnetBlocks: resnets,
+		Upsample:     upsample,
+	}, nil
+}
+
+// Patchify converts latents [B, C, H, W] to patches [B, H*W/4, C*4] using 2x2 patches
+// This is the inverse of the VAE's patchify for feeding to transformer
+func (vae *AutoencoderKLFlux2) Patchify(latents *mlx.Array) *mlx.Array {
+	shape := latents.Shape()
+	B := shape[0]
+	C := shape[1]
+	H := shape[2]
+	W := shape[3]
+
+	patchH := vae.Config.PatchSize[0]
+	patchW := vae.Config.PatchSize[1]
+
+	pH := H / patchH
+	pW := W / patchW
+
+	// [B, C, H, W] -> [B, C, pH, patchH, pW, patchW]
+	x := mlx.Reshape(latents, B, C, pH, patchH, pW, patchW)
+	// [B, C, pH, patchH, pW, patchW] -> [B, pH, pW, C, patchH, patchW]
+	x = mlx.Transpose(x, 0, 2, 4, 1, 3, 5)
+	// [B, pH, pW, C, patchH, patchW] -> [B, pH*pW, C*patchH*patchW]
+	return mlx.Reshape(x, B, pH*pW, C*patchH*patchW)
+}
+
+// Unpatchify converts patches [B, L, C*4] back to [B, C, H, W]
+func (vae *AutoencoderKLFlux2) Unpatchify(patches *mlx.Array, pH, pW, C int32) *mlx.Array {
+	shape := patches.Shape()
+	B := shape[0]
+
+	patchH := vae.Config.PatchSize[0]
+	patchW := vae.Config.PatchSize[1]
+
+	// [B, pH*pW, C*patchH*patchW] -> [B, pH, pW, C, patchH, patchW]
+	x := mlx.Reshape(patches, B, pH, pW, C, patchH, patchW)
+	// [B, pH, pW, C, patchH, patchW] -> [B, C, pH, patchH, pW, patchW]
+	x = mlx.Transpose(x, 0, 3, 1, 4, 2, 5)
+	// [B, C, pH, patchH, pW, patchW] -> [B, C, H, W]
+	H := pH * patchH
+	W := pW * patchW
+	return mlx.Reshape(x, B, C, H, W)
+}
+
+// denormalizePatchified applies inverse batch normalization to patchified latents.
+// Input: [B, L, 128] where 128 = 32 latent channels * 4 (2x2 patch)
+// Output: [B, L, 128] denormalized
+func (vae *AutoencoderKLFlux2) denormalizePatchified(x *mlx.Array) *mlx.Array {
+	shape := x.Shape()
+	C := shape[2] // 128
+
+	// Reshape stats for broadcasting [1, 1, C]
+	mean := mlx.Reshape(vae.LatentBN.RunningMean, 1, 1, C)
+	variance := mlx.Reshape(vae.LatentBN.RunningVar, 1, 1, C)
+
+	// Inverse BN (affine=False): x_denorm = x * sqrt(var + eps) + mean
+	if vae.LatentBN.Bias != nil {
+		bias := mlx.Reshape(vae.LatentBN.Bias, 1, 1, C)
+		x = mlx.Sub(x, bias)
+	}
+	if vae.LatentBN.Weight != nil {
+		weight := mlx.Reshape(vae.LatentBN.Weight, 1, 1, C)
+		x = mlx.Div(x, weight)
+	}
+	x = mlx.Mul(x, mlx.Sqrt(mlx.AddScalar(variance, vae.LatentBN.Eps)))
+	x = mlx.Add(x, mean)
+
+	return x
+}
+
+// Decode decodes latent patches to images.
+// If Tiling is set, uses tiled decoding to reduce memory for large images.
+// latents: [B, L, C*4] patchified latents from transformer
+// pH, pW: patch grid dimensions
+// Returns: [B, 3, H, W] image tensor
+func (v *AutoencoderKLFlux2) Decode(latents *mlx.Array, pH, pW int32) *mlx.Array {
+	// Denormalize patchified latents
+	z := v.denormalizePatchified(latents)
+
+	// Unpatchify: [B, L, C*4] -> [B, C, H, W]
+	z = v.Unpatchify(z, pH, pW, v.Config.LatentChannels)
+
+	// Convert NCHW -> NHWC for processing
+	z = mlx.Transpose(z, 0, 2, 3, 1)
+
+	// Use tiled decoding if enabled
+	if v.Tiling != nil {
+		mlx.Eval(z)
+		return vae.DecodeTiled(z, v.Tiling, v.decodeTile)
+	}
+
+	// Direct decode (no tiling)
+	h := v.decodeTile(z)
+	h = mlx.ClipScalar(h, 0.0, 1.0, true, true)
+	h = mlx.Transpose(h, 0, 3, 1, 2)
+	return h
+}
+
+// decodeTile decodes a single latent tile to pixels (internal helper)
+// z: [B, H, W, C] latent tile in NHWC format
+// Returns: [B, H*8, W*8, 3] pixel tile in NHWC format (before clipping)
+func (vae *AutoencoderKLFlux2) decodeTile(z *mlx.Array) *mlx.Array {
+	// Post-quant conv
+	if vae.PostQuantConv != nil {
+		z = vae.PostQuantConv.Forward(z)
+	}
+
+	// Decoder
+	h := vae.DecoderConvIn.Forward(z)
+	h = vae.DecoderMid.Forward(h)
+
+	for _, upBlock := range vae.DecoderUp {
+		h = upBlock.Forward(h)
+	}
+
+	h = vae.DecoderNormOut.Forward(h)
+	h = mlx.SiLU(h)
+	h = vae.DecoderConvOut.Forward(h)
+
+	// VAE outputs [-1, 1], convert to [0, 1]
+	h = mlx.MulScalar(h, 0.5)
+	h = mlx.AddScalar(h, 0.5)
+
+	return h
+}
+
+// loadEncoderWeights loads the encoder components for image conditioning
+func (m *AutoencoderKLFlux2) loadEncoderWeights(weights safetensors.WeightSource, cfg *VAEConfig) error {
+	var err error
+
+	// Load encoder conv_in
+	m.EncoderConvIn = &Conv2D{Stride: 1, Padding: 1}
+	if err := safetensors.LoadModule(m.EncoderConvIn, weights, "encoder.conv_in"); err != nil {
+		return fmt.Errorf("encoder.conv_in: %w", err)
+	}
+
+	// Load encoder down blocks
+	numBlocks := len(cfg.BlockOutChannels)
+	m.EncoderDown = make([]*DownEncoderBlock2D, numBlocks)
+	for i := 0; i < numBlocks; i++ {
+		prefix := fmt.Sprintf("encoder.down_blocks.%d", i)
+		hasDownsample := i < numBlocks-1
+		m.EncoderDown[i], err = loadDownEncoderBlock2D(weights, prefix, cfg.LayersPerBlock, cfg.NormNumGroups, hasDownsample)
+		if err != nil {
+			return fmt.Errorf("%s: %w", prefix, err)
+		}
+	}
+
+	// Load encoder mid block
+	m.EncoderMid, err = loadVAEMidBlock(weights, "encoder.mid_block", cfg.NormNumGroups)
+	if err != nil {
+		return fmt.Errorf("encoder.mid_block: %w", err)
+	}
+
+	// Load encoder conv_norm_out and conv_out
+	m.EncoderNormOut = &GroupNormLayer{NumGroups: cfg.NormNumGroups, Eps: 1e-5}
+	if err := safetensors.LoadModule(m.EncoderNormOut, weights, "encoder.conv_norm_out"); err != nil {
+		return fmt.Errorf("encoder.conv_norm_out: %w", err)
+	}
+
+	m.EncoderConvOut = &Conv2D{Stride: 1, Padding: 1}
+	if err := safetensors.LoadModule(m.EncoderConvOut, weights, "encoder.conv_out"); err != nil {
+		return fmt.Errorf("encoder.conv_out: %w", err)
+	}
+
+	// Load quant_conv (for encoding)
+	if cfg.UseQuantConv {
+		m.QuantConv = &Conv2D{Stride: 1, Padding: 0}
+		if err := safetensors.LoadModule(m.QuantConv, weights, "quant_conv"); err != nil {
+			return fmt.Errorf("quant_conv: %w", err)
+		}
+	}
+
+	return nil
+}
+
+// loadDownEncoderBlock2D loads a down encoder block.
+func loadDownEncoderBlock2D(weights safetensors.WeightSource, prefix string, numLayers, numGroups int32, hasDownsample bool) (*DownEncoderBlock2D, error) {
+	resnets := make([]*ResnetBlock2D, numLayers)
+	for i := int32(0); i < numLayers; i++ {
+		resPrefix := fmt.Sprintf("%s.resnets.%d", prefix, i)
+		resnet, err := loadResnetBlock2D(weights, resPrefix, numGroups)
+		if err != nil {
+			return nil, err
+		}
+		resnets[i] = resnet
+	}
+
+	var downsample *Conv2D
+	if hasDownsample {
+		downsample = &Conv2D{Stride: 2, Padding: 0}
+		if err := safetensors.LoadModule(downsample, weights, prefix+".downsamplers.0.conv"); err != nil {
+			return nil, err
+		}
+	}
+
+	return &DownEncoderBlock2D{
+		ResnetBlocks: resnets,
+		Downsample:   downsample,
+	}, nil
+}
+
+// EncodeImage encodes an image to normalized latents.
+// image: [B, 3, H, W] image tensor in [-1, 1]
+// Returns: [B, L, C*4] patchified normalized latents
+func (vae *AutoencoderKLFlux2) EncodeImage(image *mlx.Array) *mlx.Array {
+	// Convert NCHW -> NHWC
+	x := mlx.Transpose(image, 0, 2, 3, 1)
+
+	// Encoder
+	h := vae.EncoderConvIn.Forward(x)
+
+	for _, downBlock := range vae.EncoderDown {
+		h = downBlock.Forward(h)
+	}
+
+	h = vae.EncoderMid.Forward(h)
+	h = vae.EncoderNormOut.Forward(h)
+	h = mlx.SiLU(h)
+	h = vae.EncoderConvOut.Forward(h)
+
+	// Quant conv outputs [B, H, W, 2*latent_channels] (mean + logvar)
+	if vae.QuantConv != nil {
+		h = vae.QuantConv.Forward(h)
+	}
+
+	// Take only the mean (first latent_channels) - deterministic encoding
+	// h is [B, H, W, 64] -> take first 32 channels for mean
+	shape := h.Shape()
+	latentChannels := vae.Config.LatentChannels // 32
+	h = mlx.Slice(h, []int32{0, 0, 0, 0}, []int32{shape[0], shape[1], shape[2], latentChannels})
+
+	// Convert NHWC -> NCHW for patchifying
+	h = mlx.Transpose(h, 0, 3, 1, 2)
+
+	// Patchify: [B, C, H, W] -> [B, L, C*4]
+	h = vae.Patchify(h)
+
+	// Apply BatchNorm on patchified latents [B, L, 128]
+	// The BatchNorm has 128 channels matching the patchified dimension
+	h = vae.normalizePatchified(h)
+
+	return h
+}
+
+// normalizePatchified applies batch normalization to patchified latents.
+// Input: [B, L, 128] where 128 = 32 latent channels * 4 (2x2 patch)
+// Output: [B, L, 128] normalized
+func (vae *AutoencoderKLFlux2) normalizePatchified(x *mlx.Array) *mlx.Array {
+	shape := x.Shape()
+	C := shape[2] // 128
+
+	// Reshape stats for broadcasting [1, 1, C]
+	mean := mlx.Reshape(vae.LatentBN.RunningMean, 1, 1, C)
+	variance := mlx.Reshape(vae.LatentBN.RunningVar, 1, 1, C)
+
+	// Normalize: (x - mean) / sqrt(var + eps)
+	xNorm := mlx.Sub(x, mean)
+	xNorm = mlx.Div(xNorm, mlx.Sqrt(mlx.AddScalar(variance, vae.LatentBN.Eps)))
+
+	// Scale and shift (only if affine=True)
+	if vae.LatentBN.Weight != nil {
+		weight := mlx.Reshape(vae.LatentBN.Weight, 1, 1, C)
+		xNorm = mlx.Mul(xNorm, weight)
+	}
+	if vae.LatentBN.Bias != nil {
+		bias := mlx.Reshape(vae.LatentBN.Bias, 1, 1, C)
+		xNorm = mlx.Add(xNorm, bias)
+	}
+
+	return xNorm
+}

x/imagegen/models/qwen3/text_encoder.go

@@ -0,0 +1,388 @@
+// Package qwen3 provides a shared Qwen3 text encoder used by multiple image generation models.
+package qwen3
+
+import (
+	"fmt"
+	"math"
+
+	"github.com/ollama/ollama/x/imagegen/manifest"
+	"github.com/ollama/ollama/x/imagegen/mlx"
+	"github.com/ollama/ollama/x/imagegen/nn"
+	"github.com/ollama/ollama/x/imagegen/safetensors"
+	"github.com/ollama/ollama/x/imagegen/tokenizer"
+)
+
+// Config holds Qwen3 text encoder configuration
+type Config struct {
+	HiddenSize        int32   `json:"hidden_size"`
+	NumHiddenLayers   int32   `json:"num_hidden_layers"`
+	IntermediateSize  int32   `json:"intermediate_size"`
+	NumAttentionHeads int32   `json:"num_attention_heads"`
+	NumKeyValueHeads  int32   `json:"num_key_value_heads"`
+	VocabSize         int32   `json:"vocab_size"`
+	RMSNormEps        float32 `json:"rms_norm_eps"`
+	RopeTheta         float32 `json:"rope_theta"`
+	HeadDim           int32   `json:"head_dim"`
+}
+
+// Attention implements Qwen3 attention with QK norms
+type Attention struct {
+	QProj nn.LinearLayer `weight:"q_proj"`
+	KProj nn.LinearLayer `weight:"k_proj"`
+	VProj nn.LinearLayer `weight:"v_proj"`
+	OProj nn.LinearLayer `weight:"o_proj"`
+	QNorm *nn.RMSNorm    `weight:"q_norm"`
+	KNorm *nn.RMSNorm    `weight:"k_norm"`
+	// Computed fields
+	NHeads    int32
+	NKVHeads  int32
+	HeadDim   int32
+	Scale     float32
+	RopeTheta float32
+}
+
+// applyRoPEQwen3 applies the custom RoPE for Qwen3 text encoder
+func applyRoPEQwen3(x *mlx.Array, seqLen int32, theta float32) *mlx.Array {
+	shape := x.Shape()
+	B := shape[0]
+	L := shape[1]
+	H := shape[2]
+	D := shape[3]
+	half := D / 2
+
+	freqsArr := make([]float32, half)
+	logTheta := float32(math.Log(float64(theta)))
+	for i := int32(0); i < half; i++ {
+		freqsArr[i] = float32(math.Exp(float64(-logTheta * float32(i) / float32(half))))
+	}
+	freqs := mlx.NewArray(freqsArr, []int32{half})
+
+	posArr := make([]float32, seqLen)
+	for i := int32(0); i < seqLen; i++ {
+		posArr[i] = float32(i)
+	}
+	pos := mlx.NewArray(posArr, []int32{seqLen})
+
+	posExpanded := mlx.Reshape(pos, seqLen, 1)
+	freqsExpanded := mlx.Reshape(freqs, 1, half)
+	args := mlx.Mul(posExpanded, freqsExpanded)
+
+	cosVals := mlx.Cos(args)
+	sinVals := mlx.Sin(args)
+	cosVals = mlx.Reshape(cosVals, seqLen, 1, half)
+	sinVals = mlx.Reshape(sinVals, seqLen, 1, half)
+
+	x1 := mlx.Slice(x, []int32{0, 0, 0, 0}, []int32{B, L, H, half})
+	x2 := mlx.Slice(x, []int32{0, 0, 0, half}, []int32{B, L, H, D})
+
+	part1 := mlx.Sub(mlx.Mul(x1, cosVals), mlx.Mul(x2, sinVals))
+	part2 := mlx.Add(mlx.Mul(x1, sinVals), mlx.Mul(x2, cosVals))
+
+	return mlx.Concatenate([]*mlx.Array{part1, part2}, 3)
+}
+
+// Forward computes attention with causal masking and optional padding mask
+func (attn *Attention) Forward(x *mlx.Array, mask *mlx.Array, maskMode string) *mlx.Array {
+	shape := x.Shape()
+	B := shape[0]
+	L := shape[1]
+
+	q := attn.QProj.Forward(x)
+	k := attn.KProj.Forward(x)
+	v := attn.VProj.Forward(x)
+
+	q = mlx.Reshape(q, B, L, attn.NHeads, attn.HeadDim)
+	k = mlx.Reshape(k, B, L, attn.NKVHeads, attn.HeadDim)
+	v = mlx.Reshape(v, B, L, attn.NKVHeads, attn.HeadDim)
+
+	// QK norm uses 1e-6 hardcoded (Qwen3 specific)
+	q = attn.QNorm.Forward(q, 1e-6)
+	k = attn.KNorm.Forward(k, 1e-6)
+
+	q = applyRoPEQwen3(q, L, attn.RopeTheta)
+	k = applyRoPEQwen3(k, L, attn.RopeTheta)
+
+	q = mlx.Transpose(q, 0, 2, 1, 3)
+	k = mlx.Transpose(k, 0, 2, 1, 3)
+	v = mlx.Transpose(v, 0, 2, 1, 3)
+
+	if attn.NKVHeads < attn.NHeads {
+		repeats := attn.NHeads / attn.NKVHeads
+		k = repeatKV(k, repeats)
+		v = repeatKV(v, repeats)
+	}
+
+	out := mlx.ScaledDotProductAttentionWithSinks(q, k, v, attn.Scale, maskMode, mask, nil)
+
+	out = mlx.Transpose(out, 0, 2, 1, 3)
+	out = mlx.Reshape(out, B, L, attn.NHeads*attn.HeadDim)
+
+	out = attn.OProj.Forward(out)
+
+	return out
+}
+
+// repeatKV repeats key/value heads for GQA
+func repeatKV(x *mlx.Array, repeats int32) *mlx.Array {
+	if repeats == 1 {
+		return x
+	}
+	shape := x.Shape()
+	x = mlx.ExpandDims(x, 2)
+	x = mlx.Tile(x, []int32{1, 1, repeats, 1, 1})
+	return mlx.Reshape(x, shape[0], shape[1]*repeats, shape[2], shape[3])
+}
+
+// MLP implements Qwen3 SwiGLU MLP
+type MLP struct {
+	GateProj nn.LinearLayer `weight:"gate_proj"`
+	UpProj   nn.LinearLayer `weight:"up_proj"`
+	DownProj nn.LinearLayer `weight:"down_proj"`
+}
+
+// Forward applies the MLP
+func (m *MLP) Forward(x *mlx.Array) *mlx.Array {
+	gate := m.GateProj.Forward(x)
+	gate = mlx.SiLU(gate)
+	up := m.UpProj.Forward(x)
+	h := mlx.Mul(gate, up)
+	return m.DownProj.Forward(h)
+}
+
+// Block represents a single Qwen3 transformer block
+type Block struct {
+	Attention         *Attention  `weight:"self_attn"`
+	MLP               *MLP        `weight:"mlp"`
+	InputLayerNorm    *nn.RMSNorm `weight:"input_layernorm"`
+	PostAttnLayerNorm *nn.RMSNorm `weight:"post_attention_layernorm"`
+}
+
+// Forward applies the Qwen3 block
+func (qb *Block) Forward(x *mlx.Array, eps float32, mask *mlx.Array, maskMode string) *mlx.Array {
+	h := qb.InputLayerNorm.Forward(x, eps)
+	attnOut := qb.Attention.Forward(h, mask, maskMode)
+	x = mlx.Add(x, attnOut)
+
+	h = qb.PostAttnLayerNorm.Forward(x, eps)
+	mlpOut := qb.MLP.Forward(h)
+	x = mlx.Add(x, mlpOut)
+
+	return x
+}
+
+// TextEncoder is the full Qwen3 encoder
+type TextEncoder struct {
+	EmbedTokens *nn.Embedding `weight:"model.embed_tokens"`
+	Layers      []*Block      `weight:"model.layers"`
+	FinalNorm   *nn.RMSNorm   `weight:"model.norm"`
+	*Config
+}
+
+// Load loads the Qwen3 text encoder from ollama blob storage.
+func (m *TextEncoder) Load(modelManifest *manifest.ModelManifest, configPath string) error {
+	fmt.Print("  Loading text encoder... ")
+
+	// Load config from blob
+	var cfg Config
+	if err := modelManifest.ReadConfigJSON(configPath, &cfg); err != nil {
+		return fmt.Errorf("config: %w", err)
+	}
+	m.Config = &cfg
+	m.Layers = make([]*Block, cfg.NumHiddenLayers)
+
+	// Load weights from tensor blobs
+	weights, err := manifest.LoadWeightsFromManifest(modelManifest, "text_encoder")
+	if err != nil {
+		return fmt.Errorf("weights: %w", err)
+	}
+	if err := weights.Load(0); err != nil {
+		return fmt.Errorf("load weights: %w", err)
+	}
+	defer weights.ReleaseAll()
+
+	return m.loadWeights(weights)
+}
+
+// loadWeights loads weights from any WeightSource into the model
+func (m *TextEncoder) loadWeights(weights safetensors.WeightSource) error {
+	if err := safetensors.LoadModule(m, weights, ""); err != nil {
+		return fmt.Errorf("load module: %w", err)
+	}
+	m.initComputedFields()
+	fmt.Println("✓")
+	return nil
+}
+
+// initComputedFields initializes computed fields after loading weights
+func (m *TextEncoder) initComputedFields() {
+	cfg := m.Config
+	m.FinalNorm.Eps = cfg.RMSNormEps
+	for _, block := range m.Layers {
+		// Attention
+		block.Attention.NHeads = cfg.NumAttentionHeads
+		block.Attention.NKVHeads = cfg.NumKeyValueHeads
+		block.Attention.HeadDim = cfg.HeadDim
+		block.Attention.Scale = float32(1.0 / math.Sqrt(float64(cfg.HeadDim)))
+		block.Attention.RopeTheta = cfg.RopeTheta
+		block.Attention.QNorm.Eps = cfg.RMSNormEps
+		block.Attention.KNorm.Eps = cfg.RMSNormEps
+		// Block norms
+		block.InputLayerNorm.Eps = cfg.RMSNormEps
+		block.PostAttnLayerNorm.Eps = cfg.RMSNormEps
+	}
+}
+
+// Forward encodes text tokens with provided attention mask (LxL) and mask mode.
+func (te *TextEncoder) Forward(tokens *mlx.Array, attnMask *mlx.Array, maskMode string) *mlx.Array {
+	h := te.EmbedTokens.Forward(tokens)
+	eps := te.RMSNormEps
+
+	for _, layer := range te.Layers {
+		h = layer.Forward(h, eps, attnMask, maskMode)
+	}
+
+	// Apply final RMS norm
+	h = te.FinalNorm.Forward(h, eps)
+
+	return h
+}
+
+// ForwardWithLayerOutputs encodes text tokens and returns hidden states from specified layers.
+// This is used by Flux2 which needs embeddings from specific intermediate layers.
+func (te *TextEncoder) ForwardWithLayerOutputs(tokens *mlx.Array, layerIndices []int, attnMask *mlx.Array, maskMode string) []*mlx.Array {
+	h := te.EmbedTokens.Forward(tokens)
+	eps := te.RMSNormEps
+
+	outputs := make([]*mlx.Array, len(layerIndices))
+	layerSet := make(map[int]int)
+	for i, idx := range layerIndices {
+		layerSet[idx] = i
+	}
+
+	for i, layer := range te.Layers {
+		h = layer.Forward(h, eps, attnMask, maskMode)
+		if outIdx, ok := layerSet[i]; ok {
+			outputs[outIdx] = h
+		}
+	}
+
+	return outputs
+}
+
+// ApplyChatTemplate wraps prompt in Qwen3 chat format.
+// If think is true, adds the <think></think> block after the assistant tag
+// (matches tokenizer.apply_chat_template with enable_thinking=False in Python).
+func ApplyChatTemplate(prompt string, think bool) string {
+	base := "<|im_start|>user\n" + prompt + "<|im_end|>\n<|im_start|>assistant\n"
+	if think {
+		return base + "<think>\n\n</think>\n\n"
+	}
+	return base
+}
+
+// EncodePrompt encodes a text prompt using the tokenizer and encoder.
+// If think is true, includes the <think></think> block in the chat template.
+func (te *TextEncoder) EncodePrompt(tok *tokenizer.Tokenizer, prompt string, maxLen int, think bool) (*mlx.Array, *mlx.Array) {
+	formattedPrompt := ApplyChatTemplate(prompt, think)
+
+	tokens := tok.Encode(formattedPrompt, false)
+
+	if len(tokens) > maxLen {
+		tokens = tokens[:maxLen]
+	}
+
+	maskData := make([]float32, maxLen)
+	for i := 0; i < len(tokens); i++ {
+		maskData[i] = 1.0
+	}
+
+	// Get PAD token (different from EOS for Qwen3)
+	padToken := tok.PAD()
+	if padToken < 0 {
+		padToken = tok.EOS() // fallback
+	}
+
+	paddedTokens := make([]int32, maxLen)
+	copy(paddedTokens, tokens)
+	for i := len(tokens); i < maxLen; i++ {
+		paddedTokens[i] = padToken
+	}
+
+	tokensArr := mlx.NewArrayInt32(paddedTokens, []int32{1, int32(maxLen)})
+	maskArr := mlx.NewArray(maskData, []int32{1, int32(maxLen)})
+
+	// Build combined causal + PAD mask [L, L]
+	// mask[i,j] = 0 if (j <= i AND valid[j]) else -inf
+	L := int32(maxLen)
+	validLen := int32(len(tokens))
+	combinedMaskData := make([]float32, L*L)
+	negInf := float32(-1e9)
+	for i := int32(0); i < L; i++ {
+		for j := int32(0); j < L; j++ {
+			idx := i*L + j
+			if j <= i && j < validLen {
+				combinedMaskData[idx] = 0
+			} else {
+				combinedMaskData[idx] = negInf
+			}
+		}
+	}
+	maskMat := mlx.NewArray(combinedMaskData, []int32{L, L})
+
+	embeddings := te.Forward(tokensArr, maskMat, "")
+
+	return embeddings, maskArr
+}
+
+// EncodePromptWithLayers encodes a text prompt and returns embeddings from specified layers.
+// Used by Flux2 which concatenates embeddings from multiple intermediate layers.
+// If think is true, includes the <think></think> block in the chat template.
+// Returns embeddings and padded sequence length.
+func (te *TextEncoder) EncodePromptWithLayers(tok *tokenizer.Tokenizer, prompt string, maxLen int, layerIndices []int, think bool) (*mlx.Array, int32) {
+	formattedPrompt := ApplyChatTemplate(prompt, think)
+	tokens := tok.Encode(formattedPrompt, false)
+
+	if len(tokens) > maxLen {
+		tokens = tokens[:maxLen]
+	}
+
+	// Pad to maxLen
+	padToken := tok.PAD()
+	if padToken < 0 {
+		padToken = tok.EOS() // fallback
+	}
+	padded := make([]int32, maxLen)
+	copy(padded, tokens)
+	for i := len(tokens); i < maxLen; i++ {
+		padded[i] = padToken
+	}
+	tokensArr := mlx.NewArrayInt32(padded, []int32{1, int32(maxLen)})
+
+	// Build combined causal + PAD mask [L, L]
+	// mask[i,j] = 0 if (j <= i AND valid[j]) else -inf
+	// This combines causal masking with PAD token masking
+	L := int32(maxLen)
+	validLen := int32(len(tokens))
+	maskData := make([]float32, L*L)
+	negInf := float32(-1e9)
+	for i := int32(0); i < L; i++ {
+		for j := int32(0); j < L; j++ {
+			idx := i*L + j
+			if j <= i && j < validLen {
+				maskData[idx] = 0 // allowed: causal OK and not PAD
+			} else {
+				maskData[idx] = negInf // blocked: future or PAD
+			}
+		}
+	}
+	maskMat := mlx.NewArray(maskData, []int32{L, L})
+
+	layerOutputs := te.ForwardWithLayerOutputs(tokensArr, layerIndices, maskMat, "")
+
+	// Concatenate layer outputs along the hidden dimension
+	// Each output is [B, L, hidden_dim], result is [B, L, num_layers * hidden_dim]
+	embeddings := mlx.Concatenate(layerOutputs, 2)
+
+	// Return embeddings and padded length
+	return embeddings, int32(maxLen)
+}

x/imagegen/models/zimage/scheduler.go

@@ -0,0 +1,141 @@
+package zimage
+
+import (
+	"math"
+
+	"github.com/ollama/ollama/x/imagegen/mlx"
+)
+
+// FlowMatchSchedulerConfig holds scheduler configuration
+type FlowMatchSchedulerConfig struct {
+	NumTrainTimesteps  int32   `json:"num_train_timesteps"`  // 1000
+	Shift              float32 `json:"shift"`                // 3.0
+	UseDynamicShifting bool    `json:"use_dynamic_shifting"` // false
+}
+
+// DefaultFlowMatchSchedulerConfig returns default config
+func DefaultFlowMatchSchedulerConfig() *FlowMatchSchedulerConfig {
+	return &FlowMatchSchedulerConfig{
+		NumTrainTimesteps:  1000,
+		Shift:              3.0,
+		UseDynamicShifting: true, // Z-Image-Turbo uses dynamic shifting
+	}
+}
+
+// FlowMatchEulerScheduler implements the Flow Match Euler discrete scheduler
+// This is used in Z-Image-Turbo for fast sampling
+type FlowMatchEulerScheduler struct {
+	Config    *FlowMatchSchedulerConfig
+	Timesteps []float32 // Discretized timesteps
+	Sigmas    []float32 // Noise levels at each timestep
+	NumSteps  int       // Number of inference steps
+}
+
+// NewFlowMatchEulerScheduler creates a new scheduler
+func NewFlowMatchEulerScheduler(cfg *FlowMatchSchedulerConfig) *FlowMatchEulerScheduler {
+	return &FlowMatchEulerScheduler{
+		Config: cfg,
+	}
+}
+
+// SetTimesteps sets up the scheduler for the given number of inference steps
+func (s *FlowMatchEulerScheduler) SetTimesteps(numSteps int) {
+	s.SetTimestepsWithMu(numSteps, 0)
+}
+
+// SetTimestepsWithMu sets up the scheduler with dynamic mu shift
+func (s *FlowMatchEulerScheduler) SetTimestepsWithMu(numSteps int, mu float32) {
+	s.NumSteps = numSteps
+
+	// Create evenly spaced timesteps from 1.0 to 0.0 (flow matching goes t=1 to t=0)
+	// Match Python: np.linspace(1.0, 0.0, num_inference_steps + 1)
+	s.Timesteps = make([]float32, numSteps+1)
+	s.Sigmas = make([]float32, numSteps+1)
+
+	for i := 0; i <= numSteps; i++ {
+		t := 1.0 - float32(i)/float32(numSteps)
+
+		// Apply time shift if using dynamic shifting
+		if s.Config.UseDynamicShifting && mu != 0 {
+			t = s.timeShift(mu, t)
+		}
+
+		s.Timesteps[i] = t
+		s.Sigmas[i] = t
+	}
+}
+
+// timeShift applies the dynamic time shift (match Python)
+func (s *FlowMatchEulerScheduler) timeShift(mu float32, t float32) float32 {
+	if t <= 0 {
+		return 0
+	}
+	// exp(mu) / (exp(mu) + (1/t - 1))
+	expMu := float32(math.Exp(float64(mu)))
+	return expMu / (expMu + (1.0/t - 1.0))
+}
+
+// Step performs one denoising step
+// modelOutput: predicted velocity/noise from the model
+// timestepIdx: current timestep index
+// sample: current noisy sample
+// Returns: denoised sample for next step
+func (s *FlowMatchEulerScheduler) Step(modelOutput, sample *mlx.Array, timestepIdx int) *mlx.Array {
+	// Get current and next sigma
+	sigma := s.Sigmas[timestepIdx]
+	sigmaNext := s.Sigmas[timestepIdx+1]
+
+	// Euler step: x_{t-dt} = x_t + (sigma_next - sigma) * v_t
+	// where v_t is the velocity predicted by the model
+	dt := sigmaNext - sigma // This is negative (going from noise to clean)
+
+	// x_next = x + dt * velocity
+	scaledOutput := mlx.MulScalar(modelOutput, dt)
+	return mlx.Add(sample, scaledOutput)
+}
+
+// ScaleSample scales the sample for model input (identity for flow matching)
+func (s *FlowMatchEulerScheduler) ScaleSample(sample *mlx.Array, timestepIdx int) *mlx.Array {
+	// Flow matching doesn't need scaling
+	return sample
+}
+
+// GetTimestep returns the timestep value at the given index
+func (s *FlowMatchEulerScheduler) GetTimestep(idx int) float32 {
+	if idx < len(s.Timesteps) {
+		return s.Timesteps[idx]
+	}
+	return 0.0
+}
+
+// GetTimesteps returns all timesteps (implements Scheduler interface)
+func (s *FlowMatchEulerScheduler) GetTimesteps() []float32 {
+	return s.Timesteps
+}
+
+// AddNoise adds noise to clean samples for a given timestep
+// Used for img2img or inpainting
+func (s *FlowMatchEulerScheduler) AddNoise(cleanSample, noise *mlx.Array, timestepIdx int) *mlx.Array {
+	// In flow matching: x_t = (1-t) * x_0 + t * noise
+	t := s.Timesteps[timestepIdx]
+	oneMinusT := 1.0 - t
+
+	scaledClean := mlx.MulScalar(cleanSample, oneMinusT)
+	scaledNoise := mlx.MulScalar(noise, t)
+
+	return mlx.Add(scaledClean, scaledNoise)
+}
+
+// InitNoise creates initial noise for sampling (BFloat16 for GPU efficiency)
+func (s *FlowMatchEulerScheduler) InitNoise(shape []int32, seed int64) *mlx.Array {
+	return mlx.RandomNormalWithDtype(shape, uint64(seed), mlx.DtypeBFloat16)
+}
+
+// GetLatentShape returns the latent shape for a given image size
+func GetLatentShape(batchSize, height, width, latentChannels int32, patchSize int32) []int32 {
+	// Latent is 8x smaller than image (VAE downscale)
+	latentH := height / 8
+	latentW := width / 8
+
+	return []int32{batchSize, latentChannels, latentH, latentW}
+}

x/imagegen/models/zimage/text_encoder.go

@@ -0,0 +1,17 @@
+package zimage
+
+import (
+	"github.com/ollama/ollama/x/imagegen/models/qwen3"
+)
+
+// Re-export types from shared qwen3 package for backwards compatibility
+type (
+	Qwen3Config      = qwen3.Config
+	Qwen3Attention   = qwen3.Attention
+	Qwen3MLP         = qwen3.MLP
+	Qwen3Block       = qwen3.Block
+	Qwen3TextEncoder = qwen3.TextEncoder
+)
+
+// ApplyChatTemplate wraps prompt in Qwen3 chat format
+var ApplyChatTemplate = qwen3.ApplyChatTemplate

x/imagegen/models/zimage/transformer.go

@@ -0,0 +1,759 @@
+// Package zimage implements the Z-Image diffusion transformer model.
+package zimage
+
+import (
+	"fmt"
+	"math"
+
+	"github.com/ollama/ollama/x/imagegen/cache"
+	"github.com/ollama/ollama/x/imagegen/manifest"
+	"github.com/ollama/ollama/x/imagegen/mlx"
+	"github.com/ollama/ollama/x/imagegen/nn"
+	"github.com/ollama/ollama/x/imagegen/safetensors"
+)
+
+// TransformerConfig holds Z-Image transformer configuration
+type TransformerConfig struct {
+	Dim            int32   `json:"dim"`
+	NHeads         int32   `json:"n_heads"`
+	NKVHeads       int32   `json:"n_kv_heads"`
+	NLayers        int32   `json:"n_layers"`
+	NRefinerLayers int32   `json:"n_refiner_layers"`
+	InChannels     int32   `json:"in_channels"`
+	PatchSize      int32   `json:"-"` // Computed from AllPatchSize
+	CapFeatDim     int32   `json:"cap_feat_dim"`
+	NormEps        float32 `json:"norm_eps"`
+	RopeTheta      float32 `json:"rope_theta"`
+	TScale         float32 `json:"t_scale"`
+	QKNorm         bool    `json:"qk_norm"`
+	AxesDims       []int32 `json:"axes_dims"`
+	AxesLens       []int32 `json:"axes_lens"`
+	AllPatchSize   []int32 `json:"all_patch_size"` // JSON array, PatchSize = first element
+}
+
+// TimestepEmbedder creates sinusoidal timestep embeddings
+// Output dimension is 256 (fixed), used for AdaLN modulation
+type TimestepEmbedder struct {
+	Linear1       nn.LinearLayer `weight:"mlp.0"`
+	Linear2       nn.LinearLayer `weight:"mlp.2"`
+	FreqEmbedSize int32          // 256 (computed)
+}
+
+// Forward computes timestep embeddings -> [B, 256]
+func (te *TimestepEmbedder) Forward(t *mlx.Array) *mlx.Array {
+	// t: [B] timesteps
+
+	// Create sinusoidal embedding
+	half := te.FreqEmbedSize / 2
+
+	// freqs = exp(-log(10000) * arange(half) / half)
+	freqs := make([]float32, half)
+	for i := int32(0); i < half; i++ {
+		freqs[i] = float32(math.Exp(-math.Log(10000.0) * float64(i) / float64(half)))
+	}
+	freqsArr := mlx.NewArray(freqs, []int32{1, half})
+
+	// t[:, None] * freqs[None, :] -> [B, half]
+	tExpanded := mlx.ExpandDims(t, 1) // [B, 1]
+	args := mlx.Mul(tExpanded, freqsArr)
+
+	// embedding = [cos(args), sin(args)] -> [B, 256]
+	cosArgs := mlx.Cos(args)
+	sinArgs := mlx.Sin(args)
+	embedding := mlx.Concatenate([]*mlx.Array{cosArgs, sinArgs}, 1)
+
+	// MLP: linear1 -> silu -> linear2
+	h := te.Linear1.Forward(embedding)
+	h = mlx.SiLU(h)
+	h = te.Linear2.Forward(h)
+
+	return h
+}
+
+// XEmbedder embeds image patches to model dimension
+type XEmbedder struct {
+	Linear nn.LinearLayer `weight:"2-1"`
+}
+
+// Forward embeds patchified image latents
+func (xe *XEmbedder) Forward(x *mlx.Array) *mlx.Array {
+	// x: [B, L, in_channels * 4] -> [B, L, dim]
+	return xe.Linear.Forward(x)
+}
+
+// CapEmbedder projects caption features to model dimension
+type CapEmbedder struct {
+	Norm     *nn.RMSNorm    `weight:"0"`
+	Linear   nn.LinearLayer `weight:"1"`
+	PadToken *mlx.Array     // loaded separately at root level
+}
+
+// Forward projects caption embeddings: [B, L, cap_feat_dim] -> [B, L, dim]
+func (ce *CapEmbedder) Forward(capFeats *mlx.Array) *mlx.Array {
+	// RMSNorm on last axis (uses 1e-6)
+	h := ce.Norm.Forward(capFeats, 1e-6)
+	// Linear projection
+	return ce.Linear.Forward(h)
+}
+
+// FeedForward implements SwiGLU FFN
+type FeedForward struct {
+	W1     nn.LinearLayer `weight:"w1"` // gate projection
+	W2     nn.LinearLayer `weight:"w2"` // down projection
+	W3     nn.LinearLayer `weight:"w3"` // up projection
+	OutDim int32          // computed from W2
+}
+
+// Forward applies SwiGLU: silu(W1(x)) * W3(x), then W2
+func (ff *FeedForward) Forward(x *mlx.Array) *mlx.Array {
+	shape := x.Shape()
+	B := shape[0]
+	L := shape[1]
+	D := shape[2]
+
+	// Reshape for matmul
+	x = mlx.Reshape(x, B*L, D)
+
+	gate := ff.W1.Forward(x)
+	gate = mlx.SiLU(gate)
+	up := ff.W3.Forward(x)
+	h := mlx.Mul(gate, up)
+	out := ff.W2.Forward(h)
+
+	return mlx.Reshape(out, B, L, ff.OutDim)
+}
+
+// Attention implements multi-head attention with QK norm
+type Attention struct {
+	ToQ   nn.LinearLayer `weight:"to_q"`
+	ToK   nn.LinearLayer `weight:"to_k"`
+	ToV   nn.LinearLayer `weight:"to_v"`
+	ToOut nn.LinearLayer `weight:"to_out.0"`
+	NormQ *mlx.Array     `weight:"norm_q.weight"` // [head_dim] for per-head RMSNorm
+	NormK *mlx.Array     `weight:"norm_k.weight"`
+	// Fused QKV (computed at init time for efficiency, not loaded from weights)
+	ToQKV nn.LinearLayer `weight:"-"` // Fused Q+K+V projection (created by FuseQKV)
+	Fused bool           `weight:"-"` // Whether to use fused QKV path
+	// Computed fields (not loaded from weights)
+	NHeads  int32   `weight:"-"`
+	HeadDim int32   `weight:"-"`
+	Dim     int32   `weight:"-"`
+	Scale   float32 `weight:"-"`
+}
+
+// FuseQKV creates a fused QKV projection by concatenating weights.
+// This reduces 3 matmuls to 1 for a ~5-10% speedup.
+// Note: Fusion is skipped for quantized weights as it would require complex
+// dequant-concat-requant operations. The FP8 memory bandwidth savings outweigh
+// the ~5% fusion benefit.
+func (attn *Attention) FuseQKV() {
+	if attn.ToQ == nil || attn.ToK == nil || attn.ToV == nil {
+		return
+	}
+
+	// Skip fusion for quantized weights - type assert to check
+	toQ, qOk := attn.ToQ.(*nn.Linear)
+	toK, kOk := attn.ToK.(*nn.Linear)
+	toV, vOk := attn.ToV.(*nn.Linear)
+	if !qOk || !kOk || !vOk {
+		// One or more are QuantizedLinear, skip fusion
+		return
+	}
+
+	if toQ.Weight == nil || toK.Weight == nil || toV.Weight == nil {
+		return
+	}
+
+	// Concatenate weights: [dim, dim] x 3 -> [3*dim, dim]
+	// Weight shapes: ToQ.Weight [out_dim, in_dim], etc.
+	qWeight := toQ.Weight
+	kWeight := toK.Weight
+	vWeight := toV.Weight
+
+	// Concatenate along output dimension (axis 0)
+	fusedWeight := mlx.Concatenate([]*mlx.Array{qWeight, kWeight, vWeight}, 0)
+
+	// Evaluate fused weight to ensure it's materialized
+	mlx.Eval(fusedWeight)
+
+	// Create fused linear layer
+	fusedLinear := &nn.Linear{Weight: fusedWeight}
+
+	// Handle bias if present
+	if toQ.Bias != nil && toK.Bias != nil && toV.Bias != nil {
+		fusedBias := mlx.Concatenate([]*mlx.Array{toQ.Bias, toK.Bias, toV.Bias}, 0)
+		mlx.Eval(fusedBias)
+		fusedLinear.Bias = fusedBias
+	}
+
+	attn.ToQKV = fusedLinear
+	attn.Fused = true
+}
+
+// Forward computes attention
+func (attn *Attention) Forward(x *mlx.Array, cos, sin *mlx.Array) *mlx.Array {
+	shape := x.Shape()
+	B := shape[0]
+	L := shape[1]
+	D := shape[2]
+
+	xFlat := mlx.Reshape(x, B*L, D)
+
+	var q, k, v *mlx.Array
+	if attn.Fused && attn.ToQKV != nil {
+		// Fused QKV path: single matmul then split
+		qkv := attn.ToQKV.Forward(xFlat) // [B*L, 3*dim]
+
+		// Split into Q, K, V along last dimension
+		// Each has shape [B*L, dim]
+		q = mlx.Slice(qkv, []int32{0, 0}, []int32{B * L, attn.Dim})
+		k = mlx.Slice(qkv, []int32{0, attn.Dim}, []int32{B * L, 2 * attn.Dim})
+		v = mlx.Slice(qkv, []int32{0, 2 * attn.Dim}, []int32{B * L, 3 * attn.Dim})
+	} else {
+		// Separate Q, K, V projections
+		q = attn.ToQ.Forward(xFlat)
+		k = attn.ToK.Forward(xFlat)
+		v = attn.ToV.Forward(xFlat)
+	}
+
+	// Reshape to [B, L, nheads, head_dim]
+	q = mlx.Reshape(q, B, L, attn.NHeads, attn.HeadDim)
+	k = mlx.Reshape(k, B, L, attn.NHeads, attn.HeadDim)
+	v = mlx.Reshape(v, B, L, attn.NHeads, attn.HeadDim)
+
+	// QK norm
+	q = mlx.RMSNorm(q, attn.NormQ, 1e-5)
+	k = mlx.RMSNorm(k, attn.NormK, 1e-5)
+
+	// Apply RoPE if provided
+	if cos != nil && sin != nil {
+		q = applyRoPE3D(q, cos, sin)
+		k = applyRoPE3D(k, cos, sin)
+	}
+
+	// Transpose to [B, nheads, L, head_dim]
+	q = mlx.Transpose(q, 0, 2, 1, 3)
+	k = mlx.Transpose(k, 0, 2, 1, 3)
+	v = mlx.Transpose(v, 0, 2, 1, 3)
+
+	// SDPA
+	out := mlx.ScaledDotProductAttention(q, k, v, attn.Scale, false)
+
+	// Transpose back and reshape
+	out = mlx.Transpose(out, 0, 2, 1, 3)
+	out = mlx.Reshape(out, B*L, attn.Dim)
+	out = attn.ToOut.Forward(out)
+
+	return mlx.Reshape(out, B, L, attn.Dim)
+}
+
+// applyRoPE3D applies 3-axis rotary position embeddings
+// x: [B, L, nheads, head_dim]
+// cos, sin: [B, L, 1, head_dim/2]
+func applyRoPE3D(x *mlx.Array, cos, sin *mlx.Array) *mlx.Array {
+	shape := x.Shape()
+	B := shape[0]
+	L := shape[1]
+	nheads := shape[2]
+	headDim := shape[3]
+	half := headDim / 2
+
+	// Create even/odd index arrays
+	evenIdx := make([]int32, half)
+	oddIdx := make([]int32, half)
+	for i := int32(0); i < half; i++ {
+		evenIdx[i] = i * 2
+		oddIdx[i] = i*2 + 1
+	}
+	evenIndices := mlx.NewArrayInt32(evenIdx, []int32{half})
+	oddIndices := mlx.NewArrayInt32(oddIdx, []int32{half})
+
+	// Extract x1 (even indices) and x2 (odd indices) along last axis
+	x1 := mlx.Take(x, evenIndices, 3) // [B, L, nheads, half]
+	x2 := mlx.Take(x, oddIndices, 3)  // [B, L, nheads, half]
+
+	// Apply rotation: [x1*cos - x2*sin, x1*sin + x2*cos]
+	r1 := mlx.Sub(mlx.Mul(x1, cos), mlx.Mul(x2, sin))
+	r2 := mlx.Add(mlx.Mul(x1, sin), mlx.Mul(x2, cos))
+
+	// Stack and reshape to interleave: [r1_0, r2_0, r1_1, r2_1, ...]
+	r1 = mlx.ExpandDims(r1, 4)                          // [B, L, nheads, half, 1]
+	r2 = mlx.ExpandDims(r2, 4)                          // [B, L, nheads, half, 1]
+	stacked := mlx.Concatenate([]*mlx.Array{r1, r2}, 4) // [B, L, nheads, half, 2]
+	return mlx.Reshape(stacked, B, L, nheads, headDim)
+}
+
+// TransformerBlock is a single transformer block with optional AdaLN modulation
+type TransformerBlock struct {
+	Attention      *Attention     `weight:"attention"`
+	FeedForward    *FeedForward   `weight:"feed_forward"`
+	AttentionNorm1 *nn.RMSNorm    `weight:"attention_norm1"`
+	AttentionNorm2 *nn.RMSNorm    `weight:"attention_norm2"`
+	FFNNorm1       *nn.RMSNorm    `weight:"ffn_norm1"`
+	FFNNorm2       *nn.RMSNorm    `weight:"ffn_norm2"`
+	AdaLN          nn.LinearLayer `weight:"adaLN_modulation.0,optional"` // only if modulation
+	// Computed fields
+	HasModulation bool
+	Dim           int32
+}
+
+// Forward applies the transformer block
+func (tb *TransformerBlock) Forward(x *mlx.Array, adaln *mlx.Array, cos, sin *mlx.Array, eps float32) *mlx.Array {
+	if tb.AdaLN != nil && adaln != nil {
+		// Compute modulation: [B, 256] -> [B, 4*dim]
+		chunks := tb.AdaLN.Forward(adaln)
+
+		// Split into 4 parts: scale_msa, gate_msa, scale_mlp, gate_mlp
+		chunkShape := chunks.Shape()
+		chunkDim := chunkShape[1] / 4
+
+		scaleMSA := mlx.Slice(chunks, []int32{0, 0}, []int32{chunkShape[0], chunkDim})
+		gateMSA := mlx.Slice(chunks, []int32{0, chunkDim}, []int32{chunkShape[0], chunkDim * 2})
+		scaleMLP := mlx.Slice(chunks, []int32{0, chunkDim * 2}, []int32{chunkShape[0], chunkDim * 3})
+		gateMLP := mlx.Slice(chunks, []int32{0, chunkDim * 3}, []int32{chunkShape[0], chunkDim * 4})
+
+		// Expand for broadcasting: [B, 1, dim]
+		scaleMSA = mlx.ExpandDims(scaleMSA, 1)
+		gateMSA = mlx.ExpandDims(gateMSA, 1)
+		scaleMLP = mlx.ExpandDims(scaleMLP, 1)
+		gateMLP = mlx.ExpandDims(gateMLP, 1)
+
+		// Attention with modulation
+		normX := tb.AttentionNorm1.Forward(x, eps)
+		normX = mlx.Mul(normX, mlx.AddScalar(scaleMSA, 1.0))
+		attnOut := tb.Attention.Forward(normX, cos, sin)
+		attnOut = tb.AttentionNorm2.Forward(attnOut, eps)
+		x = mlx.Add(x, mlx.Mul(mlx.Tanh(gateMSA), attnOut))
+
+		// FFN with modulation
+		normFFN := tb.FFNNorm1.Forward(x, eps)
+		normFFN = mlx.Mul(normFFN, mlx.AddScalar(scaleMLP, 1.0))
+		ffnOut := tb.FeedForward.Forward(normFFN)
+		ffnOut = tb.FFNNorm2.Forward(ffnOut, eps)
+		x = mlx.Add(x, mlx.Mul(mlx.Tanh(gateMLP), ffnOut))
+	} else {
+		// No modulation (context refiner)
+		attnOut := tb.Attention.Forward(tb.AttentionNorm1.Forward(x, eps), cos, sin)
+		x = mlx.Add(x, tb.AttentionNorm2.Forward(attnOut, eps))
+
+		ffnOut := tb.FeedForward.Forward(tb.FFNNorm1.Forward(x, eps))
+		x = mlx.Add(x, tb.FFNNorm2.Forward(ffnOut, eps))
+	}
+
+	return x
+}
+
+// FinalLayer outputs the denoised patches
+type FinalLayer struct {
+	AdaLN  nn.LinearLayer `weight:"adaLN_modulation.1"` // [256] -> [dim]
+	Output nn.LinearLayer `weight:"linear"`             // [dim] -> [out_channels]
+	OutDim int32          // computed from Output
+}
+
+// Forward computes final output
+func (fl *FinalLayer) Forward(x *mlx.Array, c *mlx.Array) *mlx.Array {
+	// c: [B, 256] -> scale: [B, dim]
+	scale := mlx.SiLU(c)
+	scale = fl.AdaLN.Forward(scale)
+	scale = mlx.ExpandDims(scale, 1) // [B, 1, dim]
+
+	// LayerNorm (affine=False) then scale
+	x = layerNormNoAffine(x, 1e-6)
+	x = mlx.Mul(x, mlx.AddScalar(scale, 1.0))
+
+	// Output projection
+	shape := x.Shape()
+	B := shape[0]
+	L := shape[1]
+	D := shape[2]
+	x = mlx.Reshape(x, B*L, D)
+	x = fl.Output.Forward(x)
+
+	return mlx.Reshape(x, B, L, fl.OutDim)
+}
+
+// layerNormNoAffine applies layer norm without learnable parameters
+func layerNormNoAffine(x *mlx.Array, eps float32) *mlx.Array {
+	ndim := x.Ndim()
+	lastAxis := ndim - 1
+
+	mean := mlx.Mean(x, lastAxis, true)
+	xCentered := mlx.Sub(x, mean)
+	variance := mlx.Mean(mlx.Square(xCentered), lastAxis, true)
+	return mlx.Div(xCentered, mlx.Sqrt(mlx.AddScalar(variance, eps)))
+}
+
+// Transformer is the full Z-Image DiT model
+type Transformer struct {
+	TEmbed          *TimestepEmbedder   `weight:"t_embedder"`
+	XEmbed          *XEmbedder          `weight:"all_x_embedder"`
+	CapEmbed        *CapEmbedder        `weight:"cap_embedder"`
+	NoiseRefiners   []*TransformerBlock `weight:"noise_refiner"`
+	ContextRefiners []*TransformerBlock `weight:"context_refiner"`
+	Layers          []*TransformerBlock `weight:"layers"`
+	FinalLayer      *FinalLayer         `weight:"all_final_layer.2-1"`
+	XPadToken       *mlx.Array          `weight:"x_pad_token"`
+	CapPadToken     *mlx.Array          `weight:"cap_pad_token"`
+	*TransformerConfig
+}
+
+// Load loads the Z-Image transformer from ollama blob storage.
+func (m *Transformer) Load(modelManifest *manifest.ModelManifest) error {
+	fmt.Print("  Loading transformer... ")
+
+	// Load config from blob
+	var cfg TransformerConfig
+	if err := modelManifest.ReadConfigJSON("transformer/config.json", &cfg); err != nil {
+		return fmt.Errorf("config: %w", err)
+	}
+	if len(cfg.AllPatchSize) > 0 {
+		cfg.PatchSize = cfg.AllPatchSize[0]
+	}
+	m.TransformerConfig = &cfg
+	m.NoiseRefiners = make([]*TransformerBlock, cfg.NRefinerLayers)
+	m.ContextRefiners = make([]*TransformerBlock, cfg.NRefinerLayers)
+	m.Layers = make([]*TransformerBlock, cfg.NLayers)
+
+	weights, err := manifest.LoadWeightsFromManifest(modelManifest, "transformer")
+	if err != nil {
+		return fmt.Errorf("weights: %w", err)
+	}
+	if err := weights.Load(0); err != nil {
+		return fmt.Errorf("load weights: %w", err)
+	}
+	defer weights.ReleaseAll()
+
+	return m.loadWeights(weights)
+}
+
+// loadWeights loads weights from any WeightSource into the model
+func (m *Transformer) loadWeights(weights safetensors.WeightSource) error {
+	if err := safetensors.LoadModule(m, weights, ""); err != nil {
+		return fmt.Errorf("load module: %w", err)
+	}
+	m.initComputedFields()
+	fmt.Println("✓")
+	return nil
+}
+
+// initComputedFields initializes computed fields after loading weights
+func (m *Transformer) initComputedFields() {
+	cfg := m.TransformerConfig
+	m.TEmbed.FreqEmbedSize = 256
+	m.FinalLayer.OutDim = m.FinalLayer.Output.OutputDim()
+	m.CapEmbed.Norm.Eps = 1e-6
+
+	for _, block := range m.NoiseRefiners {
+		initTransformerBlock(block, cfg)
+	}
+	for _, block := range m.ContextRefiners {
+		initTransformerBlock(block, cfg)
+	}
+	for _, block := range m.Layers {
+		initTransformerBlock(block, cfg)
+	}
+}
+
+// FuseAllQKV fuses QKV projections in all attention layers for efficiency.
+// This reduces 3 matmuls to 1 per attention layer, providing ~5-10% speedup.
+func (m *Transformer) FuseAllQKV() {
+	for _, block := range m.NoiseRefiners {
+		block.Attention.FuseQKV()
+	}
+	for _, block := range m.ContextRefiners {
+		block.Attention.FuseQKV()
+	}
+	for _, block := range m.Layers {
+		block.Attention.FuseQKV()
+	}
+}
+
+// initTransformerBlock sets computed fields on a transformer block
+func initTransformerBlock(block *TransformerBlock, cfg *TransformerConfig) {
+	block.Dim = cfg.Dim
+	block.HasModulation = block.AdaLN != nil
+
+	// Init attention computed fields
+	attn := block.Attention
+	attn.NHeads = cfg.NHeads
+	attn.HeadDim = cfg.Dim / cfg.NHeads
+	attn.Dim = cfg.Dim
+	attn.Scale = float32(1.0 / math.Sqrt(float64(attn.HeadDim)))
+
+	// Init feedforward OutDim
+	block.FeedForward.OutDim = block.FeedForward.W2.OutputDim()
+
+	// Set eps on all RMSNorm layers
+	block.AttentionNorm1.Eps = cfg.NormEps
+	block.AttentionNorm2.Eps = cfg.NormEps
+	block.FFNNorm1.Eps = cfg.NormEps
+	block.FFNNorm2.Eps = cfg.NormEps
+}
+
+// RoPECache holds precomputed RoPE values
+type RoPECache struct {
+	ImgCos     *mlx.Array
+	ImgSin     *mlx.Array
+	CapCos     *mlx.Array
+	CapSin     *mlx.Array
+	UnifiedCos *mlx.Array
+	UnifiedSin *mlx.Array
+	ImgLen     int32
+	CapLen     int32
+	GridH      int32 // Image token grid height
+	GridW      int32 // Image token grid width
+}
+
+// PrepareRoPECache precomputes RoPE values for the given image and caption lengths.
+// hTok and wTok are the number of tokens in each dimension (latentH/patchSize, latentW/patchSize).
+func (m *Transformer) PrepareRoPECache(hTok, wTok, capLen int32) *RoPECache {
+	imgLen := hTok * wTok
+
+	// Image positions: grid over (1, H, W) starting at (capLen+1, 0, 0)
+	imgPos := createCoordinateGrid(1, hTok, wTok, capLen+1, 0, 0)
+	imgPos = mlx.ToBFloat16(imgPos)
+	// Caption positions: grid over (capLen, 1, 1) starting at (1, 0, 0)
+	capPos := createCoordinateGrid(capLen, 1, 1, 1, 0, 0)
+	capPos = mlx.ToBFloat16(capPos)
+
+	// Compute RoPE from UNIFIED positions
+	unifiedPos := mlx.Concatenate([]*mlx.Array{imgPos, capPos}, 1)
+	unifiedCos, unifiedSin := prepareRoPE3D(unifiedPos, m.TransformerConfig.AxesDims)
+
+	// Slice RoPE for image and caption parts
+	imgCos := mlx.Slice(unifiedCos, []int32{0, 0, 0, 0}, []int32{1, imgLen, 1, 64})
+	imgSin := mlx.Slice(unifiedSin, []int32{0, 0, 0, 0}, []int32{1, imgLen, 1, 64})
+	capCos := mlx.Slice(unifiedCos, []int32{0, imgLen, 0, 0}, []int32{1, imgLen + capLen, 1, 64})
+	capSin := mlx.Slice(unifiedSin, []int32{0, imgLen, 0, 0}, []int32{1, imgLen + capLen, 1, 64})
+
+	return &RoPECache{
+		ImgCos:     imgCos,
+		ImgSin:     imgSin,
+		CapCos:     capCos,
+		CapSin:     capSin,
+		UnifiedCos: unifiedCos,
+		UnifiedSin: unifiedSin,
+		ImgLen:     imgLen,
+		CapLen:     capLen,
+		GridH:      hTok,
+		GridW:      wTok,
+	}
+}
+
+// Forward runs the Z-Image transformer with precomputed RoPE
+func (m *Transformer) Forward(x *mlx.Array, t *mlx.Array, capFeats *mlx.Array, rope *RoPECache) *mlx.Array {
+	imgLen := rope.ImgLen
+
+	// Timestep embedding -> [B, 256]
+	temb := m.TEmbed.Forward(mlx.MulScalar(t, m.TransformerConfig.TScale))
+
+	// Embed image patches -> [B, L_img, dim]
+	x = m.XEmbed.Forward(x)
+
+	// Embed caption features -> [B, L_cap, dim]
+	capEmb := m.CapEmbed.Forward(capFeats)
+
+	eps := m.NormEps
+
+	// Noise refiner: refine image patches with modulation
+	for _, refiner := range m.NoiseRefiners {
+		x = refiner.Forward(x, temb, rope.ImgCos, rope.ImgSin, eps)
+	}
+
+	// Context refiner: refine caption (no modulation)
+	for _, refiner := range m.ContextRefiners {
+		capEmb = refiner.Forward(capEmb, nil, rope.CapCos, rope.CapSin, eps)
+	}
+
+	// Concatenate image and caption for joint attention
+	unified := mlx.Concatenate([]*mlx.Array{x, capEmb}, 1)
+
+	// Main transformer layers use full unified RoPE
+	for _, layer := range m.Layers {
+		unified = layer.Forward(unified, temb, rope.UnifiedCos, rope.UnifiedSin, eps)
+	}
+
+	// Extract image tokens only
+	unifiedShape := unified.Shape()
+	B := unifiedShape[0]
+	imgOut := mlx.Slice(unified, []int32{0, 0, 0}, []int32{B, imgLen, unifiedShape[2]})
+
+	// Final layer
+	return m.FinalLayer.Forward(imgOut, temb)
+}
+
+// ForwardWithCache runs the transformer with layer caching for faster inference.
+// On refresh steps (step % cacheInterval == 0), all layers are computed and cached.
+// On other steps, shallow layers (0 to cacheLayers-1) reuse cached outputs.
+func (m *Transformer) ForwardWithCache(
+	x *mlx.Array,
+	t *mlx.Array,
+	capFeats *mlx.Array,
+	rope *RoPECache,
+	stepCache *cache.StepCache,
+	step int,
+	cacheInterval int,
+) *mlx.Array {
+	imgLen := rope.ImgLen
+	cacheLayers := stepCache.NumLayers()
+	eps := m.NormEps
+
+	// Timestep embedding -> [B, 256]
+	temb := m.TEmbed.Forward(mlx.MulScalar(t, m.TransformerConfig.TScale))
+
+	// Embed image patches -> [B, L_img, dim]
+	x = m.XEmbed.Forward(x)
+
+	// Context refiners: compute once on step 0, reuse forever
+	// (caption embedding doesn't depend on timestep or latents)
+	var capEmb *mlx.Array
+	if stepCache.GetConstant() != nil {
+		capEmb = stepCache.GetConstant()
+	} else {
+		capEmb = m.CapEmbed.Forward(capFeats)
+		for _, refiner := range m.ContextRefiners {
+			capEmb = refiner.Forward(capEmb, nil, rope.CapCos, rope.CapSin, eps)
+		}
+		stepCache.SetConstant(capEmb)
+	}
+
+	// Noise refiners: always compute (depend on x which changes each step)
+	for _, refiner := range m.NoiseRefiners {
+		x = refiner.Forward(x, temb, rope.ImgCos, rope.ImgSin, eps)
+	}
+
+	// Concatenate image and caption for joint attention
+	unified := mlx.Concatenate([]*mlx.Array{x, capEmb}, 1)
+
+	// Determine if this is a cache refresh step
+	refreshCache := stepCache.ShouldRefresh(step, cacheInterval)
+
+	// Main transformer layers with caching
+	for i, layer := range m.Layers {
+		if i < cacheLayers && !refreshCache && stepCache.Get(i) != nil {
+			// Use cached output for shallow layers
+			unified = stepCache.Get(i)
+		} else {
+			// Compute layer
+			unified = layer.Forward(unified, temb, rope.UnifiedCos, rope.UnifiedSin, eps)
+			// Cache shallow layer outputs on refresh steps
+			if i < cacheLayers && refreshCache {
+				stepCache.Set(i, unified)
+			}
+		}
+	}
+
+	// Extract image tokens only
+	unifiedShape := unified.Shape()
+	B := unifiedShape[0]
+	imgOut := mlx.Slice(unified, []int32{0, 0, 0}, []int32{B, imgLen, unifiedShape[2]})
+
+	// Final layer
+	return m.FinalLayer.Forward(imgOut, temb)
+}
+
+// createCoordinateGrid creates 3D position grid [1, d0*d1*d2, 3]
+func createCoordinateGrid(d0, d1, d2, s0, s1, s2 int32) *mlx.Array {
+	// Create meshgrid and stack
+	total := d0 * d1 * d2
+	coords := make([]float32, total*3)
+
+	idx := 0
+	for i := int32(0); i < d0; i++ {
+		for j := int32(0); j < d1; j++ {
+			for k := int32(0); k < d2; k++ {
+				coords[idx*3+0] = float32(s0 + i)
+				coords[idx*3+1] = float32(s1 + j)
+				coords[idx*3+2] = float32(s2 + k)
+				idx++
+			}
+		}
+	}
+
+	return mlx.NewArray(coords, []int32{1, total, 3})
+}
+
+// prepareRoPE3D computes cos/sin for 3-axis RoPE
+// positions: [B, L, 3] with (h, w, t) coordinates
+// axesDims: [32, 48, 48] - dimensions for each axis
+// Returns: cos, sin each [B, L, 1, head_dim/2]
+func prepareRoPE3D(positions *mlx.Array, axesDims []int32) (*mlx.Array, *mlx.Array) {
+	// Compute frequencies for each axis
+	// dims = [32, 48, 48], so halves = [16, 24, 24]
+	ropeTheta := float32(256.0)
+
+	freqs := make([]*mlx.Array, 3)
+	for axis := 0; axis < 3; axis++ {
+		half := axesDims[axis] / 2
+		f := make([]float32, half)
+		for i := int32(0); i < half; i++ {
+			f[i] = float32(math.Exp(-math.Log(float64(ropeTheta)) * float64(i) / float64(half)))
+		}
+		freqs[axis] = mlx.NewArray(f, []int32{1, 1, 1, half})
+	}
+
+	// Extract position coordinates
+	shape := positions.Shape()
+	B := shape[0]
+	L := shape[1]
+
+	// positions[:, :, 0] -> h positions
+	posH := mlx.Slice(positions, []int32{0, 0, 0}, []int32{B, L, 1})
+	posW := mlx.Slice(positions, []int32{0, 0, 1}, []int32{B, L, 2})
+	posT := mlx.Slice(positions, []int32{0, 0, 2}, []int32{B, L, 3})
+
+	// Compute args: pos * freqs for each axis
+	posH = mlx.ExpandDims(posH, 3) // [B, L, 1, 1]
+	posW = mlx.ExpandDims(posW, 3)
+	posT = mlx.ExpandDims(posT, 3)
+
+	argsH := mlx.Mul(posH, freqs[0]) // [B, L, 1, 16]
+	argsW := mlx.Mul(posW, freqs[1]) // [B, L, 1, 24]
+	argsT := mlx.Mul(posT, freqs[2]) // [B, L, 1, 24]
+
+	// Concatenate: [B, L, 1, 16+24+24=64]
+	args := mlx.Concatenate([]*mlx.Array{argsH, argsW, argsT}, 3)
+
+	// Compute cos and sin
+	return mlx.Cos(args), mlx.Sin(args)
+}
+
+// PatchifyLatents converts latents [B, C, H, W] to patches [B, L, C*patch^2]
+// Matches Python: x.reshape(C, 1, 1, H_tok, 2, W_tok, 2).transpose(1,2,3,5,4,6,0).reshape(1,-1,C*4)
+func PatchifyLatents(latents *mlx.Array, patchSize int32) *mlx.Array {
+	shape := latents.Shape()
+	C := shape[1]
+	H := shape[2]
+	W := shape[3]
+
+	pH := H / patchSize // H_tok
+	pW := W / patchSize // W_tok
+
+	// Match Python exactly: reshape treating B=1 as part of contiguous data
+	// [1, C, H, W] -> [C, 1, 1, pH, 2, pW, 2]
+	x := mlx.Reshape(latents, C, 1, 1, pH, patchSize, pW, patchSize)
+
+	// Python: transpose(1, 2, 3, 5, 4, 6, 0)
+	// [C, 1, 1, pH, 2, pW, 2] -> [1, 1, pH, pW, 2, 2, C]
+	x = mlx.Transpose(x, 1, 2, 3, 5, 4, 6, 0)
+
+	// [1, 1, pH, pW, 2, 2, C] -> [1, pH*pW, C*4]
+	return mlx.Reshape(x, 1, pH*pW, C*patchSize*patchSize)
+}
+
+// UnpatchifyLatents converts patches [B, L, C*patch^2] back to [B, C, H, W]
+// Matches Python: out.reshape(1,1,H_tok,W_tok,2,2,C).transpose(6,0,1,2,4,3,5).reshape(1,C,H,W)
+func UnpatchifyLatents(patches *mlx.Array, patchSize, H, W, C int32) *mlx.Array {
+	pH := H / patchSize
+	pW := W / patchSize
+
+	// [1, L, C*4] -> [1, 1, pH, pW, 2, 2, C]
+	x := mlx.Reshape(patches, 1, 1, pH, pW, patchSize, patchSize, C)
+
+	// Python: transpose(6, 0, 1, 2, 4, 3, 5)
+	// [1, 1, pH, pW, 2, 2, C] -> [C, 1, 1, pH, 2, pW, 2]
+	x = mlx.Transpose(x, 6, 0, 1, 2, 4, 3, 5)
+
+	// [C, 1, 1, pH, 2, pW, 2] -> [1, C, H, W]
+	return mlx.Reshape(x, 1, C, H, W)
+}

x/imagegen/models/zimage/vae.go

@@ -0,0 +1,820 @@
+package zimage
+
+import (
+	"fmt"
+	"math"
+
+	"github.com/ollama/ollama/x/imagegen/manifest"
+	"github.com/ollama/ollama/x/imagegen/mlx"
+	"github.com/ollama/ollama/x/imagegen/safetensors"
+	"github.com/ollama/ollama/x/imagegen/vae"
+)
+
+// VAEConfig holds VAE decoder configuration
+type VAEConfig struct {
+	InChannels       int32   `json:"in_channels"`
+	OutChannels      int32   `json:"out_channels"`
+	LatentChannels   int32   `json:"latent_channels"`
+	BlockOutChannels []int32 `json:"block_out_channels"`
+	LayersPerBlock   int32   `json:"layers_per_block"`
+	NormNumGroups    int32   `json:"norm_num_groups"`
+	ScalingFactor    float32 `json:"scaling_factor"`
+	ShiftFactor      float32 `json:"shift_factor"`
+}
+
+// GroupNormLayer implements group normalization
+type GroupNormLayer struct {
+	Weight    *mlx.Array
+	Bias      *mlx.Array
+	NumGroups int32
+	Eps       float32
+}
+
+// NewGroupNorm creates a group norm layer
+func NewGroupNorm(weight, bias *mlx.Array, numGroups int32) *GroupNormLayer {
+	return &GroupNormLayer{
+		Weight:    weight,
+		Bias:      bias,
+		NumGroups: numGroups,
+		Eps:       1e-5,
+	}
+}
+
+// Forward applies group normalization
+// Input and output are in NHWC format [B, H, W, C]
+func (gn *GroupNormLayer) Forward(x *mlx.Array) *mlx.Array {
+	// x: [B, H, W, C] (NHWC format)
+	shape := x.Shape()
+	B := shape[0]
+	H := shape[1]
+	W := shape[2]
+	C := shape[3]
+
+	// For large spatial sizes, use tiled computation to avoid CUDA grid limits
+	// CUDA grid.y max is 65535, so H*W/16 must be <= 65535, meaning H*W <= ~1M
+	// To be safe, tile when H*W > 512*512 = 262144
+	if H*W > 512*512 {
+		return gn.forwardTiled(x, B, H, W, C)
+	}
+
+	return gn.forwardSmall(x, B, H, W, C)
+}
+
+// forwardSmall is the standard GroupNorm for tensors that fit within CUDA grid limits
+func (gn *GroupNormLayer) forwardSmall(x *mlx.Array, B, H, W, C int32) *mlx.Array {
+	// Reshape to [B, H, W, groups, C/groups]
+	groupSize := C / gn.NumGroups
+	x = mlx.Reshape(x, B, H, W, gn.NumGroups, groupSize)
+
+	// Compute mean and variance per group (over H, W, and C/groups dimensions)
+	mean := mlx.Mean(x, 1, true)
+	mean = mlx.Mean(mean, 2, true)
+	mean = mlx.Mean(mean, 4, true)
+
+	xCentered := mlx.Sub(x, mean)
+
+	// Variance over same axes
+	sq := mlx.Square(xCentered)
+	variance := mlx.Mean(sq, 1, true)
+	variance = mlx.Mean(variance, 2, true)
+	variance = mlx.Mean(variance, 4, true)
+
+	// Normalize
+	xNorm := mlx.Div(xCentered, mlx.Sqrt(mlx.AddScalar(variance, gn.Eps)))
+
+	// Reshape back to [B, H, W, C]
+	xNorm = mlx.Reshape(xNorm, B, H, W, C)
+
+	// Scale and shift (weight and bias are [C])
+	if gn.Weight != nil {
+		weight := mlx.Reshape(gn.Weight, 1, 1, 1, C)
+		xNorm = mlx.Mul(xNorm, weight)
+	}
+	if gn.Bias != nil {
+		bias := mlx.Reshape(gn.Bias, 1, 1, 1, C)
+		xNorm = mlx.Add(xNorm, bias)
+	}
+
+	return xNorm
+}
+
+// forwardTiled handles large tensors by processing in H-tiles to avoid CUDA grid limits
+func (gn *GroupNormLayer) forwardTiled(x *mlx.Array, B, H, W, C int32) *mlx.Array {
+	groupSize := C / gn.NumGroups
+
+	// Keep the input - we need it for slicing tiles later
+	// Track if we were the ones who kept it, so we can restore state after
+	wasKept := x.Kept()
+	mlx.Keep(x)
+
+	// Compute per-group mean and variance using flattened spatial dimensions
+	// Build the entire compute graph first, then eval once
+	// Reshape to [B, H*W, groups, groupSize]
+	xFlat := mlx.Reshape(x, B, H*W, gn.NumGroups, groupSize)
+
+	// Mean over spatial (axis 1) and groupSize (axis 3) dimensions
+	// Result shape: [B, 1, groups, 1]
+	mean1 := mlx.Mean(xFlat, 1, true)
+	mean := mlx.Mean(mean1, 3, true)
+
+	// Variance using E[X^2] - E[X]^2
+	xSq := mlx.Square(xFlat)
+	meanSq1 := mlx.Mean(xSq, 1, true)
+	meanSq := mlx.Mean(meanSq1, 3, true)
+	meanSquared := mlx.Square(mean)
+	variance := mlx.Sub(meanSq, meanSquared)
+
+	// invStd = 1/sqrt(var + eps)
+	varPlusEps := mlx.AddScalar(variance, gn.Eps)
+	stdDev := mlx.Sqrt(varPlusEps)
+	one := mlx.Full(1.0, 1)
+	invStd := mlx.Div(one, stdDev)
+
+	// Eval mean and invStd together - these are what we need for the tile loop
+	mlx.Keep(mean, invStd)
+	mlx.Eval(mean, invStd)
+
+	// Tile along H dimension
+	tileH := int32(512 * 512 / W)
+	if tileH < 1 {
+		tileH = 1
+	}
+	if tileH > H {
+		tileH = H
+	}
+
+	// Prepare weight and bias reshaped for 4D broadcast [1, 1, groups, groupSize]
+	var weightGN, biasGN *mlx.Array
+	if gn.Weight != nil {
+		weightGN = mlx.Reshape(gn.Weight, 1, 1, gn.NumGroups, groupSize)
+		mlx.Keep(weightGN)
+		mlx.Eval(weightGN)
+	}
+	if gn.Bias != nil {
+		biasGN = mlx.Reshape(gn.Bias, 1, 1, gn.NumGroups, groupSize)
+		mlx.Keep(biasGN)
+		mlx.Eval(biasGN)
+	}
+
+	var tiles []*mlx.Array
+	for hStart := int32(0); hStart < H; hStart += tileH {
+		hEnd := hStart + tileH
+		if hEnd > H {
+			hEnd = H
+		}
+		tileHeight := hEnd - hStart
+		spatialSize := tileHeight * W
+
+		// Build the compute graph for this tile (no intermediate Evals)
+		// Extract tile and flatten spatial dims: [B, tileH*W, groups, groupSize]
+		tile := mlx.Slice(x, []int32{0, hStart, 0, 0}, []int32{B, hEnd, W, C})
+		tileFlat := mlx.Reshape(tile, B, spatialSize, gn.NumGroups, groupSize)
+
+		// Normalize: (x - mean) * invStd
+		tileCentered := mlx.Sub(tileFlat, mean)
+		tileNorm := mlx.Mul(tileCentered, invStd)
+
+		// Apply scale and shift in 4D space
+		if weightGN != nil {
+			tileNorm = mlx.Mul(tileNorm, weightGN)
+		}
+		if biasGN != nil {
+			tileNorm = mlx.Add(tileNorm, biasGN)
+		}
+
+		// Reshape back to [B, tileH, W, C]
+		tileOut := mlx.Reshape(tileNorm, B, tileHeight, W, C)
+
+		// Now eval and keep this tile
+		mlx.Keep(tileOut)
+		mlx.Eval(tileOut)
+
+		tiles = append(tiles, tileOut)
+	}
+
+	// Concatenate tiles along H axis
+	var result *mlx.Array
+	if len(tiles) == 1 {
+		result = tiles[0]
+	} else {
+		result = mlx.Concatenate(tiles, 1)
+		mlx.Eval(result)
+		// Free the individual tiles now that they're concatenated
+		for _, t := range tiles {
+			t.Free()
+		}
+	}
+
+	// Clean up kept arrays
+	// Restore x's kept state - only free if we were the ones who kept it
+	if !wasKept {
+		x.Free()
+	}
+	mean.Free()
+	invStd.Free()
+	if weightGN != nil {
+		weightGN.Free()
+	}
+	if biasGN != nil {
+		biasGN.Free()
+	}
+
+	return result
+}
+
+// Conv2D represents a 2D convolution layer
+// Works natively in NHWC format (MLX's native format)
+type Conv2D struct {
+	Weight  *mlx.Array // [out_channels, kH, kW, in_channels] (OHWI for MLX)
+	Bias    *mlx.Array // [out_channels]
+	Stride  int32
+	Padding int32
+}
+
+// NewConv2D creates a Conv2D layer
+// weight comes in as [out_channels, in_channels, kH, kW] (OIHW from PyTorch)
+// we transpose to [out_channels, kH, kW, in_channels] (OHWI for MLX)
+func NewConv2D(weight, bias *mlx.Array, stride, padding int32) *Conv2D {
+	// Transpose weight from OIHW to OHWI
+	// [O, I, H, W] -> [O, H, W, I]
+	weightOHWI := mlx.Transpose(weight, 0, 2, 3, 1)
+	return &Conv2D{
+		Weight:  weightOHWI,
+		Bias:    bias,
+		Stride:  stride,
+		Padding: padding,
+	}
+}
+
+// Forward applies convolution
+// Input and output are in NHWC format [N, H, W, C]
+func (conv *Conv2D) Forward(x *mlx.Array) *mlx.Array {
+	// Conv in NHWC format (MLX native)
+	out := mlx.Conv2d(x, conv.Weight, conv.Stride, conv.Padding)
+
+	if conv.Bias != nil {
+		// Bias is [C], reshape to [1, 1, 1, C] for NHWC broadcast
+		bias := mlx.Reshape(conv.Bias, 1, 1, 1, conv.Bias.Dim(0))
+		out = mlx.Add(out, bias)
+	}
+
+	return out
+}
+
+// ResnetBlock2D implements a ResNet block for VAE
+type ResnetBlock2D struct {
+	Norm1        *GroupNormLayer
+	Conv1        *Conv2D
+	Norm2        *GroupNormLayer
+	Conv2        *Conv2D
+	ConvShortcut *Conv2D // nil if in_channels == out_channels
+}
+
+// NewResnetBlock2D creates a ResNet block
+func NewResnetBlock2D(weights safetensors.WeightSource, prefix string, numGroups int32) (*ResnetBlock2D, error) {
+	norm1Weight, err := weights.GetTensor(prefix + ".norm1.weight")
+	if err != nil {
+		return nil, err
+	}
+	norm1Bias, err := weights.GetTensor(prefix + ".norm1.bias")
+	if err != nil {
+		return nil, err
+	}
+
+	conv1Weight, err := weights.GetTensor(prefix + ".conv1.weight")
+	if err != nil {
+		return nil, err
+	}
+	conv1Bias, err := weights.GetTensor(prefix + ".conv1.bias")
+	if err != nil {
+		return nil, err
+	}
+
+	norm2Weight, err := weights.GetTensor(prefix + ".norm2.weight")
+	if err != nil {
+		return nil, err
+	}
+	norm2Bias, err := weights.GetTensor(prefix + ".norm2.bias")
+	if err != nil {
+		return nil, err
+	}
+
+	conv2Weight, err := weights.GetTensor(prefix + ".conv2.weight")
+	if err != nil {
+		return nil, err
+	}
+	conv2Bias, err := weights.GetTensor(prefix + ".conv2.bias")
+	if err != nil {
+		return nil, err
+	}
+
+	block := &ResnetBlock2D{
+		Norm1: NewGroupNorm(norm1Weight, norm1Bias, numGroups),
+		Conv1: NewConv2D(conv1Weight, conv1Bias, 1, 1),
+		Norm2: NewGroupNorm(norm2Weight, norm2Bias, numGroups),
+		Conv2: NewConv2D(conv2Weight, conv2Bias, 1, 1),
+	}
+
+	if weights.HasTensor(prefix + ".conv_shortcut.weight") {
+		shortcutWeight, err := weights.GetTensor(prefix + ".conv_shortcut.weight")
+		if err != nil {
+			return nil, err
+		}
+		shortcutBias, err := weights.GetTensor(prefix + ".conv_shortcut.bias")
+		if err != nil {
+			return nil, err
+		}
+		block.ConvShortcut = NewConv2D(shortcutWeight, shortcutBias, 1, 0)
+	}
+
+	return block, nil
+}
+
+// Forward applies the ResNet block with staged evaluation
+func (rb *ResnetBlock2D) Forward(x *mlx.Array) *mlx.Array {
+	var h *mlx.Array
+
+	// Stage 1: norm1
+	{
+		h = rb.Norm1.Forward(x)
+		mlx.Eval(h)
+	}
+
+	// Stage 2: silu + conv1
+	{
+		prev := h
+		h = mlx.SiLU(h)
+		h = rb.Conv1.Forward(h)
+		prev.Free()
+		mlx.Eval(h)
+	}
+
+	// Stage 3: norm2
+	{
+		prev := h
+		h = rb.Norm2.Forward(h)
+		prev.Free()
+		mlx.Eval(h)
+	}
+
+	// Stage 4: silu + conv2
+	{
+		prev := h
+		h = mlx.SiLU(h)
+		h = rb.Conv2.Forward(h)
+		prev.Free()
+		mlx.Eval(h)
+	}
+
+	// Residual connection
+	{
+		prev := h
+		if rb.ConvShortcut != nil {
+			shortcut := rb.ConvShortcut.Forward(x)
+			h = mlx.Add(h, shortcut)
+		} else {
+			h = mlx.Add(h, x)
+		}
+		prev.Free()
+		mlx.Eval(h)
+	}
+
+	return h
+}
+
+// VAEAttentionBlock implements self-attention for VAE
+type VAEAttentionBlock struct {
+	GroupNorm   *GroupNormLayer
+	ToQWeight   *mlx.Array
+	ToQBias     *mlx.Array
+	ToKWeight   *mlx.Array
+	ToKBias     *mlx.Array
+	ToVWeight   *mlx.Array
+	ToVBias     *mlx.Array
+	ToOutWeight *mlx.Array
+	ToOutBias   *mlx.Array
+	NumHeads    int32
+}
+
+// NewVAEAttentionBlock creates an attention block
+func NewVAEAttentionBlock(weights safetensors.WeightSource, prefix string, numGroups int32) (*VAEAttentionBlock, error) {
+	normWeight, err := weights.GetTensor(prefix + ".group_norm.weight")
+	if err != nil {
+		return nil, err
+	}
+	normBias, err := weights.GetTensor(prefix + ".group_norm.bias")
+	if err != nil {
+		return nil, err
+	}
+
+	toQWeight, err := weights.GetTensor(prefix + ".to_q.weight")
+	if err != nil {
+		return nil, err
+	}
+	toQBias, err := weights.GetTensor(prefix + ".to_q.bias")
+	if err != nil {
+		return nil, err
+	}
+
+	toKWeight, err := weights.GetTensor(prefix + ".to_k.weight")
+	if err != nil {
+		return nil, err
+	}
+	toKBias, err := weights.GetTensor(prefix + ".to_k.bias")
+	if err != nil {
+		return nil, err
+	}
+
+	toVWeight, err := weights.GetTensor(prefix + ".to_v.weight")
+	if err != nil {
+		return nil, err
+	}
+	toVBias, err := weights.GetTensor(prefix + ".to_v.bias")
+	if err != nil {
+		return nil, err
+	}
+
+	toOutWeight, err := weights.GetTensor(prefix + ".to_out.0.weight")
+	if err != nil {
+		return nil, err
+	}
+	toOutBias, err := weights.GetTensor(prefix + ".to_out.0.bias")
+	if err != nil {
+		return nil, err
+	}
+
+	return &VAEAttentionBlock{
+		GroupNorm:   NewGroupNorm(normWeight, normBias, numGroups),
+		ToQWeight:   mlx.Transpose(toQWeight, 1, 0),
+		ToQBias:     toQBias,
+		ToKWeight:   mlx.Transpose(toKWeight, 1, 0),
+		ToKBias:     toKBias,
+		ToVWeight:   mlx.Transpose(toVWeight, 1, 0),
+		ToVBias:     toVBias,
+		ToOutWeight: mlx.Transpose(toOutWeight, 1, 0),
+		ToOutBias:   toOutBias,
+		NumHeads:    1,
+	}, nil
+}
+
+// Forward applies attention with staged evaluation
+// Input and output are in NHWC format [B, H, W, C]
+func (ab *VAEAttentionBlock) Forward(x *mlx.Array) *mlx.Array {
+	residual := x
+	shape := x.Shape()
+	B := shape[0]
+	H := shape[1]
+	W := shape[2]
+	C := shape[3]
+
+	var h *mlx.Array
+
+	// Stage 1: GroupNorm + reshape to [B, H*W, C]
+	{
+		h = ab.GroupNorm.Forward(x)
+		h = mlx.Reshape(h, B, H*W, C)
+		mlx.Eval(h)
+	}
+
+	var out *mlx.Array
+
+	// Stage 2: Q, K, V projections + attention
+	{
+		q := mlx.Linear(h, ab.ToQWeight)
+		q = mlx.Add(q, ab.ToQBias)
+		k := mlx.Linear(h, ab.ToKWeight)
+		k = mlx.Add(k, ab.ToKBias)
+		v := mlx.Linear(h, ab.ToVWeight)
+		v = mlx.Add(v, ab.ToVBias)
+		h.Free()
+
+		q = mlx.ExpandDims(q, 1)
+		k = mlx.ExpandDims(k, 1)
+		v = mlx.ExpandDims(v, 1)
+
+		scale := float32(1.0 / math.Sqrt(float64(C)))
+		out = mlx.ScaledDotProductAttention(q, k, v, scale, false)
+		out = mlx.Squeeze(out, 1)
+		mlx.Eval(out)
+	}
+
+	// Stage 3: Output projection + reshape + residual
+	{
+		prev := out
+		out = mlx.Linear(out, ab.ToOutWeight)
+		out = mlx.Add(out, ab.ToOutBias)
+		out = mlx.Reshape(out, B, H, W, C)
+		out = mlx.Add(out, residual)
+		prev.Free()
+		mlx.Eval(out)
+	}
+
+	return out
+}
+
+// UpDecoderBlock2D implements an upsampling decoder block
+type UpDecoderBlock2D struct {
+	ResnetBlocks []*ResnetBlock2D
+	Upsample     *Conv2D
+}
+
+// NewUpDecoderBlock2D creates an up decoder block
+func NewUpDecoderBlock2D(weights safetensors.WeightSource, prefix string, numLayers, numGroups int32, hasUpsample bool) (*UpDecoderBlock2D, error) {
+	resnets := make([]*ResnetBlock2D, numLayers)
+	for i := int32(0); i < numLayers; i++ {
+		resPrefix := fmt.Sprintf("%s.resnets.%d", prefix, i)
+		resnet, err := NewResnetBlock2D(weights, resPrefix, numGroups)
+		if err != nil {
+			return nil, err
+		}
+		resnets[i] = resnet
+	}
+
+	var upsample *Conv2D
+	if hasUpsample {
+		upWeight, err := weights.GetTensor(prefix + ".upsamplers.0.conv.weight")
+		if err != nil {
+			return nil, err
+		}
+		upBias, err := weights.GetTensor(prefix + ".upsamplers.0.conv.bias")
+		if err != nil {
+			return nil, err
+		}
+		upsample = NewConv2D(upWeight, upBias, 1, 1)
+	}
+
+	return &UpDecoderBlock2D{
+		ResnetBlocks: resnets,
+		Upsample:     upsample,
+	}, nil
+}
+
+// Forward applies the up decoder block with staged evaluation to reduce peak memory
+func (ub *UpDecoderBlock2D) Forward(x *mlx.Array) *mlx.Array {
+	for _, resnet := range ub.ResnetBlocks {
+		prev := x
+		x = resnet.Forward(x) // ResNet handles its own pools
+		prev.Free()
+	}
+
+	if ub.Upsample != nil {
+		// Stage 1: Upsample2x (nearest neighbor)
+		{
+			prev := x
+			x = Upsample2x(x)
+			prev.Free()
+			mlx.Eval(x)
+		}
+
+		// Stage 2: Upsample conv
+		{
+			prev := x
+			x = ub.Upsample.Forward(x)
+			prev.Free()
+			mlx.Eval(x)
+		}
+	}
+
+	return x
+}
+
+// VAEMidBlock is the middle block with attention
+type VAEMidBlock struct {
+	Resnet1   *ResnetBlock2D
+	Attention *VAEAttentionBlock
+	Resnet2   *ResnetBlock2D
+}
+
+// NewVAEMidBlock creates the mid block
+func NewVAEMidBlock(weights safetensors.WeightSource, prefix string, numGroups int32) (*VAEMidBlock, error) {
+	resnet1, err := NewResnetBlock2D(weights, prefix+".resnets.0", numGroups)
+	if err != nil {
+		return nil, err
+	}
+
+	attention, err := NewVAEAttentionBlock(weights, prefix+".attentions.0", numGroups)
+	if err != nil {
+		return nil, err
+	}
+
+	resnet2, err := NewResnetBlock2D(weights, prefix+".resnets.1", numGroups)
+	if err != nil {
+		return nil, err
+	}
+
+	return &VAEMidBlock{
+		Resnet1:   resnet1,
+		Attention: attention,
+		Resnet2:   resnet2,
+	}, nil
+}
+
+// Forward applies the mid block with staged evaluation
+func (mb *VAEMidBlock) Forward(x *mlx.Array) *mlx.Array {
+	prev := x
+	x = mb.Resnet1.Forward(x) // ResNet handles its own pools
+	prev.Free()
+
+	// Attention handles its own pools
+	prev = x
+	x = mb.Attention.Forward(x)
+	prev.Free()
+
+	prev = x
+	x = mb.Resnet2.Forward(x) // ResNet handles its own pools
+	prev.Free()
+
+	return x
+}
+
+// VAEDecoder is the full VAE decoder
+type VAEDecoder struct {
+	Config      *VAEConfig
+	ConvIn      *Conv2D
+	MidBlock    *VAEMidBlock
+	UpBlocks    []*UpDecoderBlock2D
+	ConvNormOut *GroupNormLayer
+	ConvOut     *Conv2D
+
+	// Tiling configuration (nil = no tiling)
+	Tiling *vae.TilingConfig
+}
+
+// Load loads the VAE decoder from ollama blob storage.
+func (m *VAEDecoder) Load(modelManifest *manifest.ModelManifest) error {
+	// Load config from blob
+	var cfg VAEConfig
+	if err := modelManifest.ReadConfigJSON("vae/config.json", &cfg); err != nil {
+		return fmt.Errorf("config: %w", err)
+	}
+	m.Config = &cfg
+
+	// Load weights from tensor blobs
+	weights, err := manifest.LoadWeightsFromManifest(modelManifest, "vae")
+	if err != nil {
+		return fmt.Errorf("weights: %w", err)
+	}
+	if err := weights.Load(0); err != nil {
+		return fmt.Errorf("load weights: %w", err)
+	}
+	defer weights.ReleaseAll()
+
+	return m.loadWeights(weights, &cfg)
+}
+
+// loadWeights loads VAE weights from any WeightSource
+func (m *VAEDecoder) loadWeights(weights safetensors.WeightSource, cfg *VAEConfig) error {
+	var err error
+
+	// Load conv_in
+	fmt.Print("  Loading conv_in... ")
+	convInWeight, err := weights.GetTensor("decoder.conv_in.weight")
+	if err != nil {
+		return err
+	}
+	convInBias, err := weights.GetTensor("decoder.conv_in.bias")
+	if err != nil {
+		return err
+	}
+	m.ConvIn = NewConv2D(convInWeight, convInBias, 1, 1)
+	fmt.Println("✓")
+
+	// Load mid block
+	fmt.Print("  Loading mid block... ")
+	m.MidBlock, err = NewVAEMidBlock(weights, "decoder.mid_block", cfg.NormNumGroups)
+	if err != nil {
+		return err
+	}
+	fmt.Println("✓")
+
+	// Load up blocks
+	fmt.Print("  Loading up blocks... ")
+	numBlocks := len(cfg.BlockOutChannels)
+	m.UpBlocks = make([]*UpDecoderBlock2D, numBlocks)
+	for i := 0; i < numBlocks; i++ {
+		prefix := fmt.Sprintf("decoder.up_blocks.%d", i)
+		hasUpsample := i < numBlocks-1
+		m.UpBlocks[i], err = NewUpDecoderBlock2D(weights, prefix, cfg.LayersPerBlock+1, cfg.NormNumGroups, hasUpsample)
+		if err != nil {
+			return err
+		}
+	}
+	fmt.Printf("✓ [%d blocks]\n", numBlocks)
+
+	// Load conv_norm_out
+	fmt.Print("  Loading conv_norm_out... ")
+	normWeight, err := weights.GetTensor("decoder.conv_norm_out.weight")
+	if err != nil {
+		return err
+	}
+	normBias, err := weights.GetTensor("decoder.conv_norm_out.bias")
+	if err != nil {
+		return err
+	}
+	m.ConvNormOut = NewGroupNorm(normWeight, normBias, cfg.NormNumGroups)
+	fmt.Println("✓")
+
+	// Load conv_out
+	fmt.Print("  Loading conv_out... ")
+	convOutWeight, err := weights.GetTensor("decoder.conv_out.weight")
+	if err != nil {
+		return err
+	}
+	convOutBias, err := weights.GetTensor("decoder.conv_out.bias")
+	if err != nil {
+		return err
+	}
+	m.ConvOut = NewConv2D(convOutWeight, convOutBias, 1, 1)
+	fmt.Println("✓")
+
+	return nil
+}
+
+// Decode decodes latents to images.
+// Input latents are in NCHW format, output is in NCHW format.
+// If Tiling is set, uses tiled decoding to reduce memory for large images.
+func (v *VAEDecoder) Decode(latents *mlx.Array) *mlx.Array {
+	// Scale latents
+	z := mlx.DivScalar(latents, v.Config.ScalingFactor)
+	z = mlx.AddScalar(z, v.Config.ShiftFactor)
+	// Convert NCHW -> NHWC for internal processing
+	z = mlx.Transpose(z, 0, 2, 3, 1)
+
+	// Use tiled decoding if enabled
+	if v.Tiling != nil {
+		mlx.Eval(z)
+		return vae.DecodeTiled(z, v.Tiling, v.decodeTile)
+	}
+
+	// Direct decode
+	h := v.decodeTile(z)
+	h = mlx.ClipScalar(h, 0.0, 1.0, true, true)
+	// Convert NHWC -> NCHW for output
+	h = mlx.Transpose(h, 0, 3, 1, 2)
+	mlx.Eval(h)
+	return h
+}
+
+// decodeTile decodes a single latent tile to pixels.
+// Input: [B, H, W, C] latent tile in NHWC format (already scaled)
+// Output: [B, H*8, W*8, 3] pixel tile in NHWC format
+func (v *VAEDecoder) decodeTile(z *mlx.Array) *mlx.Array {
+	h := v.ConvIn.Forward(z)
+	mlx.Eval(h)
+
+	prev := h
+	h = v.MidBlock.Forward(h)
+	prev.Free()
+
+	for _, upBlock := range v.UpBlocks {
+		prev = h
+		h = upBlock.Forward(h)
+		prev.Free()
+	}
+
+	prev = h
+	h = v.ConvNormOut.Forward(h)
+	mlx.Eval(h) // Eval after GroupNorm to avoid grid dimension issues
+	prev.Free()
+
+	prev = h
+	h = mlx.SiLU(h)
+	h = v.ConvOut.Forward(h)
+	mlx.Eval(h)
+	prev.Free()
+
+	// VAE outputs [-1, 1], convert to [0, 1]
+	h = mlx.MulScalar(h, 0.5)
+	h = mlx.AddScalar(h, 0.5)
+
+	return h
+}
+
+// Upsample2x performs 2x nearest neighbor upsampling using Take.
+// Input and output are in NHWC format: [B, H, W, C] -> [B, H*2, W*2, C]
+// Uses Take with repeated indices to produce contiguous output.
+func Upsample2x(x *mlx.Array) *mlx.Array {
+	shape := x.Shape()
+	H := shape[1]
+	W := shape[2]
+
+	// Create indices [0, 0, 1, 1, 2, 2, ...] for nearest neighbor
+	// For H dimension
+	hIdx := mlx.ArangeInt(0, H, 1, mlx.DtypeInt32)
+	hIdx = mlx.Reshape(hIdx, H, 1)
+	hIdx = mlx.BroadcastTo(hIdx, []int32{H, 2})
+	hIdx = mlx.Reshape(hIdx, H*2)
+
+	// For W dimension
+	wIdx := mlx.ArangeInt(0, W, 1, mlx.DtypeInt32)
+	wIdx = mlx.Reshape(wIdx, W, 1)
+	wIdx = mlx.BroadcastTo(wIdx, []int32{W, 2})
+	wIdx = mlx.Reshape(wIdx, W*2)
+
+	// Take along H axis (axis 1 in NHWC)
+	x = mlx.Take(x, hIdx, 1)
+	// Take along W axis (axis 2 in NHWC)
+	x = mlx.Take(x, wIdx, 2)
+
+	return x
+}

x/imagegen/models/zimage/zimage.go

@@ -0,0 +1,488 @@
+// Package zimage implements the Z-Image diffusion transformer model.
+package zimage
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/ollama/ollama/x/imagegen/cache"
+	"github.com/ollama/ollama/x/imagegen/manifest"
+	"github.com/ollama/ollama/x/imagegen/mlx"
+	"github.com/ollama/ollama/x/imagegen/tokenizer"
+	"github.com/ollama/ollama/x/imagegen/vae"
+)
+
+// GenerateConfig holds all options for image generation.
+type GenerateConfig struct {
+	Prompt         string
+	NegativePrompt string                     // Empty = no CFG
+	CFGScale       float32                    // Only used if NegativePrompt is set (default: 4.0)
+	Width          int32                      // Image width (default: 1024)
+	Height         int32                      // Image height (default: 1024)
+	Steps          int                        // Denoising steps (default: 9 for turbo)
+	Seed           int64                      // Random seed
+	Progress       func(step, totalSteps int) // Optional progress callback
+	CapturePath    string                     // GPU capture path (debug)
+
+	// TeaCache options (timestep embedding aware caching)
+	TeaCache          bool    // TeaCache is always enabled for faster inference
+	TeaCacheThreshold float32 // Threshold for cache reuse (default: 0.1, lower = more aggressive)
+
+	// Fused QKV (fuse Q/K/V projections into single matmul)
+	FusedQKV bool // Enable fused QKV projection (default: false)
+}
+
+// Model represents a Z-Image diffusion model.
+type Model struct {
+	ModelName   string
+	Tokenizer   *tokenizer.Tokenizer
+	TextEncoder *Qwen3TextEncoder
+	Transformer *Transformer
+	VAEDecoder  *VAEDecoder
+	qkvFused    bool // Track if QKV has been fused (do only once)
+}
+
+// Load loads the Z-Image model from ollama blob storage.
+func (m *Model) Load(modelName string) error {
+	fmt.Printf("Loading Z-Image model from manifest: %s...\n", modelName)
+	start := time.Now()
+
+	if mlx.GPUIsAvailable() {
+		mlx.SetDefaultDeviceGPU()
+		mlx.EnableCompile()
+	}
+
+	m.ModelName = modelName
+
+	// Load manifest
+	manifest, err := manifest.LoadManifest(modelName)
+	if err != nil {
+		return fmt.Errorf("load manifest: %w", err)
+	}
+
+	// Load tokenizer from manifest with config
+	fmt.Print("  Loading tokenizer... ")
+	tokData, err := manifest.ReadConfig("tokenizer/tokenizer.json")
+	if err != nil {
+		return fmt.Errorf("tokenizer: %w", err)
+	}
+
+	// Try to read tokenizer config files from manifest
+	tokConfig := &tokenizer.TokenizerConfig{}
+	if data, err := manifest.ReadConfig("tokenizer/tokenizer_config.json"); err == nil {
+		tokConfig.TokenizerConfigJSON = data
+	}
+	if data, err := manifest.ReadConfig("tokenizer/generation_config.json"); err == nil {
+		tokConfig.GenerationConfigJSON = data
+	}
+	if data, err := manifest.ReadConfig("tokenizer/special_tokens_map.json"); err == nil {
+		tokConfig.SpecialTokensMapJSON = data
+	}
+
+	tok, err := tokenizer.LoadFromBytesWithConfig(tokData, tokConfig)
+	if err != nil {
+		return fmt.Errorf("tokenizer: %w", err)
+	}
+	m.Tokenizer = tok
+	fmt.Println("✓")
+
+	// Load text encoder
+	m.TextEncoder = &Qwen3TextEncoder{}
+	if err := m.TextEncoder.Load(manifest, "text_encoder/config.json"); err != nil {
+		return fmt.Errorf("text encoder: %w", err)
+	}
+	mlx.Eval(mlx.Collect(m.TextEncoder)...)
+	fmt.Printf("  (%.1f GB, peak %.1f GB)\n",
+		float64(mlx.MetalGetActiveMemory())/(1024*1024*1024),
+		float64(mlx.MetalGetPeakMemory())/(1024*1024*1024))
+
+	// Load transformer
+	m.Transformer = &Transformer{}
+	if err := m.Transformer.Load(manifest); err != nil {
+		return fmt.Errorf("transformer: %w", err)
+	}
+	mlx.Eval(mlx.Collect(m.Transformer)...)
+	fmt.Printf("  (%.1f GB, peak %.1f GB)\n",
+		float64(mlx.MetalGetActiveMemory())/(1024*1024*1024),
+		float64(mlx.MetalGetPeakMemory())/(1024*1024*1024))
+
+	// Load VAE decoder
+	m.VAEDecoder = &VAEDecoder{}
+	if err := m.VAEDecoder.Load(manifest); err != nil {
+		return fmt.Errorf("VAE decoder: %w", err)
+	}
+	mlx.Eval(mlx.Collect(m.VAEDecoder)...)
+	fmt.Printf("  (%.1f GB, peak %.1f GB)\n",
+		float64(mlx.MetalGetActiveMemory())/(1024*1024*1024),
+		float64(mlx.MetalGetPeakMemory())/(1024*1024*1024))
+
+	mem := mlx.MetalGetActiveMemory()
+	fmt.Printf("  Loaded in %.2fs (%.1f GB VRAM)\n", time.Since(start).Seconds(), float64(mem)/(1024*1024*1024))
+
+	return nil
+}
+
+// Generate creates an image from a prompt.
+func (m *Model) Generate(prompt string, width, height int32, steps int, seed int64) (*mlx.Array, error) {
+	return m.GenerateFromConfig(context.Background(), &GenerateConfig{
+		Prompt: prompt,
+		Width:  width,
+		Height: height,
+		Steps:  steps,
+		Seed:   seed,
+	})
+}
+
+// GenerateWithProgress creates an image with progress callback.
+func (m *Model) GenerateWithProgress(prompt string, width, height int32, steps int, seed int64, progress func(step, totalSteps int)) (*mlx.Array, error) {
+	return m.GenerateFromConfig(context.Background(), &GenerateConfig{
+		Prompt:   prompt,
+		Width:    width,
+		Height:   height,
+		Steps:    steps,
+		Seed:     seed,
+		Progress: progress,
+	})
+}
+
+// GenerateWithCFG creates an image with classifier-free guidance.
+func (m *Model) GenerateWithCFG(prompt, negativePrompt string, width, height int32, steps int, seed int64, cfgScale float32, progress func(step, totalSteps int)) (*mlx.Array, error) {
+	return m.GenerateFromConfig(context.Background(), &GenerateConfig{
+		Prompt:         prompt,
+		NegativePrompt: negativePrompt,
+		CFGScale:       cfgScale,
+		Width:          width,
+		Height:         height,
+		Steps:          steps,
+		Seed:           seed,
+		Progress:       progress,
+	})
+}
+
+// GenerateFromConfig generates an image using the unified config struct.
+func (m *Model) GenerateFromConfig(ctx context.Context, cfg *GenerateConfig) (*mlx.Array, error) {
+	start := time.Now()
+	result, err := m.generate(ctx, cfg)
+	if err != nil {
+		return nil, err
+	}
+	if cfg.NegativePrompt != "" {
+		fmt.Printf("Generated with CFG (scale=%.1f) in %.2fs (%d steps)\n", cfg.CFGScale, time.Since(start).Seconds(), cfg.Steps)
+	} else {
+		fmt.Printf("Generated in %.2fs (%d steps)\n", time.Since(start).Seconds(), cfg.Steps)
+	}
+	return result, nil
+}
+
+// GenerateImage implements runner.ImageModel interface.
+func (m *Model) GenerateImage(ctx context.Context, prompt string, width, height int32, steps int, seed int64, progress func(step, total int)) (*mlx.Array, error) {
+	return m.GenerateFromConfig(ctx, &GenerateConfig{
+		Prompt:   prompt,
+		Width:    width,
+		Height:   height,
+		Steps:    steps,
+		Seed:     seed,
+		Progress: progress,
+	})
+}
+
+// generate is the internal denoising pipeline.
+func (m *Model) generate(ctx context.Context, cfg *GenerateConfig) (*mlx.Array, error) {
+	// Apply defaults
+	if cfg.Width <= 0 {
+		cfg.Width = 1024
+	}
+	if cfg.Height <= 0 {
+		cfg.Height = 1024
+	}
+	if cfg.Steps <= 0 {
+		cfg.Steps = 9 // Z-Image turbo default
+	}
+	if cfg.CFGScale <= 0 {
+		cfg.CFGScale = 4.0
+	}
+	// TeaCache enabled by default
+	cfg.TeaCache = true
+	if cfg.TeaCacheThreshold <= 0 {
+		cfg.TeaCacheThreshold = 0.15
+	}
+
+	// Enable fused QKV if requested (only fuse once)
+	if cfg.FusedQKV && !m.qkvFused {
+		m.Transformer.FuseAllQKV()
+		m.qkvFused = true
+		fmt.Println("  Fused QKV enabled")
+	}
+
+	useCFG := cfg.NegativePrompt != ""
+	tcfg := m.Transformer.TransformerConfig
+	latentH := cfg.Height / 8
+	latentW := cfg.Width / 8
+	hTok := latentH / tcfg.PatchSize
+	wTok := latentW / tcfg.PatchSize
+
+	// Text encoding with padding to multiple of 32
+	var posEmb, negEmb *mlx.Array
+	{
+		posEmb, _ = m.TextEncoder.EncodePrompt(m.Tokenizer, cfg.Prompt, 512, false)
+		if useCFG {
+			negEmb, _ = m.TextEncoder.EncodePrompt(m.Tokenizer, cfg.NegativePrompt, 512, false)
+		}
+
+		// Pad both to same length (multiple of 32)
+		maxLen := posEmb.Shape()[1]
+		if useCFG && negEmb.Shape()[1] > maxLen {
+			maxLen = negEmb.Shape()[1]
+		}
+		if pad := (32 - (maxLen % 32)) % 32; pad > 0 {
+			maxLen += pad
+		}
+
+		posEmb = padToLength(posEmb, maxLen)
+		if useCFG {
+			negEmb = padToLength(negEmb, maxLen)
+			mlx.Keep(posEmb, negEmb)
+			mlx.Eval(posEmb, negEmb)
+		} else {
+			mlx.Keep(posEmb)
+			mlx.Eval(posEmb)
+		}
+	}
+
+	// Scheduler
+	scheduler := NewFlowMatchEulerScheduler(DefaultFlowMatchSchedulerConfig())
+	scheduler.SetTimestepsWithMu(cfg.Steps, CalculateShift(hTok*wTok))
+
+	// Init latents [B, C, H, W]
+	var latents *mlx.Array
+	{
+		latents = scheduler.InitNoise([]int32{1, tcfg.InChannels, latentH, latentW}, cfg.Seed)
+		mlx.Eval(latents)
+	}
+
+	// RoPE cache
+	var ropeCache *RoPECache
+	{
+		ropeCache = m.Transformer.PrepareRoPECache(hTok, wTok, posEmb.Shape()[1])
+		mlx.Keep(ropeCache.ImgCos, ropeCache.ImgSin, ropeCache.CapCos, ropeCache.CapSin,
+			ropeCache.UnifiedCos, ropeCache.UnifiedSin)
+		mlx.Eval(ropeCache.UnifiedCos)
+	}
+
+	// Pre-compute batched embeddings for CFG (outside the loop for efficiency)
+	var batchedEmb *mlx.Array
+	if useCFG {
+		// Concatenate embeddings once: [1, L, D] + [1, L, D] -> [2, L, D]
+		batchedEmb = mlx.Concatenate([]*mlx.Array{posEmb, negEmb}, 0)
+		mlx.Keep(batchedEmb)
+		mlx.Eval(batchedEmb)
+	}
+
+	// TeaCache for timestep-aware caching
+	// For CFG mode, we cache pos/neg separately, skip early steps, and always compute CFG fresh
+	var teaCache *cache.TeaCache
+	if cfg.TeaCache {
+		skipEarly := 0
+		if useCFG {
+			skipEarly = 3 // Skip first 3 steps for CFG to preserve structure
+		}
+		teaCache = cache.NewTeaCache(&cache.TeaCacheConfig{
+			Threshold:      cfg.TeaCacheThreshold,
+			RescaleFactor:  1.0,
+			SkipEarlySteps: skipEarly,
+		})
+		if useCFG {
+			fmt.Printf("  TeaCache enabled (CFG mode): threshold=%.2f, skip first %d steps\n", cfg.TeaCacheThreshold, skipEarly)
+		} else {
+			fmt.Printf("  TeaCache enabled: threshold=%.2f\n", cfg.TeaCacheThreshold)
+		}
+	}
+
+	// cleanup frees all kept arrays when we need to abort early
+	cleanup := func() {
+		posEmb.Free()
+		if negEmb != nil {
+			negEmb.Free()
+		}
+		ropeCache.ImgCos.Free()
+		ropeCache.ImgSin.Free()
+		ropeCache.CapCos.Free()
+		ropeCache.CapSin.Free()
+		ropeCache.UnifiedCos.Free()
+		ropeCache.UnifiedSin.Free()
+		if batchedEmb != nil {
+			batchedEmb.Free()
+		}
+		if teaCache != nil {
+			teaCache.Free()
+		}
+		latents.Free()
+	}
+
+	// Denoising loop
+	if cfg.Progress != nil {
+		cfg.Progress(0, cfg.Steps) // Start at 0%
+	}
+	for i := 0; i < cfg.Steps; i++ {
+		// Check for cancellation
+		if ctx != nil {
+			select {
+			case <-ctx.Done():
+				cleanup()
+				return nil, ctx.Err()
+			default:
+			}
+		}
+		stepStart := time.Now()
+
+		// GPU capture on step 2 if requested
+		if cfg.CapturePath != "" && i == 1 {
+			mlx.MetalStartCapture(cfg.CapturePath)
+		}
+
+		tCurr := scheduler.Timesteps[i]
+		var noisePred *mlx.Array
+
+		// TeaCache: check if we should compute or reuse cached output
+		shouldCompute := teaCache == nil || teaCache.ShouldCompute(i, tCurr)
+
+		if shouldCompute {
+			timestep := mlx.ToBFloat16(mlx.NewArray([]float32{1.0 - tCurr}, []int32{1}))
+			patches := PatchifyLatents(latents, tcfg.PatchSize)
+
+			var output *mlx.Array
+			if useCFG {
+				// CFG Batching: single forward pass with batch=2
+				// Tile patches: [1, L, D] -> [2, L, D]
+				batchedPatches := mlx.Tile(patches, []int32{2, 1, 1})
+				// Tile timestep: [1] -> [2]
+				batchedTimestep := mlx.Tile(timestep, []int32{2})
+
+				// Single batched forward pass (RoPE broadcasts from [1,L,H,D] to [2,L,H,D])
+				batchedOutput := m.Transformer.Forward(batchedPatches, batchedTimestep, batchedEmb, ropeCache)
+
+				// Split output: [2, L, D] -> pos [1, L, D], neg [1, L, D]
+				outputShape := batchedOutput.Shape()
+				L := outputShape[1]
+				D := outputShape[2]
+				posOutput := mlx.Slice(batchedOutput, []int32{0, 0, 0}, []int32{1, L, D})
+				negOutput := mlx.Slice(batchedOutput, []int32{1, 0, 0}, []int32{2, L, D})
+
+				// Convert to noise predictions (unpatchify and negate)
+				posPred := UnpatchifyLatents(posOutput, tcfg.PatchSize, latentH, latentW, tcfg.InChannels)
+				posPred = mlx.Neg(posPred)
+				negPred := UnpatchifyLatents(negOutput, tcfg.PatchSize, latentH, latentW, tcfg.InChannels)
+				negPred = mlx.Neg(negPred)
+
+				// Cache pos/neg separately for TeaCache
+				if teaCache != nil {
+					teaCache.UpdateCFGCache(posPred, negPred, tCurr)
+					mlx.Keep(teaCache.Arrays()...)
+				}
+
+				// Apply CFG: noisePred = neg + scale * (pos - neg)
+				diff := mlx.Sub(posPred, negPred)
+				scaledDiff := mlx.MulScalar(diff, cfg.CFGScale)
+				noisePred = mlx.Add(negPred, scaledDiff)
+			} else {
+				// Non-CFG forward pass
+				output = m.Transformer.Forward(patches, timestep, posEmb, ropeCache)
+				noisePred = UnpatchifyLatents(output, tcfg.PatchSize, latentH, latentW, tcfg.InChannels)
+				noisePred = mlx.Neg(noisePred)
+
+				// Update TeaCache
+				if teaCache != nil {
+					teaCache.UpdateCache(noisePred, tCurr)
+					mlx.Keep(teaCache.Arrays()...)
+				}
+			}
+		} else if useCFG && teaCache != nil && teaCache.HasCFGCache() {
+			// CFG mode: get cached pos/neg and compute CFG fresh
+			posPred, negPred := teaCache.GetCFGCached()
+			diff := mlx.Sub(posPred, negPred)
+			scaledDiff := mlx.MulScalar(diff, cfg.CFGScale)
+			noisePred = mlx.Add(negPred, scaledDiff)
+			fmt.Printf("    [TeaCache: reusing cached pos/neg outputs]\n")
+		} else {
+			// Non-CFG mode: reuse cached noise prediction
+			noisePred = teaCache.GetCached()
+			fmt.Printf("    [TeaCache: reusing cached output]\n")
+		}
+
+		oldLatents := latents
+		latents = scheduler.Step(noisePred, latents, i)
+
+		mlx.Eval(latents)
+		oldLatents.Free()
+
+		if cfg.CapturePath != "" && i == 1 {
+			mlx.MetalStopCapture()
+		}
+
+		activeMem := float64(mlx.MetalGetActiveMemory()) / (1024 * 1024 * 1024)
+		peakMem := float64(mlx.MetalGetPeakMemory()) / (1024 * 1024 * 1024)
+		fmt.Printf("  Step %d/%d: t=%.4f (%.2fs) [%.1f GB active, %.1f GB peak]\n",
+			i+1, cfg.Steps, tCurr, time.Since(stepStart).Seconds(), activeMem, peakMem)
+
+		if cfg.Progress != nil {
+			cfg.Progress(i+1, cfg.Steps) // Report completed step
+		}
+	}
+
+	// Free denoising temporaries before VAE decode
+	posEmb.Free()
+	if negEmb != nil {
+		negEmb.Free()
+	}
+	ropeCache.ImgCos.Free()
+	ropeCache.ImgSin.Free()
+	ropeCache.CapCos.Free()
+	ropeCache.CapSin.Free()
+	ropeCache.UnifiedCos.Free()
+	ropeCache.UnifiedSin.Free()
+	if batchedEmb != nil {
+		batchedEmb.Free()
+	}
+	if teaCache != nil {
+		hits, misses := teaCache.Stats()
+		fmt.Printf("  TeaCache stats: %d hits, %d misses (%.1f%% cache rate)\n",
+			hits, misses, float64(hits)/float64(hits+misses)*100)
+		teaCache.Free()
+	}
+
+	// VAE decode - enable tiling for larger images to reduce memory
+	// VAE attention is O(n²) on latent pixels, tiling helps significantly
+	if latentH > 64 || latentW > 64 {
+		m.VAEDecoder.Tiling = vae.DefaultTilingConfig()
+	}
+	decoded := m.VAEDecoder.Decode(latents)
+	latents.Free()
+
+	return decoded, nil
+}
+
+// padToLength pads a sequence tensor to the target length by repeating the last token.
+func padToLength(x *mlx.Array, targetLen int32) *mlx.Array {
+	shape := x.Shape()
+	currentLen := shape[1]
+	if currentLen >= targetLen {
+		return x
+	}
+	padLen := targetLen - currentLen
+	lastToken := mlx.Slice(x, []int32{0, currentLen - 1, 0}, []int32{shape[0], currentLen, shape[2]})
+	padding := mlx.Tile(lastToken, []int32{1, padLen, 1})
+	return mlx.Concatenate([]*mlx.Array{x, padding}, 1)
+}
+
+// CalculateShift computes the mu shift value for dynamic scheduling
+func CalculateShift(imgSeqLen int32) float32 {
+	baseSeqLen := float32(256)
+	maxSeqLen := float32(4096)
+	baseShift := float32(0.5)
+	maxShift := float32(1.15)
+
+	m := (maxShift - baseShift) / (maxSeqLen - baseSeqLen)
+	b := baseShift - m*baseSeqLen
+	return float32(imgSeqLen)*m + b
+}

x/imagegen/nn/nn.go

@@ -0,0 +1,258 @@
+// Package nn provides neural network layer types.
+package nn
+
+import "github.com/ollama/ollama/x/imagegen/mlx"
+
+// Layer is the interface for neural network layers with a Forward method.
+type Layer interface {
+	Forward(x *mlx.Array) *mlx.Array
+}
+
+// LinearLayer is an interface for linear layers (both regular and quantized).
+// This allows swapping between Linear and QuantizedLinear at runtime.
+type LinearLayer interface {
+	Forward(x *mlx.Array) *mlx.Array
+	OutputDim() int32 // Returns the output dimension of the layer
+}
+
+// Linear applies an affine transformation: y = x @ W.T + b
+// Weight is stored as [out_features, in_features], matching PyTorch/MLX convention.
+type Linear struct {
+	Weight *mlx.Array `weight:"weight"`          // [out_features, in_features]
+	Bias   *mlx.Array `weight:"bias,optional"`   // [out_features] or nil
+}
+
+// NewLinear creates a linear layer.
+// Weight should be [out_features, in_features].
+func NewLinear(weight *mlx.Array, bias *mlx.Array) *Linear {
+	return &Linear{Weight: weight, Bias: bias}
+}
+
+// NewQuantizedLinear creates a quantized linear layer directly from bf16 weights.
+// Quantizes the weight immediately and evaluates to break lazy dependencies.
+// Note: For modes like "nvfp4", qbiases will be nil.
+func NewQuantizedLinear(weight *mlx.Array, bias *mlx.Array, groupSize, bits int, mode string) *QuantizedLinear {
+	qw, scales, qbiases := mlx.Quantize(weight, groupSize, bits, mode)
+	// Eval immediately so bf16 weight can be freed
+	// Handle modes that don't return qbiases (e.g., nvfp4)
+	if qbiases != nil {
+		mlx.Eval(qw, scales, qbiases)
+	} else {
+		mlx.Eval(qw, scales)
+	}
+	return &QuantizedLinear{
+		Weight:    qw,
+		Scales:    scales,
+		QBiases:   qbiases,
+		Bias:      bias,
+		GroupSize: groupSize,
+		Bits:      bits,
+		Mode:      mode,
+	}
+}
+
+// Forward applies the linear transformation: x @ W.T + bias
+func (l *Linear) Forward(x *mlx.Array) *mlx.Array {
+	w := mlx.Transpose(l.Weight, 1, 0)
+	if l.Bias != nil {
+		return mlx.AddMM(l.Bias, x, w, 1.0, 1.0)
+	}
+	return mlx.Linear(x, w)
+}
+
+// OutputDim returns the output dimension of the linear layer.
+func (l *Linear) OutputDim() int32 {
+	return l.Weight.Shape()[0]
+}
+
+// ToQuantized converts this Linear to a QuantizedLinear.
+func (l *Linear) ToQuantized(groupSize, bits int, mode string) *QuantizedLinear {
+	qw, scales, qbiases := mlx.Quantize(l.Weight, groupSize, bits, mode)
+	return &QuantizedLinear{
+		Weight:    qw,
+		Scales:    scales,
+		QBiases:   qbiases,
+		Bias:      l.Bias,
+		GroupSize: groupSize,
+		Bits:      bits,
+		Mode:      mode,
+	}
+}
+
+// QuantizedLinear applies an affine transformation using quantized weights.
+// Equivalent to mlx.nn.QuantizedLinear.
+// Supports multiple quantization modes:
+//   - "affine": scale + zero-point bias (QBiases required)
+//   - "nvfp4": NVIDIA FP4 with E4M3 scales (QBiases nil)
+type QuantizedLinear struct {
+	Weight    *mlx.Array // Quantized weight data
+	Scales    *mlx.Array // Scale factors for dequantization
+	QBiases   *mlx.Array // Quantization biases (NOT layer bias), nil for nvfp4
+	Bias      *mlx.Array // Layer bias [output_dims] or nil
+	GroupSize int
+	Bits      int
+	Mode      string
+}
+
+// Forward applies the quantized linear transformation.
+func (ql *QuantizedLinear) Forward(x *mlx.Array) *mlx.Array {
+	out := mlx.QuantizedMatmul(x, ql.Weight, ql.Scales, ql.QBiases, true, ql.GroupSize, ql.Bits, ql.Mode)
+	if ql.Bias != nil {
+		out = mlx.Add(out, ql.Bias)
+	}
+	return out
+}
+
+// OutputDim returns the output dimension of the quantized linear layer.
+// For mxfp8/mxfp4, quantized weight shape is [out_features, in_features / group_size].
+// The output dimension is the first dimension of the weight.
+func (ql *QuantizedLinear) OutputDim() int32 {
+	return ql.Weight.Shape()[0]
+}
+
+// RMSNorm represents an RMS normalization layer.
+type RMSNorm struct {
+	Weight *mlx.Array `weight:"weight"`
+	Eps    float32    // optional: used if Forward called with eps=0
+}
+
+// NewRMSNorm creates an RMSNorm layer (for models not using weight loader).
+func NewRMSNorm(weight *mlx.Array, eps float32) *RMSNorm {
+	return &RMSNorm{Weight: weight, Eps: eps}
+}
+
+// Forward applies RMS normalization. If eps=0, uses stored Eps.
+func (rn *RMSNorm) Forward(x *mlx.Array, eps float32) *mlx.Array {
+	if eps == 0 {
+		eps = rn.Eps
+	}
+	return mlx.RMSNorm(x, rn.Weight, eps)
+}
+
+// Embedding represents an embedding layer.
+type Embedding struct {
+	Weight *mlx.Array `weight:"weight"`
+}
+
+// NewEmbedding creates an embedding layer.
+func NewEmbedding(weight *mlx.Array) *Embedding {
+	return &Embedding{Weight: weight}
+}
+
+// Forward looks up embeddings by indices.
+func (e *Embedding) Forward(indices *mlx.Array) *mlx.Array {
+	return mlx.Take(e.Weight, indices, 0)
+}
+
+// RepeatKV repeats K/V tensors for grouped query attention
+// x: [B, num_kv_heads, S, head_dim] -> [B, num_heads, S, head_dim]
+func RepeatKV(x *mlx.Array, repeatFactor int32) *mlx.Array {
+	if repeatFactor == 1 {
+		return x
+	}
+	shape := x.Shape()
+	// [B, num_kv_heads, S, head_dim] -> [B, num_kv_heads, 1, S, head_dim]
+	x = mlx.ExpandDims(x, 2)
+	// Repeat along the new axis
+	reps := []int32{1, 1, repeatFactor, 1, 1}
+	x = mlx.Tile(x, reps)
+	// Reshape: [B, num_kv_heads, repeat, S, head_dim] -> [B, num_kv_heads * repeat, S, head_dim]
+	return mlx.Reshape(x, shape[0], shape[1]*repeatFactor, shape[2], shape[3])
+}
+
+// ApplyCausalMask applies causal (lower triangular) mask to attention scores
+func ApplyCausalMask(scores *mlx.Array) *mlx.Array {
+	// scores: [B, num_heads, S, S]
+	shape := scores.Shape()
+	seqLen := shape[2]
+
+	// Create causal mask: 1 for positions to keep, 0 for positions to mask
+	mask := mlx.Tri(seqLen, seqLen, 0)
+
+	// Where mask is 0, set score to -inf
+	negInf := mlx.NewScalarArray(float32(-1e9))
+
+	// Broadcast mask to match scores shape
+	mask = mlx.ExpandDims(mlx.ExpandDims(mask, 0), 0) // [1, 1, S, S]
+
+	// Use where: if mask > 0, keep scores, else -inf
+	return mlx.Where(mask, scores, negInf)
+}
+
+// ApplyCausalMaskWithOffset applies causal mask for cached attention
+// scores: [B, num_heads, queryLen, keyLen] where keyLen = cacheLen + queryLen
+// offset: the starting position of the new queries (i.e., cache length)
+func ApplyCausalMaskWithOffset(scores *mlx.Array, offset int32) *mlx.Array {
+	if offset == 0 {
+		return ApplyCausalMask(scores)
+	}
+
+	shape := scores.Shape()
+	queryLen := shape[2]
+	keyLen := shape[3]
+
+	// For cached attention, new queries can attend to all cached keys plus
+	// new keys up to and including their position.
+	mask := mlx.Tri(queryLen, keyLen, int(offset))
+
+	negInf := mlx.NewScalarArray(float32(-1e9))
+	mask = mlx.ExpandDims(mlx.ExpandDims(mask, 0), 0) // [1, 1, queryLen, keyLen]
+
+	return mlx.Where(mask, scores, negInf)
+}
+
+// LayerNorm represents a standard layer normalization layer (with bias).
+type LayerNorm struct {
+	Weight *mlx.Array `weight:"weight"`
+	Bias   *mlx.Array `weight:"bias"`
+	Eps    float32
+}
+
+// Forward applies layer normalization: (x - mean) / sqrt(var + eps) * weight + bias
+func (ln *LayerNorm) Forward(x *mlx.Array) *mlx.Array {
+	eps := ln.Eps
+	if eps == 0 {
+		eps = 1e-5
+	}
+	// Compute mean and variance along last dimension
+	mean := mlx.Mean(x, -1, true)
+	centered := mlx.Sub(x, mean)
+	variance := mlx.Mean(mlx.Mul(centered, centered), -1, true)
+	normalized := mlx.Mul(centered, mlx.RSqrt(mlx.AddScalar(variance, eps)))
+
+	// Scale and shift
+	out := mlx.Mul(normalized, ln.Weight)
+	if ln.Bias != nil {
+		out = mlx.Add(out, ln.Bias)
+	}
+	return out
+}
+
+// MultiLinearLayer is an interface for per-head linear layers.
+// This allows swapping between MultiLinear (bf16) and pre-dequantized weights.
+type MultiLinearLayer interface {
+	Forward(x *mlx.Array) *mlx.Array
+}
+
+// MultiLinear performs per-head linear projections.
+// Weight shape: [num_heads, output_dims, input_dims]
+// Input shape: [B, num_heads, L, input_dims]
+// Output shape: [B, num_heads, L, output_dims]
+type MultiLinear struct {
+	Weight *mlx.Array `weight:"weight"`
+}
+
+// NewMultiLinear creates a MultiLinear layer with the given weight.
+func NewMultiLinear(weight *mlx.Array) *MultiLinear {
+	return &MultiLinear{Weight: weight}
+}
+
+// Forward applies per-head linear transformation: x @ weight.T per head via broadcasting.
+func (ml *MultiLinear) Forward(x *mlx.Array) *mlx.Array {
+	// Weight: [num_heads, output_dims, input_dims]
+	// x: [B, num_heads, L, input_dims]
+	// wT: [num_heads, input_dims, output_dims]
+	// Result: [B, num_heads, L, output_dims]
+	wT := mlx.Transpose(ml.Weight, 0, 2, 1)
+	return mlx.Matmul(x, wT)
+}

x/imagegen/nn/nn_test.go

@@ -0,0 +1,376 @@
+package nn
+
+import (
+	"fmt"
+	"math"
+	"os"
+	"path/filepath"
+	"runtime"
+	"testing"
+
+	"github.com/ollama/ollama/x/imagegen/mlx"
+)
+
+// TestMain initializes MLX before running tests.
+// If MLX libraries are not available, tests are skipped.
+func TestMain(m *testing.M) {
+	// Change to repo root so ./build/lib/ollama/ path works
+	_, thisFile, _, _ := runtime.Caller(0)
+	repoRoot := filepath.Join(filepath.Dir(thisFile), "..", "..", "..")
+	if err := os.Chdir(repoRoot); err != nil {
+		fmt.Printf("Failed to change to repo root: %v\n", err)
+		os.Exit(1)
+	}
+
+	if err := mlx.InitMLX(); err != nil {
+		fmt.Printf("Skipping nn tests: %v\n", err)
+		os.Exit(0)
+	}
+	os.Exit(m.Run())
+}
+
+// TestLinearNoBias verifies Linear without bias computes x @ w.T correctly.
+func TestLinearNoBias(t *testing.T) {
+	// Weight: [out=2, in=3] -> transposed at forward time
+	weight := mlx.NewArrayFloat32([]float32{
+		1, 2, 3, // row 0
+		4, 5, 6, // row 1
+	}, []int32{2, 3})
+	mlx.Eval(weight)
+
+	linear := NewLinear(weight, nil)
+
+	// Input: [1, 3]
+	x := mlx.NewArrayFloat32([]float32{1, 1, 1}, []int32{1, 3})
+	mlx.Eval(x)
+
+	out := linear.Forward(x)
+	mlx.Eval(out)
+
+	// Expected: [1,1,1] @ [[1,4],[2,5],[3,6]] = [6, 15]
+	data := out.Data()
+	if len(data) != 2 || data[0] != 6 || data[1] != 15 {
+		t.Errorf("expected [6, 15], got %v", data)
+	}
+}
+
+// TestLinearWithBias verifies Linear with bias computes x @ w.T + b correctly.
+func TestLinearWithBias(t *testing.T) {
+	weight := mlx.NewArrayFloat32([]float32{
+		1, 2, 3,
+		4, 5, 6,
+	}, []int32{2, 3})
+	bias := mlx.NewArrayFloat32([]float32{10, 20}, []int32{2})
+	mlx.Eval(weight, bias)
+
+	linear := NewLinear(weight, bias)
+
+	x := mlx.NewArrayFloat32([]float32{1, 1, 1}, []int32{1, 3})
+	mlx.Eval(x)
+
+	out := linear.Forward(x)
+	mlx.Eval(out)
+
+	// Expected: [6, 15] + [10, 20] = [16, 35]
+	data := out.Data()
+	if len(data) != 2 || data[0] != 16 || data[1] != 35 {
+		t.Errorf("expected [16, 35], got %v", data)
+	}
+}
+
+// TestLinearBatched verifies Linear works with batched input.
+func TestLinearBatched(t *testing.T) {
+	weight := mlx.NewArrayFloat32([]float32{
+		1, 0,
+		0, 1,
+	}, []int32{2, 2}) // Identity
+	mlx.Eval(weight)
+
+	linear := NewLinear(weight, nil)
+
+	// Batch of 3 inputs
+	x := mlx.NewArrayFloat32([]float32{
+		1, 2,
+		3, 4,
+		5, 6,
+	}, []int32{3, 2})
+	mlx.Eval(x)
+
+	out := linear.Forward(x)
+	mlx.Eval(out)
+
+	// Identity should return same values
+	data := out.Data()
+	expected := []float32{1, 2, 3, 4, 5, 6}
+	for i, v := range expected {
+		if data[i] != v {
+			t.Errorf("at %d: expected %f, got %f", i, v, data[i])
+		}
+	}
+}
+
+// TestRMSNorm verifies RMSNorm computation.
+func TestRMSNorm(t *testing.T) {
+	weight := mlx.NewArrayFloat32([]float32{1, 1, 1, 1}, []int32{4})
+	mlx.Eval(weight)
+
+	norm := NewRMSNorm(weight, 1e-5)
+
+	// Input with known RMS
+	x := mlx.NewArrayFloat32([]float32{2, 2, 2, 2}, []int32{1, 4})
+	mlx.Eval(x)
+
+	out := norm.Forward(x, 0) // eps=0 uses stored Eps
+	mlx.Eval(out)
+
+	// RMS of [2,2,2,2] = 2, so normalized = [1,1,1,1]
+	data := out.Data()
+	for i, v := range data {
+		if math.Abs(float64(v-1.0)) > 1e-4 {
+			t.Errorf("at %d: expected ~1.0, got %f", i, v)
+		}
+	}
+}
+
+// TestRMSNormWithScale verifies RMSNorm applies weight scaling.
+func TestRMSNormWithScale(t *testing.T) {
+	weight := mlx.NewArrayFloat32([]float32{2, 2, 2, 2}, []int32{4})
+	mlx.Eval(weight)
+
+	norm := NewRMSNorm(weight, 1e-5)
+
+	x := mlx.NewArrayFloat32([]float32{2, 2, 2, 2}, []int32{1, 4})
+	mlx.Eval(x)
+
+	out := norm.Forward(x, 0) // eps=0 uses stored Eps
+	mlx.Eval(out)
+
+	// Normalized [1,1,1,1] * weight [2,2,2,2] = [2,2,2,2]
+	data := out.Data()
+	for i, v := range data {
+		if math.Abs(float64(v-2.0)) > 1e-4 {
+			t.Errorf("at %d: expected ~2.0, got %f", i, v)
+		}
+	}
+}
+
+// TestEmbedding verifies embedding lookup.
+func TestEmbedding(t *testing.T) {
+	// Embedding table: 4 tokens, dim 3
+	weight := mlx.NewArrayFloat32([]float32{
+		0, 0, 0, // token 0
+		1, 1, 1, // token 1
+		2, 2, 2, // token 2
+		3, 3, 3, // token 3
+	}, []int32{4, 3})
+	mlx.Eval(weight)
+
+	emb := NewEmbedding(weight)
+
+	// Look up tokens [1, 3, 0]
+	indices := mlx.NewArrayInt32([]int32{1, 3, 0}, []int32{3})
+	mlx.Eval(indices)
+
+	out := emb.Forward(indices)
+	mlx.Eval(out)
+
+	data := out.Data()
+	expected := []float32{1, 1, 1, 3, 3, 3, 0, 0, 0}
+	for i, v := range expected {
+		if data[i] != v {
+			t.Errorf("at %d: expected %f, got %f", i, v, data[i])
+		}
+	}
+}
+
+// TestRepeatKV verifies K/V repetition for GQA.
+func TestRepeatKV(t *testing.T) {
+	// [B=1, num_kv_heads=2, S=2, head_dim=2]
+	x := mlx.NewArrayFloat32([]float32{
+		// head 0
+		1, 2, // pos 0
+		3, 4, // pos 1
+		// head 1
+		5, 6, // pos 0
+		7, 8, // pos 1
+	}, []int32{1, 2, 2, 2})
+	mlx.Eval(x)
+
+	// Repeat factor 2: 2 kv heads -> 4 heads
+	out := RepeatKV(x, 2)
+	mlx.Eval(out)
+
+	shape := out.Shape()
+	if shape[0] != 1 || shape[1] != 4 || shape[2] != 2 || shape[3] != 2 {
+		t.Errorf("expected shape [1,4,2,2], got %v", shape)
+	}
+
+	data := out.Data()
+	// After repeat: head0, head0, head1, head1
+	expected := []float32{
+		1, 2, 3, 4, // head 0 (original)
+		1, 2, 3, 4, // head 0 (repeat)
+		5, 6, 7, 8, // head 1 (original)
+		5, 6, 7, 8, // head 1 (repeat)
+	}
+	for i, v := range expected {
+		if data[i] != v {
+			t.Errorf("at %d: expected %f, got %f", i, v, data[i])
+		}
+	}
+}
+
+// TestRepeatKVNoOp verifies RepeatKV with factor 1 returns input unchanged.
+func TestRepeatKVNoOp(t *testing.T) {
+	x := mlx.NewArrayFloat32([]float32{1, 2, 3, 4}, []int32{1, 1, 2, 2})
+	mlx.Eval(x)
+
+	out := RepeatKV(x, 1)
+	// Should return same pointer
+	if out != x {
+		t.Error("RepeatKV with factor 1 should return input unchanged")
+	}
+}
+
+// TestApplyCausalMask verifies causal masking.
+func TestApplyCausalMask(t *testing.T) {
+	// [B=1, heads=1, S=3, S=3] - all ones
+	scores := mlx.Ones(1, 1, 3, 3)
+	mlx.Eval(scores)
+
+	out := ApplyCausalMask(scores)
+	mlx.Eval(out)
+
+	data := out.Data()
+	// Lower triangular should be 1, upper should be -1e9
+	// Row 0: [1, -inf, -inf]
+	// Row 1: [1, 1, -inf]
+	// Row 2: [1, 1, 1]
+	if data[0] != 1 || data[1] >= 0 || data[2] >= 0 {
+		t.Errorf("row 0 wrong: %v", data[0:3])
+	}
+	if data[3] != 1 || data[4] != 1 || data[5] >= 0 {
+		t.Errorf("row 1 wrong: %v", data[3:6])
+	}
+	if data[6] != 1 || data[7] != 1 || data[8] != 1 {
+		t.Errorf("row 2 wrong: %v", data[6:9])
+	}
+}
+
+// TestApplyCausalMaskWithOffset verifies causal masking with cache offset.
+func TestApplyCausalMaskWithOffset(t *testing.T) {
+	// Simulating: cache has 2 tokens, adding 1 new query
+	// scores: [B=1, heads=1, queryLen=1, keyLen=3]
+	scores := mlx.Ones(1, 1, 1, 3)
+	mlx.Eval(scores)
+
+	out := ApplyCausalMaskWithOffset(scores, 2)
+	mlx.Eval(out)
+
+	data := out.Data()
+	// With offset=2, query at position 2 can attend to all 3 positions
+	if data[0] != 1 || data[1] != 1 || data[2] != 1 {
+		t.Errorf("expected [1, 1, 1], got %v", data)
+	}
+}
+
+// TestApplyCausalMaskWithOffsetZero verifies offset=0 falls back to regular causal.
+func TestApplyCausalMaskWithOffsetZero(t *testing.T) {
+	scores := mlx.Ones(1, 1, 2, 2)
+	mlx.Eval(scores)
+
+	out := ApplyCausalMaskWithOffset(scores, 0)
+	mlx.Eval(out)
+
+	data := out.Data()
+	// Standard causal: [1, -inf], [1, 1]
+	if data[0] != 1 || data[1] >= 0 {
+		t.Errorf("row 0 wrong: %v", data[0:2])
+	}
+	if data[2] != 1 || data[3] != 1 {
+		t.Errorf("row 1 wrong: %v", data[2:4])
+	}
+}
+
+// BenchmarkLinearSmall benchmarks small Linear forward pass.
+func BenchmarkLinearSmall(b *testing.B) {
+	weight := mlx.RandomNormal([]int32{256, 256}, 42)
+	mlx.Eval(weight)
+
+	linear := NewLinear(weight, nil)
+
+	x := mlx.RandomNormal([]int32{1, 256}, 43)
+	mlx.Eval(x)
+
+	b.ResetTimer()
+	for range b.N {
+		out := linear.Forward(x)
+		mlx.Eval(out)
+	}
+}
+
+// BenchmarkLinearLarge benchmarks larger Linear forward pass.
+func BenchmarkLinearLarge(b *testing.B) {
+	weight := mlx.RandomNormal([]int32{4096, 4096}, 42)
+	mlx.Eval(weight)
+
+	linear := NewLinear(weight, nil)
+
+	x := mlx.RandomNormal([]int32{1, 4096}, 43)
+	mlx.Eval(x)
+
+	b.ResetTimer()
+	for range b.N {
+		out := linear.Forward(x)
+		mlx.Eval(out)
+	}
+}
+
+// BenchmarkRMSNorm benchmarks RMSNorm forward pass.
+func BenchmarkRMSNorm(b *testing.B) {
+	weight := mlx.Ones(4096)
+	mlx.Eval(weight)
+
+	norm := NewRMSNorm(weight, 1e-5)
+
+	x := mlx.RandomNormal([]int32{1, 4096}, 42)
+	mlx.Eval(x)
+
+	b.ResetTimer()
+	for range b.N {
+		out := norm.Forward(x, 0)
+		mlx.Eval(out)
+	}
+}
+
+// BenchmarkEmbedding benchmarks embedding lookup.
+func BenchmarkEmbedding(b *testing.B) {
+	// Typical vocab size
+	weight := mlx.RandomNormal([]int32{32000, 4096}, 42)
+	mlx.Eval(weight)
+
+	emb := NewEmbedding(weight)
+
+	// Single token lookup
+	indices := mlx.NewArrayInt32([]int32{1000}, []int32{1})
+	mlx.Eval(indices)
+
+	b.ResetTimer()
+	for range b.N {
+		out := emb.Forward(indices)
+		mlx.Eval(out)
+	}
+}
+
+// BenchmarkRepeatKV benchmarks K/V repetition.
+func BenchmarkRepeatKV(b *testing.B) {
+	// Typical GQA setup: 8 kv heads -> 32 heads
+	x := mlx.RandomNormal([]int32{1, 8, 512, 128}, 42)
+	mlx.Eval(x)
+
+	b.ResetTimer()
+	for range b.N {
+		out := RepeatKV(x, 4)
+		mlx.Eval(out)
+	}
+}

x/imagegen/runner.go

@@ -0,0 +1,181 @@
+// Package imagegen provides a unified MLX runner for both LLM and image generation models.
+package imagegen
+
+import (
+	"context"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"log/slog"
+	"net/http"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
+	"github.com/ollama/ollama/envconfig"
+	"github.com/ollama/ollama/x/imagegen/mlx"
+	"github.com/ollama/ollama/x/internal/mlxthread"
+)
+
+// Execute is the entry point for the unified MLX runner subprocess.
+func Execute(args []string) error {
+	// Set up logging with appropriate level from environment
+	slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: envconfig.LogLevel()})))
+
+	fs := flag.NewFlagSet("mlx-runner", flag.ExitOnError)
+	modelName := fs.String("model", "", "path to model")
+	port := fs.Int("port", 0, "port to listen on")
+
+	if err := fs.Parse(args); err != nil {
+		return err
+	}
+
+	if *modelName == "" {
+		return fmt.Errorf("--model is required")
+	}
+	if *port == 0 {
+		return fmt.Errorf("--port is required")
+	}
+
+	// Detect model type from capabilities
+	mode := detectModelMode(*modelName)
+	slog.Info("starting mlx runner", "model", *modelName, "port", *port, "mode", mode)
+
+	if mode != ModeImageGen {
+		return fmt.Errorf("imagegen runner only supports image generation models")
+	}
+
+	worker, err := mlxthread.Start("imagegen", func() error {
+		if err := mlx.InitMLX(); err != nil {
+			slog.Error("unable to initialize MLX", "error", err)
+			return err
+		}
+		slog.Info("MLX library initialized")
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+
+	// Create and start server
+	var server *server
+	if err := worker.Do(context.Background(), func() error {
+		var err error
+		server, err = newServer(*modelName, *port)
+		if err != nil {
+			return fmt.Errorf("failed to create server: %w", err)
+		}
+		server.mlxThread = worker
+		return nil
+	}); err != nil {
+		return err
+	}
+
+	// Set up HTTP handlers
+	mux := http.NewServeMux()
+	mux.HandleFunc("/health", server.healthHandler)
+	mux.HandleFunc("/completion", server.completionHandler)
+
+	httpServer := &http.Server{
+		Addr:    fmt.Sprintf("127.0.0.1:%d", *port),
+		Handler: mux,
+	}
+
+	// Handle shutdown
+	done := make(chan struct{})
+	go func() {
+		sigCh := make(chan os.Signal, 1)
+		signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
+		<-sigCh
+		slog.Info("shutting down mlx runner")
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		if err := httpServer.Shutdown(ctx); err != nil {
+			slog.Warn("graceful shutdown timed out", "error", err)
+			if err := httpServer.Close(); err != nil {
+				slog.Warn("failed to close http server", "error", err)
+			}
+		}
+		if err := worker.Stop(ctx, func() {
+			mlx.ClearCache()
+		}); err != nil {
+			slog.Warn("failed to stop mlx worker", "error", err)
+		}
+		close(done)
+	}()
+
+	slog.Info("mlx runner listening", "addr", httpServer.Addr)
+	if err := httpServer.ListenAndServe(); err != http.ErrServerClosed {
+		return err
+	}
+
+	<-done
+	return nil
+}
+
+// detectModelMode determines whether a model is an LLM or image generation model.
+func detectModelMode(modelName string) ModelMode {
+	// Check for image generation model by looking at model_index.json
+	modelType := DetectModelType(modelName)
+	if modelType != "" {
+		// Known image generation model types
+		switch modelType {
+		case "ZImagePipeline", "FluxPipeline", "Flux2KleinPipeline":
+			return ModeImageGen
+		}
+	}
+
+	// Default to LLM mode for safetensors models without known image gen types
+	return ModeLLM
+}
+
+// server holds the model and handles HTTP requests.
+type server struct {
+	modelName string
+	port      int
+	mlxThread *mlxthread.Thread
+
+	// Image generation model.
+	imageModel ImageModel
+}
+
+// newServer creates a new server instance for image generation models.
+func newServer(modelName string, port int) (*server, error) {
+	s := &server{
+		modelName: modelName,
+		port:      port,
+	}
+
+	if err := s.loadImageModel(); err != nil {
+		return nil, fmt.Errorf("failed to load image model: %w", err)
+	}
+
+	return s, nil
+}
+
+func (s *server) healthHandler(w http.ResponseWriter, r *http.Request) {
+	resp := HealthResponse{Status: "ok"}
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(resp)
+}
+
+func (s *server) completionHandler(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodPost {
+		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
+		return
+	}
+
+	var req Request
+	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+
+	if err := s.mlxThread.Do(r.Context(), func() error {
+		s.handleImageCompletion(w, r, req)
+		return nil
+	}); err != nil && r.Context().Err() == nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+	}
+}

x/imagegen/safetensors/loader.go

@@ -0,0 +1,429 @@
+package safetensors
+
+import (
+	"fmt"
+	"reflect"
+	"strings"
+
+	"github.com/ollama/ollama/x/imagegen/mlx"
+	"github.com/ollama/ollama/x/imagegen/nn"
+)
+
+// WeightSource is an interface for loading weights.
+// Both ModelWeights (directory-based) and ManifestWeights (blob-based) implement this.
+type WeightSource interface {
+	GetTensor(name string) (*mlx.Array, error)
+	ListTensors() []string
+	HasTensor(name string) bool
+	Quantization() string // Returns "NVFP4", "INT4", "INT8", or ""
+	GroupSize() int       // Returns quantization group size, or 0 if not specified
+}
+
+// QuantizationParams returns groupSize, bits, mode for a quantization type.
+// MLX quantization modes:
+//   - "affine": scale + zero-point bias, group_size=32/64/128
+//   - "nvfp4": NVIDIA FP4 with E4M3 scales, group_size=16 (no bias)
+//   - "mxfp8": Microsoft MX FP8 with E4M3 scales, group_size=32 (no bias)
+func QuantizationParams(quantization string) (groupSize, bits int, mode string) {
+	switch strings.ToUpper(quantization) {
+	case "NVFP4":
+		// NVIDIA FP4: group_size=16, bits=4, E4M3 scales (no qbias)
+		return 16, 4, "nvfp4"
+	case "FP4", "Q4", "INT4":
+		// 4-bit quantization with affine mode (scale + qbias)
+		return 32, 4, "affine"
+	case "MXFP8":
+		// Microsoft MX FP8: group_size=32, bits=8, E4M3 scales (no qbias)
+		return 32, 8, "mxfp8"
+	case "FP8", "Q8", "INT8":
+		// 8-bit quantization with affine mode (default for quantized models)
+		return 64, 8, "affine"
+	case "":
+		return 0, 0, ""
+	default:
+		return 32, 8, "affine" // Default to affine
+	}
+}
+
+// Transformer allows structs to transform weight arrays before assignment.
+// Implement this to apply operations like transpose during loading.
+type Transformer interface {
+	Transform(field string, arr *mlx.Array) *mlx.Array
+}
+
+// LoadModule loads weights into a struct using reflection and struct tags.
+//
+// Struct tags use the format: `weight:"path[,optional]"`
+//   - path: the weight name suffix (appended to prefix)
+//   - optional: if present, missing weights don't cause errors
+//   - "-": skip this field entirely
+//   - no tag on struct pointer: recurse with current prefix
+//   - no tag on *mlx.Array: skip (computed fields don't need loading)
+//
+// For slices of struct pointers, the loader iterates with .0, .1, .2... suffixes.
+// The slice must be pre-allocated to the correct length.
+//
+// Example:
+//
+//	type Attention struct {
+//	    QProj *nn.Linear  `weight:"self_attn.q_proj"`
+//	    KProj *nn.Linear  `weight:"self_attn.k_proj"`
+//	    Cache *mlx.Array  // no tag = skipped (computed field)
+//	}
+//
+//	err := LoadModule(&attn, weights, "model.layers.0")
+func LoadModule(dst any, weights WeightSource, prefix string) error {
+	v := reflect.ValueOf(dst)
+	if v.Kind() != reflect.Ptr || v.IsNil() {
+		return fmt.Errorf("LoadModule: dst must be a non-nil pointer")
+	}
+	v = v.Elem()
+	if v.Kind() != reflect.Struct {
+		return fmt.Errorf("LoadModule: dst must be a pointer to struct, got %v", v.Kind())
+	}
+
+	var errs []string
+	loadStruct(v, weights, prefix, &errs, false)
+
+	if len(errs) > 0 {
+		return fmt.Errorf("LoadModule: missing weights:\n  %s", strings.Join(errs, "\n  "))
+	}
+	return nil
+}
+
+// loadStruct recursively loads weights into a struct value.
+func loadStruct(v reflect.Value, weights WeightSource, prefix string, errs *[]string, parentOptional bool) {
+	t := v.Type()
+
+	for i := 0; i < t.NumField(); i++ {
+		field := t.Field(i)
+		fieldVal := v.Field(i)
+
+		// Skip unexported fields
+		if !fieldVal.CanSet() {
+			continue
+		}
+
+		// Parse tag
+		tag, hasTag := field.Tag.Lookup("weight")
+		if tag == "-" {
+			continue
+		}
+
+		// Parse tag options
+		optional := parentOptional
+		weightPath := tag
+		if idx := strings.Index(tag, ","); idx != -1 {
+			weightPath = tag[:idx]
+			if strings.Contains(tag[idx+1:], "optional") {
+				optional = true
+			}
+		}
+
+		// Build full path
+		fullPath := joinPath(prefix, weightPath)
+
+		// For struct pointers without a tag, recurse with current prefix
+		if !hasTag && fieldVal.Kind() == reflect.Ptr {
+			elemType := fieldVal.Type().Elem()
+			if elemType.Kind() == reflect.Struct && elemType != reflect.TypeOf(mlx.Array{}) {
+				if fieldVal.IsNil() {
+					fieldVal.Set(reflect.New(elemType))
+				}
+				loadStruct(fieldVal.Elem(), weights, prefix, errs, optional)
+				continue
+			}
+		}
+
+		// Handle nn.LinearLayer interface fields specially
+		linearLayerType := reflect.TypeOf((*nn.LinearLayer)(nil)).Elem()
+		if field.Type == linearLayerType {
+			if !hasTag {
+				continue // no tag = skip
+			}
+			layer, err := LoadLinearLayer(weights, fullPath)
+			if err != nil {
+				if !optional {
+					*errs = append(*errs, fullPath+": "+err.Error())
+				}
+				continue
+			}
+			fieldVal.Set(reflect.ValueOf(layer))
+			continue
+		}
+
+		// Handle nn.MultiLinearLayer interface fields specially
+		multiLinearLayerType := reflect.TypeOf((*nn.MultiLinearLayer)(nil)).Elem()
+		if field.Type == multiLinearLayerType {
+			if !hasTag {
+				continue // no tag = skip
+			}
+			layer, err := LoadMultiLinearLayer(weights, fullPath)
+			if err != nil {
+				if !optional {
+					*errs = append(*errs, fullPath+": "+err.Error())
+				}
+				continue
+			}
+			fieldVal.Set(reflect.ValueOf(layer))
+			continue
+		}
+
+		// Handle by kind
+		switch fieldVal.Kind() {
+		case reflect.Ptr:
+			elemType := fieldVal.Type().Elem()
+
+			// *mlx.Array - load directly (but skip if no tag - computed fields)
+			if fieldVal.Type() == reflect.TypeOf((*mlx.Array)(nil)) {
+				if !hasTag {
+					continue // no tag on *mlx.Array = computed field, skip
+				}
+				arr, err := weights.GetTensor(fullPath)
+				if err != nil {
+					if !optional {
+						*errs = append(*errs, fullPath)
+					}
+					continue
+				}
+				// Transform before assigning if parent implements Transformer
+				if t, ok := v.Addr().Interface().(Transformer); ok {
+					arr = t.Transform(field.Name, arr)
+				}
+				fieldVal.Set(reflect.ValueOf(arr))
+				continue
+			}
+
+			// Pointer to struct - allocate and recurse
+			if elemType.Kind() == reflect.Struct {
+				if optional && !hasWeightsWithPrefix(weights, fullPath) {
+					continue
+				}
+				if fieldVal.IsNil() {
+					fieldVal.Set(reflect.New(elemType))
+				}
+				loadStruct(fieldVal.Elem(), weights, fullPath, errs, optional)
+			}
+
+		case reflect.Slice:
+			elemType := fieldVal.Type().Elem()
+			if elemType.Kind() == reflect.Ptr && elemType.Elem().Kind() == reflect.Struct {
+				loadSlice(fieldVal, weights, fullPath, errs)
+			}
+		}
+	}
+}
+
+// hasWeightsWithPrefix checks if any weights exist with the given prefix.
+func hasWeightsWithPrefix(weights WeightSource, prefix string) bool {
+	for _, name := range weights.ListTensors() {
+		if strings.HasPrefix(name, prefix+".") || name == prefix {
+			return true
+		}
+	}
+	return false
+}
+
+// loadSlice loads weights into each element of a slice of struct pointers.
+func loadSlice(v reflect.Value, weights WeightSource, prefix string, errs *[]string) {
+	elemStructType := v.Type().Elem().Elem()
+
+	for i := 0; i < v.Len(); i++ {
+		elem := v.Index(i)
+		if elem.IsNil() {
+			elem.Set(reflect.New(elemStructType))
+		}
+		loadStruct(elem.Elem(), weights, fmt.Sprintf("%s.%d", prefix, i), errs, false)
+	}
+}
+
+// joinPath joins path segments with dots, handling empty segments.
+func joinPath(prefix, suffix string) string {
+	if prefix == "" {
+		return suffix
+	}
+	if suffix == "" {
+		return prefix
+	}
+	return prefix + "." + suffix
+}
+
+// LoadMultiLinearLayer loads a per-head linear layer from weights.
+// Weight shape should be [num_heads, output_dims, input_dims].
+// If quantized, always dequantizes since batched quantized matmul isn't supported.
+func LoadMultiLinearLayer(weights WeightSource, path string) (nn.MultiLinearLayer, error) {
+	// Check if this is a quantized layer by looking for scale tensor
+	scalePath := path + ".weight_scale"
+	hasScale := weights.HasTensor(scalePath)
+
+	weight, err := weights.GetTensor(path + ".weight")
+	if err != nil {
+		return nil, fmt.Errorf("failed to load weight %s: %w", path, err)
+	}
+
+	if hasScale {
+		scales, err := weights.GetTensor(scalePath)
+		if err != nil {
+			return nil, fmt.Errorf("failed to load scales %s: %w", scalePath, err)
+		}
+
+		var qbiases *mlx.Array
+		qbiasPath := path + ".weight_qbias"
+		if weights.HasTensor(qbiasPath) {
+			qbiases, _ = weights.GetTensor(qbiasPath)
+		}
+
+		// Always dequantize for MultiLinear - no batched quantized matmul support
+		// Detect bits from tensor shapes (supports mixed-precision Q4/Q8)
+		weightShape := weight.Shape()
+		scalesShape := scales.Shape()
+		weightCols := int(weightShape[len(weightShape)-1])
+		scalesCols := int(scalesShape[len(scalesShape)-1])
+
+		// Detect quantization from tensor shapes
+		// groupSize = weightCols * packFactor / scalesCols
+		// Note: groupSize4 = 2 * groupSize8 always, so ambiguous cases need metadata
+		groupSize4 := weightCols * 8 / scalesCols
+		groupSize8 := weightCols * 4 / scalesCols
+
+		var bits, groupSize int
+		// Use metadata to help disambiguate when shapes are ambiguous
+		// (e.g., Q4 with group_size=64 has same shapes as Q8 with group_size=32)
+		quantType := strings.ToUpper(weights.Quantization())
+		isQ8Type := quantType == "Q8" || quantType == "FP8" || quantType == "INT8"
+
+		if groupSize4 == 32 {
+			// Unambiguous: Q4 with group_size=32
+			bits = 4
+			groupSize = 32
+		} else if groupSize8 == 64 {
+			// Unambiguous: Q8 with group_size=64
+			bits = 8
+			groupSize = 64
+		} else if groupSize4 == 64 && groupSize8 == 32 {
+			// Ambiguous: could be Q4/gs=64 or Q8/gs=32, use metadata
+			if isQ8Type {
+				bits = 8
+				groupSize = 32
+			} else {
+				bits = 4
+				groupSize = 64
+			}
+		} else {
+			// Fallback: use global quantization params
+			_, bits, _ = QuantizationParams(weights.Quantization())
+			packFactor := 32 / bits
+			groupSize = weightCols * packFactor / scalesCols
+		}
+		weight = mlx.Dequantize(weight, scales, qbiases, groupSize, bits, "affine")
+	}
+
+	return nn.NewMultiLinear(weight), nil
+}
+
+// LoadLinearLayer loads a linear layer from weights, automatically detecting if it's quantized.
+// If {path}.weight_scale exists, creates a QuantizedLinear layer (or dequantizes if no kernel support).
+func LoadLinearLayer(weights WeightSource, path string) (nn.LinearLayer, error) {
+	// Check if this is a quantized layer by looking for scale tensor
+	scalePath := path + ".weight_scale"
+	hasScale := weights.HasTensor(scalePath)
+	if hasScale {
+		weight, err := weights.GetTensor(path + ".weight")
+		if err != nil {
+			return nil, fmt.Errorf("failed to load quantized weight %s: %w", path, err)
+		}
+
+		scales, err := weights.GetTensor(scalePath)
+		if err != nil {
+			return nil, fmt.Errorf("failed to load scales %s: %w", scalePath, err)
+		}
+
+		// Bias is optional
+		var bias *mlx.Array
+		biasPath := path + ".bias"
+		if weights.HasTensor(biasPath) {
+			bias, _ = weights.GetTensor(biasPath)
+		}
+
+		var qbiases *mlx.Array
+		qbiasPath := path + ".weight_qbias"
+		if weights.HasTensor(qbiasPath) {
+			qbiases, _ = weights.GetTensor(qbiasPath)
+		}
+
+		// Detect bits from tensor shapes (supports mixed-precision Q4/Q8)
+		weightShape := weight.Shape()
+		scalesShape := scales.Shape()
+		weightCols := int(weightShape[len(weightShape)-1])
+		scalesCols := int(scalesShape[len(scalesShape)-1])
+
+		// Detect quantization from tensor shapes
+		// groupSize = weightCols * packFactor / scalesCols
+		// Note: groupSize4 = 2 * groupSize8 always, so ambiguous cases need metadata
+		groupSize4 := weightCols * 8 / scalesCols
+		groupSize8 := weightCols * 4 / scalesCols
+
+		var bits, groupSize int
+		mode := "affine"
+		// Use metadata to help disambiguate when shapes are ambiguous
+		// (e.g., Q4 with group_size=64 has same shapes as Q8 with group_size=32)
+		quantType := strings.ToUpper(weights.Quantization())
+		isQ8Type := quantType == "Q8" || quantType == "FP8" || quantType == "INT8"
+
+		if groupSize4 == 32 {
+			// Unambiguous: Q4 with group_size=32
+			bits = 4
+			groupSize = 32
+		} else if groupSize8 == 64 {
+			// Unambiguous: Q8 with group_size=64
+			bits = 8
+			groupSize = 64
+		} else if groupSize4 == 64 && groupSize8 == 32 {
+			// Ambiguous: could be Q4/gs=64 or Q8/gs=32, use metadata
+			if isQ8Type {
+				bits = 8
+				groupSize = 32
+			} else {
+				bits = 4
+				groupSize = 64
+			}
+		} else {
+			// Fallback: use global quantization params
+			_, bits, mode = QuantizationParams(weights.Quantization())
+			packFactor := 32 / bits
+			groupSize = weightCols * packFactor / scalesCols
+		}
+
+		// NVFP4 and MXFP8 don't have native quantized matmul kernels in MLX,
+		// so we always dequantize at load time. Affine modes (FP4, FP8) have kernel support.
+		if mlx.MetalIsAvailable() && mode != "nvfp4" && mode != "mxfp8" {
+			return &nn.QuantizedLinear{
+				Weight:    weight,
+				Scales:    scales,
+				QBiases:   qbiases,
+				Bias:      bias,
+				GroupSize: groupSize,
+				Bits:      bits,
+				Mode:      mode,
+			}, nil
+		}
+
+		dequantized := mlx.Dequantize(weight, scales, qbiases, groupSize, bits, mode)
+		return nn.NewLinear(dequantized, bias), nil
+	}
+
+	// Load as regular Linear
+	weight, err := weights.GetTensor(path + ".weight")
+	if err != nil {
+		return nil, fmt.Errorf("failed to load weight %s: %w", path, err)
+	}
+
+	// Bias is optional
+	var bias *mlx.Array
+	biasPath := path + ".bias"
+	if weights.HasTensor(biasPath) {
+		bias, _ = weights.GetTensor(biasPath)
+	}
+
+	return nn.NewLinear(weight, bias), nil
+}

x/imagegen/safetensors/safetensors.go

@@ -0,0 +1,318 @@
+package safetensors
+
+import (
+	"encoding/binary"
+	"encoding/json"
+	"fmt"
+	"os"
+	"path/filepath"
+	"sort"
+	"strings"
+
+	"github.com/ollama/ollama/x/imagegen/mlx"
+)
+
+// SafetensorHeader represents the JSON header of a safetensors file
+type SafetensorHeader map[string]TensorInfo
+
+// TensorInfo contains metadata about a tensor
+type TensorInfo struct {
+	Dtype       string  `json:"dtype"`
+	Shape       []int32 `json:"shape"`
+	DataOffsets [2]int  `json:"data_offsets"`
+}
+
+// parseSafetensorHeader reads only the JSON header from a safetensors file.
+func parseSafetensorHeader(path string) (SafetensorHeader, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return nil, fmt.Errorf("failed to open file: %w", err)
+	}
+	defer f.Close()
+
+	var headerSize uint64
+	if err := binary.Read(f, binary.LittleEndian, &headerSize); err != nil {
+		return nil, fmt.Errorf("failed to read header size: %w", err)
+	}
+
+	headerBytes := make([]byte, headerSize)
+	if _, err := f.Read(headerBytes); err != nil {
+		return nil, fmt.Errorf("failed to read header: %w", err)
+	}
+
+	var header SafetensorHeader
+	if err := json.Unmarshal(headerBytes, &header); err != nil {
+		return nil, fmt.Errorf("failed to parse header: %w", err)
+	}
+
+	delete(header, "__metadata__")
+	return header, nil
+}
+
+// dtypeFromString converts safetensors dtype string to mlx.Dtype
+func dtypeFromString(s string) mlx.Dtype {
+	switch strings.ToUpper(s) {
+	case "F32", "FLOAT32":
+		return mlx.DtypeFloat32
+	case "F16", "FLOAT16":
+		return mlx.DtypeFloat16
+	case "BF16", "BFLOAT16":
+		return mlx.DtypeBFloat16
+	case "I32", "INT32":
+		return mlx.DtypeInt32
+	case "I64", "INT64":
+		return mlx.DtypeInt64
+	case "U8", "UINT8":
+		return mlx.DtypeUint8
+	case "F8_E4M3", "F8_E5M2", "F8_E4M3FN", "F8_E5M2FNUZ":
+		return mlx.DtypeUint8 // FP8 types stored as raw uint8 bytes
+	default:
+		return mlx.DtypeFloat32
+	}
+}
+
+// ModelWeights manages weights from multiple safetensor files.
+type ModelWeights struct {
+	dir         string                          // Model directory
+	tensorFiles map[string]string               // tensor name -> file path
+	tensorInfo  map[string]TensorInfo           // tensor name -> metadata
+	nativeCache map[string]*mlx.SafetensorsFile // file path -> loaded native handle
+	cache       map[string]*mlx.Array           // tensor name -> array (after Load)
+}
+
+// LoadModelWeights scans safetensor files and builds a tensor index.
+// This only reads JSON headers, not tensor data.
+func LoadModelWeights(dir string) (*ModelWeights, error) {
+	mw := &ModelWeights{
+		dir:         dir,
+		tensorFiles: make(map[string]string),
+		tensorInfo:  make(map[string]TensorInfo),
+		nativeCache: make(map[string]*mlx.SafetensorsFile),
+	}
+
+	entries, err := os.ReadDir(dir)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read directory: %w", err)
+	}
+
+	for _, entry := range entries {
+		if strings.HasSuffix(entry.Name(), ".safetensors") {
+			path := filepath.Join(dir, entry.Name())
+
+			header, err := parseSafetensorHeader(path)
+			if err != nil {
+				return nil, fmt.Errorf("failed to parse %s: %w", entry.Name(), err)
+			}
+
+			for name, info := range header {
+				mw.tensorFiles[name] = path
+				mw.tensorInfo[name] = info
+			}
+		}
+	}
+
+	if len(mw.tensorFiles) == 0 {
+		return nil, fmt.Errorf("no safetensor files found in %s", dir)
+	}
+
+	return mw, nil
+}
+
+// LoadModelWeightsFromPaths loads weights from specific safetensor file paths.
+// Used for loading from blob storage where files are not in a directory.
+func LoadModelWeightsFromPaths(paths []string) (*ModelWeights, error) {
+	mw := &ModelWeights{
+		tensorFiles: make(map[string]string),
+		tensorInfo:  make(map[string]TensorInfo),
+		nativeCache: make(map[string]*mlx.SafetensorsFile),
+	}
+
+	for _, path := range paths {
+		header, err := parseSafetensorHeader(path)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse %s: %w", path, err)
+		}
+
+		for name, info := range header {
+			mw.tensorFiles[name] = path
+			mw.tensorInfo[name] = info
+		}
+	}
+
+	if len(mw.tensorFiles) == 0 {
+		return nil, fmt.Errorf("no tensors found in provided paths")
+	}
+
+	return mw, nil
+}
+
+// Load loads all tensors into cache with the specified dtype.
+// If dtype is 0, tensors are loaded in their original dtype.
+// Automatically uses streaming (memory-efficient) when dtype conversion is needed,
+// or native loading when tensors are already in the target dtype.
+func (mw *ModelWeights) Load(dtype mlx.Dtype) error {
+	if dtype == 0 {
+		return mw.loadNative()
+	}
+
+	// Check if any tensor needs conversion
+	needsConversion := false
+	for name := range mw.tensorFiles {
+		info := mw.tensorInfo[name]
+		if dtypeFromString(info.Dtype) != dtype {
+			needsConversion = true
+			break
+		}
+	}
+
+	if needsConversion {
+		return mw.loadStreaming(dtype)
+	}
+	return mw.loadNative()
+}
+
+// loadNative loads all tensors using the native memory-mapped loader.
+func (mw *ModelWeights) loadNative() error {
+	mw.cache = make(map[string]*mlx.Array)
+
+	fileToTensors := make(map[string][]string)
+	for name, path := range mw.tensorFiles {
+		fileToTensors[path] = append(fileToTensors[path], name)
+	}
+
+	for path, names := range fileToTensors {
+		native, err := mlx.LoadSafetensorsNative(path)
+		if err != nil {
+			return fmt.Errorf("failed to load %s: %w", path, err)
+		}
+
+		for _, name := range names {
+			arr := native.Get(name)
+			if arr == nil {
+				native.Free()
+				return fmt.Errorf("tensor %q not found in %s", name, path)
+			}
+			mw.cache[name] = arr
+		}
+
+		mw.nativeCache[path] = native
+	}
+
+	return nil
+}
+
+// loadStreaming loads tensors with dtype conversion.
+// Uses the same pattern as Python: replace each entry in the map after conversion,
+// so the original tensor loses its reference and can be freed.
+func (mw *ModelWeights) loadStreaming(dtype mlx.Dtype) error {
+	mw.cache = make(map[string]*mlx.Array)
+
+	fileToTensors := make(map[string][]string)
+	for name, path := range mw.tensorFiles {
+		fileToTensors[path] = append(fileToTensors[path], name)
+	}
+
+	for path, names := range fileToTensors {
+		native, err := mlx.LoadSafetensorsNative(path)
+		if err != nil {
+			return fmt.Errorf("failed to load %s: %w", path, err)
+		}
+
+		for _, name := range names {
+			src := native.Get(name)
+			if src == nil {
+				native.Free()
+				return fmt.Errorf("tensor %q not found in %s", name, path)
+			}
+
+			dst := mlx.AsType(src, dtype)
+			mlx.Eval(dst)
+			native.Set(name, dst)
+			mw.cache[name] = dst
+		}
+
+		native.Free()
+	}
+
+	return nil
+}
+
+// Get returns a tensor from cache. Call Load() first.
+func (mw *ModelWeights) Get(name string) (*mlx.Array, error) {
+	if mw.cache == nil {
+		return nil, fmt.Errorf("cache not initialized: call Load() first")
+	}
+	arr, ok := mw.cache[name]
+	if !ok {
+		return nil, fmt.Errorf("tensor %q not found in cache", name)
+	}
+	return arr, nil
+}
+
+// GetTensor loads a tensor using the native loader without caching.
+// For bulk loading, use Load() + Get() instead.
+func (mw *ModelWeights) GetTensor(name string) (*mlx.Array, error) {
+	if mw.cache != nil {
+		if arr, ok := mw.cache[name]; ok {
+			return arr, nil
+		}
+	}
+
+	path, ok := mw.tensorFiles[name]
+	if !ok {
+		return nil, fmt.Errorf("tensor %q not found", name)
+	}
+
+	native, ok := mw.nativeCache[path]
+	if !ok {
+		var err error
+		native, err = mlx.LoadSafetensorsNative(path)
+		if err != nil {
+			return nil, fmt.Errorf("failed to load %s: %w", path, err)
+		}
+		mw.nativeCache[path] = native
+	}
+
+	return native.Get(name), nil
+}
+
+// GetTensorInfo returns metadata about a tensor without loading it.
+func (mw *ModelWeights) GetTensorInfo(name string) (TensorInfo, bool) {
+	info, ok := mw.tensorInfo[name]
+	return info, ok
+}
+
+// ListTensors returns all tensor names.
+func (mw *ModelWeights) ListTensors() []string {
+	names := make([]string, 0, len(mw.tensorFiles))
+	for name := range mw.tensorFiles {
+		names = append(names, name)
+	}
+	sort.Strings(names)
+	return names
+}
+
+// HasTensor checks if a tensor exists.
+func (mw *ModelWeights) HasTensor(name string) bool {
+	_, ok := mw.tensorFiles[name]
+	return ok
+}
+
+// Quantization returns empty string for directory-based weights (not quantized).
+func (mw *ModelWeights) Quantization() string {
+	return ""
+}
+
+// GroupSize returns 0 for directory-based weights (use default).
+func (mw *ModelWeights) GroupSize() int {
+	return 0
+}
+
+// ReleaseAll releases all cached native file handles.
+func (mw *ModelWeights) ReleaseAll() {
+	for path, native := range mw.nativeCache {
+		native.Free()
+		delete(mw.nativeCache, path)
+	}
+}
+

x/imagegen/safetensors/safetensors_test.go

@@ -0,0 +1,165 @@
+package safetensors
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/ollama/ollama/x/imagegen/mlx"
+)
+
+func TestLoadModelWeights(t *testing.T) {
+	// Skip if no model available
+	modelDir := "../weights/gpt-oss-20b"
+	if _, err := os.Stat(modelDir); os.IsNotExist(err) {
+		t.Skip("model weights not available")
+	}
+
+	mw, err := LoadModelWeights(modelDir)
+	if err != nil {
+		t.Fatalf("LoadModelWeights: %v", err)
+	}
+	defer mw.ReleaseAll()
+
+	// Check we found tensors
+	tensors := mw.ListTensors()
+	if len(tensors) == 0 {
+		t.Fatal("no tensors found")
+	}
+	t.Logf("found %d tensors", len(tensors))
+
+	// Check HasTensor
+	if !mw.HasTensor(tensors[0]) {
+		t.Errorf("HasTensor(%q) = false", tensors[0])
+	}
+	if mw.HasTensor("nonexistent.weight") {
+		t.Error("HasTensor returned true for nonexistent tensor")
+	}
+}
+
+func TestGetTensor(t *testing.T) {
+	modelDir := "../weights/gpt-oss-20b"
+	if _, err := os.Stat(modelDir); os.IsNotExist(err) {
+		t.Skip("model weights not available")
+	}
+
+	mw, err := LoadModelWeights(modelDir)
+	if err != nil {
+		t.Fatalf("LoadModelWeights: %v", err)
+	}
+	defer mw.ReleaseAll()
+
+	tensors := mw.ListTensors()
+	if len(tensors) == 0 {
+		t.Skip("no tensors")
+	}
+
+	// Load first tensor
+	arr, err := mw.GetTensor(tensors[0])
+	if err != nil {
+		t.Fatalf("GetTensor(%q): %v", tensors[0], err)
+	}
+
+	// Verify it has a shape
+	shape := arr.Shape()
+	if len(shape) == 0 {
+		t.Error("tensor has no shape")
+	}
+	t.Logf("%s: shape=%v dtype=%v", tensors[0], shape, arr.Dtype())
+}
+
+func TestLoadWithDtype(t *testing.T) {
+	modelDir := "../weights/gpt-oss-20b"
+	if _, err := os.Stat(modelDir); os.IsNotExist(err) {
+		t.Skip("model weights not available")
+	}
+
+	mw, err := LoadModelWeights(modelDir)
+	if err != nil {
+		t.Fatalf("LoadModelWeights: %v", err)
+	}
+	defer mw.ReleaseAll()
+
+	// Load all tensors as bfloat16
+	if err := mw.Load(mlx.DtypeBFloat16); err != nil {
+		t.Fatalf("Load: %v", err)
+	}
+
+	// Get a tensor from cache
+	tensors := mw.ListTensors()
+	arr, err := mw.Get(tensors[0])
+	if err != nil {
+		t.Fatalf("Get: %v", err)
+	}
+
+	// Verify dtype (unless it was already bf16)
+	t.Logf("%s: dtype=%v", tensors[0], arr.Dtype())
+}
+
+func TestLookupTensor(t *testing.T) {
+	modelDir := "../weights/gpt-oss-20b"
+	if _, err := os.Stat(modelDir); os.IsNotExist(err) {
+		t.Skip("model weights not available")
+	}
+
+	mw, err := LoadModelWeights(modelDir)
+	if err != nil {
+		t.Fatalf("LoadModelWeights: %v", err)
+	}
+	defer mw.ReleaseAll()
+
+	// HasTensor returns false for nonexistent
+	if mw.HasTensor("nonexistent") {
+		t.Error("HasTensor should return false for nonexistent")
+	}
+
+	// HasTensor returns true for existing tensor
+	tensors := mw.ListTensors()
+	if !mw.HasTensor(tensors[0]) {
+		t.Error("HasTensor should return true for existing tensor")
+	}
+}
+
+func TestParseSafetensorHeader(t *testing.T) {
+	modelDir := "../weights/gpt-oss-20b"
+	if _, err := os.Stat(modelDir); os.IsNotExist(err) {
+		t.Skip("model weights not available")
+	}
+
+	// Find a safetensors file
+	entries, err := os.ReadDir(modelDir)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var stFile string
+	for _, e := range entries {
+		if filepath.Ext(e.Name()) == ".safetensors" {
+			stFile = filepath.Join(modelDir, e.Name())
+			break
+		}
+	}
+	if stFile == "" {
+		t.Skip("no safetensors file found")
+	}
+
+	header, err := parseSafetensorHeader(stFile)
+	if err != nil {
+		t.Fatalf("parseSafetensorHeader: %v", err)
+	}
+
+	if len(header) == 0 {
+		t.Error("header is empty")
+	}
+
+	// Check a tensor has valid info
+	for name, info := range header {
+		if info.Dtype == "" {
+			t.Errorf("%s: empty dtype", name)
+		}
+		if len(info.Shape) == 0 {
+			t.Errorf("%s: empty shape", name)
+		}
+		break // just check one
+	}
+}

x/imagegen/server.go

@@ -0,0 +1,463 @@
+package imagegen
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"log/slog"
+	"math/rand"
+	"net"
+	"net/http"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"runtime"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/ollama/ollama/envconfig"
+	"github.com/ollama/ollama/format"
+	"github.com/ollama/ollama/llm"
+	"github.com/ollama/ollama/ml"
+	"github.com/ollama/ollama/x/imagegen/manifest"
+)
+
+// Server wraps an MLX runner subprocess to implement llm.LlamaServer.
+//
+// This implementation is compatible with Ollama's scheduler and can be loaded/unloaded
+// like any other model. It is used for image generation models.
+type Server struct {
+	mu          sync.Mutex
+	cmd         *exec.Cmd
+	port        int
+	modelName   string
+	vramSize    uint64
+	done        chan error
+	client      *http.Client
+	lastErr     string // Last stderr line for error reporting
+	lastErrLock sync.Mutex
+}
+
+// NewServer prepares a new MLX runner server for image generation models.
+// The subprocess is not started until Load() is called.
+func NewServer(modelName string) (*Server, error) {
+	// Validate platform support before attempting to start
+	if err := CheckPlatformSupport(); err != nil {
+		return nil, err
+	}
+
+	return &Server{
+		modelName: modelName,
+		done:      make(chan error, 1),
+		client:    &http.Client{Timeout: 10 * time.Minute},
+	}, nil
+}
+
+// ModelPath returns the path to the model.
+func (s *Server) ModelPath() string {
+	return s.modelName
+}
+
+// Load checks whether the model fits in GPU memory and starts the subprocess.
+func (s *Server) Load(ctx context.Context, _ ml.SystemInfo, gpus []ml.DeviceInfo, requireFull bool) ([]ml.DeviceID, error) {
+	// Estimate VRAM based on tensor size from manifest
+	if modelManifest, err := manifest.LoadManifest(s.modelName); err == nil {
+		s.vramSize = uint64(modelManifest.TotalTensorSize())
+	} else {
+		s.vramSize = 8 * 1024 * 1024 * 1024
+	}
+
+	if len(gpus) > 0 {
+		available := gpus[0].FreeMemory
+		overhead := gpus[0].MinimumMemory() + envconfig.GpuOverhead()
+		if available > overhead {
+			available -= overhead
+		} else {
+			available = 0
+		}
+
+		if s.vramSize > available {
+			if requireFull {
+				return nil, llm.ErrLoadRequiredFull
+			}
+			return nil, fmt.Errorf("model requires %s but only %s are available (after %s overhead)", format.HumanBytes2(s.vramSize), format.HumanBytes2(available), format.HumanBytes2(overhead))
+		}
+	}
+
+	// Find a free port
+	port := 0
+	if a, err := net.ResolveTCPAddr("tcp", "localhost:0"); err == nil {
+		if l, err := net.ListenTCP("tcp", a); err == nil {
+			port = l.Addr().(*net.TCPAddr).Port
+			l.Close()
+		}
+	}
+	if port == 0 {
+		port = rand.Intn(65535-49152) + 49152
+	}
+	s.port = port
+
+	// Get the current executable path (we use the same binary with runner subcommand)
+	exe, err := os.Executable()
+	if err != nil {
+		return nil, fmt.Errorf("unable to lookup executable path: %w", err)
+	}
+	if eval, err := filepath.EvalSymlinks(exe); err == nil {
+		exe = eval
+	}
+
+	// Spawn subprocess: ollama runner --imagegen-engine --model <path> --port <port>
+	cmd := exec.Command(exe, "runner", "--imagegen-engine", "--model", s.modelName, "--port", strconv.Itoa(port))
+	cmd.Env = os.Environ()
+	configureMLXSubprocessEnv(cmd, ml.LibraryPaths(gpus))
+
+	s.cmd = cmd
+
+	// Forward subprocess stdout/stderr to server logs
+	stdout, _ := cmd.StdoutPipe()
+	stderr, _ := cmd.StderrPipe()
+	go func() {
+		scanner := bufio.NewScanner(stdout)
+		for scanner.Scan() {
+			slog.Info("mlx-runner", "msg", scanner.Text())
+		}
+	}()
+	go func() {
+		scanner := bufio.NewScanner(stderr)
+		for scanner.Scan() {
+			line := scanner.Text()
+			slog.Warn("mlx-runner", "msg", line)
+			s.lastErrLock.Lock()
+			s.lastErr = line
+			s.lastErrLock.Unlock()
+		}
+	}()
+
+	slog.Info("starting mlx runner subprocess", "model", s.modelName, "port", s.port)
+	if err := cmd.Start(); err != nil {
+		return nil, fmt.Errorf("failed to start mlx runner: %w", err)
+	}
+
+	// Reap subprocess when it exits
+	go func() {
+		err := cmd.Wait()
+		s.done <- err
+	}()
+
+	return nil, nil
+}
+
+// Ping checks if the subprocess is healthy.
+func (s *Server) Ping(ctx context.Context) error {
+	url := fmt.Sprintf("http://127.0.0.1:%d/health", s.port)
+	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
+	if err != nil {
+		return err
+	}
+	resp, err := s.client.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("health check failed: %d", resp.StatusCode)
+	}
+	return nil
+}
+
+func mlxLibraryPathEnv() string {
+	switch runtime.GOOS {
+	case "windows":
+		return "PATH"
+	case "darwin":
+		return "DYLD_LIBRARY_PATH"
+	default:
+		return "LD_LIBRARY_PATH"
+	}
+}
+
+func configureMLXSubprocessEnv(cmd *exec.Cmd, libraryPaths []string) {
+	if len(libraryPaths) == 0 {
+		return
+	}
+
+	// Search order for the imagegen runner is:
+	//   1. bundled lib/ollama root
+	//   2. backend-specific library dirs selected during GPU discovery
+	//   3. any existing caller-provided library path values
+	pathEnv := mlxLibraryPathEnv()
+	pathEnvPaths := append([]string{}, libraryPaths...)
+	if existingPath, ok := os.LookupEnv(pathEnv); ok {
+		pathEnvPaths = append(pathEnvPaths, filepath.SplitList(existingPath)...)
+	}
+	setSubprocessEnv(cmd, pathEnv, strings.Join(pathEnvPaths, string(filepath.ListSeparator)))
+	slog.Debug("mlx subprocess library path", pathEnv, strings.Join(pathEnvPaths, string(filepath.ListSeparator)))
+
+	ollamaLibraryPaths := append([]string{}, libraryPaths...)
+	if existingPath, ok := os.LookupEnv("OLLAMA_LIBRARY_PATH"); ok {
+		ollamaLibraryPaths = append(ollamaLibraryPaths, filepath.SplitList(existingPath)...)
+	}
+	setSubprocessEnv(cmd, "OLLAMA_LIBRARY_PATH", strings.Join(ollamaLibraryPaths, string(filepath.ListSeparator)))
+	slog.Debug("mlx subprocess library path", "OLLAMA_LIBRARY_PATH", strings.Join(ollamaLibraryPaths, string(filepath.ListSeparator)))
+}
+
+func setSubprocessEnv(cmd *exec.Cmd, key, value string) {
+	for i := range cmd.Env {
+		name, _, ok := strings.Cut(cmd.Env[i], "=")
+		if ok && strings.EqualFold(name, key) {
+			cmd.Env[i] = key + "=" + value
+			return
+		}
+	}
+	cmd.Env = append(cmd.Env, key+"="+value)
+}
+
+// getLastErr returns the last stderr line.
+func (s *Server) getLastErr() string {
+	s.lastErrLock.Lock()
+	defer s.lastErrLock.Unlock()
+	return s.lastErr
+}
+
+// WaitUntilRunning waits for the subprocess to be ready.
+func (s *Server) WaitUntilRunning(ctx context.Context) error {
+	timeout := time.After(envconfig.LoadTimeout())
+	ticker := time.NewTicker(100 * time.Millisecond)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case err := <-s.done:
+			errMsg := s.getLastErr()
+			if errMsg != "" {
+				return fmt.Errorf("mlx runner failed: %s (exit: %v)", errMsg, err)
+			}
+			return fmt.Errorf("mlx runner exited unexpectedly: %w", err)
+		case <-timeout:
+			errMsg := s.getLastErr()
+			if errMsg != "" {
+				return fmt.Errorf("timeout waiting for mlx runner: %s", errMsg)
+			}
+			return errors.New("timeout waiting for mlx runner to start")
+		case <-ticker.C:
+			if err := s.Ping(ctx); err == nil {
+				slog.Info("mlx runner is ready", "port", s.port)
+				return nil
+			}
+		}
+	}
+}
+
+// Completion handles both text and image generation requests.
+func (s *Server) Completion(ctx context.Context, req llm.CompletionRequest, fn func(llm.CompletionResponse)) error {
+	seed := req.Seed
+	if seed == 0 {
+		seed = time.Now().UnixNano()
+	}
+
+	// Extract raw image bytes from llm.ImageData slice
+	var images [][]byte
+	for _, img := range req.Images {
+		images = append(images, img.Data)
+	}
+
+	// Build request for subprocess
+	creq := Request{
+		Prompt: req.Prompt,
+		Width:  req.Width,
+		Height: req.Height,
+		Steps:  int(req.Steps),
+		Seed:   seed,
+		Images: images,
+	}
+
+	// Pass LLM options if present
+	if req.Options != nil {
+		creq.Options = &RequestOptions{
+			NumPredict:  req.Options.NumPredict,
+			Temperature: float64(req.Options.Temperature),
+			TopP:        float64(req.Options.TopP),
+			TopK:        req.Options.TopK,
+			Stop:        req.Options.Stop,
+		}
+	}
+
+	body, err := json.Marshal(creq)
+	if err != nil {
+		return err
+	}
+
+	url := fmt.Sprintf("http://127.0.0.1:%d/completion", s.port)
+	httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
+	if err != nil {
+		return err
+	}
+	httpReq.Header.Set("Content-Type", "application/json")
+
+	resp, err := s.client.Do(httpReq)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+		return fmt.Errorf("%s", strings.TrimSpace(string(body)))
+	}
+
+	scanner := bufio.NewScanner(resp.Body)
+	scanner.Buffer(make([]byte, 1024*1024), 16*1024*1024) // 16MB max
+	for scanner.Scan() {
+		// Parse subprocess response
+		var raw struct {
+			Image              string `json:"image,omitempty"`
+			Content            string `json:"content,omitempty"`
+			Done               bool   `json:"done"`
+			Step               int    `json:"step,omitempty"`
+			Total              int    `json:"total,omitempty"`
+			StopReason         string `json:"stop_reason,omitempty"`
+			PromptEvalCount    int    `json:"prompt_eval_count,omitempty"`
+			PromptEvalDuration int    `json:"prompt_eval_duration,omitempty"`
+			EvalCount          int    `json:"eval_count,omitempty"`
+			EvalDuration       int    `json:"eval_duration,omitempty"`
+		}
+		if err := json.Unmarshal(scanner.Bytes(), &raw); err != nil {
+			slog.Debug("mlx response parse error", "error", err, "line", string(scanner.Bytes()))
+			continue
+		}
+
+		// Log stop reason when generation completes
+		if raw.Done && raw.StopReason != "" {
+			slog.Info("mlx generation completed", "stop_reason", raw.StopReason)
+		}
+
+		// Convert to llm.CompletionResponse
+		cresp := llm.CompletionResponse{
+			Content:            raw.Content,
+			Done:               raw.Done,
+			Step:               raw.Step,
+			TotalSteps:         raw.Total,
+			Image:              raw.Image,
+			PromptEvalCount:    raw.PromptEvalCount,
+			PromptEvalDuration: time.Duration(raw.PromptEvalDuration),
+			EvalCount:          raw.EvalCount,
+			EvalDuration:       time.Duration(raw.EvalDuration),
+		}
+
+		fn(cresp)
+		if cresp.Done {
+			return nil
+		}
+	}
+
+	// Scanner exited without receiving Done - connection was likely closed
+	scanErr := scanner.Err()
+	if scanErr != nil {
+		slog.Error("mlx scanner error", "error", scanErr)
+	} else {
+		slog.Warn("mlx scanner EOF without Done response - subprocess may have crashed")
+	}
+
+	// Check if subprocess is still alive
+	if s.HasExited() {
+		slog.Error("mlx subprocess has exited unexpectedly")
+		if errMsg := s.getLastErr(); errMsg != "" {
+			return fmt.Errorf("mlx runner closed response before completion: %s", errMsg)
+		}
+	}
+
+	if scanErr != nil {
+		return scanErr
+	}
+
+	return errors.New("mlx runner closed response before completion")
+}
+
+// Close terminates the subprocess.
+func (s *Server) Close() error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if s.cmd != nil && s.cmd.Process != nil {
+		slog.Info("stopping mlx runner subprocess", "pid", s.cmd.Process.Pid)
+		s.cmd.Process.Signal(os.Interrupt)
+
+		// Wait briefly for graceful shutdown
+		select {
+		case <-s.done:
+		case <-time.After(5 * time.Second):
+			s.cmd.Process.Kill()
+		}
+		s.cmd = nil
+	}
+	return nil
+}
+
+// MemorySize returns the total and VRAM memory usage.
+func (s *Server) MemorySize() (total, vram uint64) {
+	return s.vramSize, s.vramSize
+}
+
+// VRAMByGPU returns VRAM usage for a specific GPU.
+func (s *Server) VRAMByGPU(id ml.DeviceID) uint64 {
+	return s.vramSize
+}
+
+// ContextLength returns the context length (not applicable for image generation).
+func (s *Server) ContextLength() int {
+	return 0
+}
+
+// Embedding returns embeddings for the input.
+func (s *Server) Embedding(ctx context.Context, input string) ([]float32, int, error) {
+	return nil, 0, errors.New("embeddings not supported for MLX models")
+}
+
+// Tokenize tokenizes the input content.
+func (s *Server) Tokenize(ctx context.Context, content string) ([]int, error) {
+	return nil, errors.New("tokenization not supported for image generation models")
+}
+
+// Detokenize converts tokens back to text.
+func (s *Server) Detokenize(ctx context.Context, tokens []int) (string, error) {
+	return "", errors.New("detokenization not supported for MLX models")
+}
+
+// Pid returns the process ID of the subprocess.
+func (s *Server) Pid() int {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.cmd != nil && s.cmd.Process != nil {
+		return s.cmd.Process.Pid
+	}
+	return -1
+}
+
+// GetPort returns the port the subprocess is listening on.
+func (s *Server) GetPort() int {
+	return s.port
+}
+
+// GetDeviceInfos returns device information.
+func (s *Server) GetDeviceInfos(ctx context.Context) []ml.DeviceInfo {
+	return nil
+}
+
+// HasExited returns whether the subprocess has exited.
+func (s *Server) HasExited() bool {
+	select {
+	case <-s.done:
+		return true
+	default:
+		return false
+	}
+}
+
+// Ensure Server implements llm.LlamaServer
+var _ llm.LlamaServer = (*Server)(nil)

x/imagegen/server_test.go

@@ -0,0 +1,102 @@
+package imagegen
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"strings"
+	"testing"
+
+	"github.com/ollama/ollama/llm"
+)
+
+type roundTripFunc func(*http.Request) (*http.Response, error)
+
+func (fn roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
+	return fn(req)
+}
+
+func newCompletionTestServer(handler func(*http.Request) string) *Server {
+	return &Server{
+		port: 11434,
+		done: make(chan error, 1),
+		client: &http.Client{
+			Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
+				body := handler(req)
+				return &http.Response{
+					StatusCode: http.StatusOK,
+					Header:     make(http.Header),
+					Body:       io.NopCloser(strings.NewReader(body)),
+					Request:    req,
+				}, nil
+			}),
+		},
+	}
+}
+
+func TestCompletionReturnsImageData(t *testing.T) {
+	s := newCompletionTestServer(func(r *http.Request) string {
+		if r.URL.Path != "/completion" {
+			t.Fatalf("path = %q, want /completion", r.URL.Path)
+		}
+
+		var req Request
+		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+			t.Fatal(err)
+		}
+		if req.Prompt != "test prompt" || req.Width != 512 || req.Height != 256 || req.Steps != 7 || req.Seed != 42 {
+			t.Fatalf("unexpected request: %+v", req)
+		}
+		if len(req.Images) != 1 || string(req.Images[0]) != "input-image" {
+			t.Fatalf("images = %q, want input-image", req.Images)
+		}
+
+		return `{"step":1,"total":2}` + "\n" +
+			`{"done":true,"image":"base64png"}` + "\n"
+	})
+
+	var responses []llm.CompletionResponse
+	err := s.Completion(context.Background(), llm.CompletionRequest{
+		Prompt: "test prompt",
+		Width:  512,
+		Height: 256,
+		Steps:  7,
+		Seed:   42,
+		Images: []llm.ImageData{{Data: []byte("input-image")}},
+	}, func(resp llm.CompletionResponse) {
+		responses = append(responses, resp)
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(responses) != 2 {
+		t.Fatalf("responses = %d, want 2", len(responses))
+	}
+	if responses[0].Step != 1 || responses[0].TotalSteps != 2 || responses[0].Done {
+		t.Fatalf("progress response = %+v", responses[0])
+	}
+	if !responses[1].Done || responses[1].Image != "base64png" {
+		t.Fatalf("final response = %+v", responses[1])
+	}
+}
+
+func TestCompletionEOFBeforeDoneReturnsError(t *testing.T) {
+	s := newCompletionTestServer(func(r *http.Request) string {
+		return `{"step":1,"total":2}` + "\n"
+	})
+
+	var responses []llm.CompletionResponse
+	err := s.Completion(context.Background(), llm.CompletionRequest{Prompt: "test prompt"}, func(resp llm.CompletionResponse) {
+		responses = append(responses, resp)
+	})
+	if err == nil {
+		t.Fatal("expected error")
+	}
+	if !strings.Contains(err.Error(), "closed response before completion") {
+		t.Fatalf("error = %v", err)
+	}
+	if len(responses) != 1 || responses[0].Done {
+		t.Fatalf("responses = %+v, want one non-done progress response", responses)
+	}
+}

x/imagegen/tokenizer/README.md

@@ -0,0 +1,85 @@
+# Tokenizer
+
+Tokenizer for LLM inference supporting BPE, SentencePiece, and WordPiece algorithms. The goal of this package is to see if a pure Go tokenizer can be fast and correct. It primarily supports the `imagegen` models however it (or parts of it) could be considered to replace Ollama's tokenizer in the `model` package.
+
+## Features
+
+- **BPE (Byte Pair Encoding)** - GPT-2/Llama style with byte-level encoding
+- **SentencePiece** - Gemma style with `▁` space handling
+- **WordPiece** - BERT style with `##` continuation tokens
+- **Parallel encoding** - Automatic parallelization for inputs >4KB
+- **HuggingFace compatible** - Loads `tokenizer.json` directly
+
+## Usage
+
+```go
+import "github.com/ollama/ollama/x/imagegen/tokenizer"
+
+// Load from HuggingFace model directory
+tok, err := tokenizer.Load("./weights/Llama-3.2-1B")
+if err != nil {
+    log.Fatal(err)
+}
+
+// Encode text to token IDs
+ids := tok.Encode("Hello, world!", false) // false = don't add BOS
+
+// Decode back to text
+text := tok.Decode(ids)
+
+// Check special tokens
+if tok.IsEOS(ids[len(ids)-1]) {
+    // End of sequence
+}
+```
+
+## Performance
+
+Benchmarks on Apple M3 Max:
+
+| Input Size | Encode | Decode | Tokens |
+|------------|--------|--------|--------|
+| 1 KB | 14.5 MB/s | 267 MB/s | 231 |
+| 10 KB | 10.9 MB/s | 321 MB/s | 2,301 |
+| 100 KB | 8.9 MB/s | 311 MB/s | 23,001 |
+| 1 MB | 9.6 MB/s | 321 MB/s | 230,001 |
+
+Comparison with other implementations (10 MB input):
+
+| Implementation | Encode Speed | Notes |
+|----------------|--------------|-------|
+| Engine (this) | ~10 MB/s | stdlib RE2, parallel >4KB |
+| tiktoken (Rust) | ~17 MB/s | Highly optimized regex |
+| Ollama (Go) | ~2-3 MB/s | regexp2 backtracking |
+
+## Performance Opportunities
+
+Potential optimizations not yet implemented:
+
+| Optimization | Expected Gain | Complexity |
+|--------------|---------------|------------|
+| Aho-Corasick for special tokens | 2-3x for many special tokens | Medium |
+| Custom regex engine (like tiktoken) | 1.5-2x | High |
+| SIMD byte scanning | 1.3-1.5x for pretokenizer | Medium |
+| Assembly BPE merge loop | 1.2-1.5x | High |
+| Memoization for repeated substrings | Variable | Low |
+
+Current bottleneck is the pretokenizer regex (~60% of encode time). tiktoken achieves ~17 MB/s with a hand-tuned Rust regex engine.
+
+## Not Yet Implemented
+
+| Feature | Used By | Notes |
+|---------|---------|-------|
+| Unigram tokenizer | T5, ALBERT, mBART | Different algorithm (not BPE) |
+| Unicode normalizers | Some multilingual models | NFD, NFKC, lowercase, etc. |
+| Custom pretokenizers | Model-specific | Beyond standard patterns |
+
+Most HuggingFace models use BPE or SentencePiece, which are fully supported. WordPiece (BERT-style) is also supported with standard `[UNK]` fallback for out-of-vocabulary characters.
+
+## Files
+
+| File | Description |
+|------|-------------|
+| `tokenizer.go` | Main implementation (~1000 lines) |
+| `tokenizer_test.go` | Tests and benchmarks |
+| `testdata/` | Mini tokenizer for unit tests |

x/imagegen/tokenizer/testdata/mini_llama.json

@@ -0,0 +1 @@
+{"model": {"type": "BPE", "vocab": {"!": 0, "\"": 1, "#": 2, "$": 3, "%": 4, "&": 5, "'": 6, "(": 7, ")": 8, "*": 9, "+": 10, ",": 11, "-": 12, ".": 13, "/": 14, "0": 15, "1": 16, "2": 17, "3": 18, "4": 19, "5": 20, "6": 21, "7": 22, "8": 23, "9": 24, ":": 25, ";": 26, "<": 27, "=": 28, ">": 29, "?": 30, "@": 31, "A": 32, "B": 33, "C": 34, "D": 35, "E": 36, "F": 37, "G": 38, "H": 39, "I": 40, "J": 41, "K": 42, "L": 43, "M": 44, "N": 45, "O": 46, "P": 47, "Q": 48, "R": 49, "S": 50, "T": 51, "U": 52, "V": 53, "fé": 59958, "W": 54, "X": 55, "Y": 56, "Z": 57, "[": 58, "\\": 59, "]": 60, "^": 61, "_": 62, "`": 63, "a": 64, "b": 65, "c": 66, "d": 67, "e": 68, "f": 69, "g": 70, "h": 71, "i": 72, "j": 73, "k": 74, "l": 75, "m": 76, "n": 77, "o": 78, "p": 79, "r": 81, "q": 80, "s": 82, "t": 83, "u": 84, "v": 85, "w": 86, "x": 87, "y": 88, "z": 89, "{": 90, "|": 91, "}": 92, "~": 93, "¡": 94, "¢": 95, "£": 96, "¤": 97, "¥": 98, "¦": 99, "§": 100, "¨": 101, "World": 10343, "©": 102, "ª": 103, "«": 104, "¬": 105, "®": 106, "world": 14957, "¯": 107, "°": 108, "±": 109, "²": 110, "³": 111, "´": 112, "µ": 113, "¶": 114, "·": 115, "¸": 116, "¹": 117, "º": 118, "»": 119, "¼": 120, "½": 121, "¾": 122, "¿": 123, "À": 124, "Á": 125, "Â": 126, "Ã": 127, "Ä": 128, "Å": 129, "Æ": 130, "Ç": 131, "È": 132, "É": 133, "Ê": 134, "Ë": 135, "Ì": 136, "Í": 137, "Î": 138, "Ï": 139, "Ð": 140, "Ñ": 141, "Ò": 142, "Ó": 143, "Ô": 144, "Õ": 145, "Ö": 146, "×": 147, "Ø": 148, "Ù": 149, "Ú": 150, "Û": 151, "Ü": 152, "Ý": 153, "Þ": 154, "ß": 155, "à": 156, "á": 157, "â": 158, "ã": 159, "ä": 160, "å": 161, "æ": 162, "ç": 163, "è": 164, "é": 165, "ê": 166, "ë": 167, "ì": 168, "Ġhello": 24748, "í": 169, "î": 170, "ï": 171, "ð": 172, "ñ": 173, "Hello": 9906, "ò": 174, "ó": 175, "ô": 176, "õ": 177, "ö": 178, "Ġ{}": 4792, "÷": 179, "ø": 180, "ù": 181, "ú": 182, "û": 183, "ü": 184, "ý": 185, "þ": 186, "ÿ": 187, "Ā": 188, "ā": 189, "Ă": 190, "ă": 191, "Ċ": 198, "Ą": 192, "ą": 193, "Ć": 194, "ć": 195, "Ĉ": 196, "ĉ": 197, "ċ": 199, "Č": 200, "č": 201, "Ď": 202, "ď": 203, "Đ": 204, "đ": 205, "Ē": 206, "ē": 207, "Ĕ": 208, "ĕ": 209, "Ė": 210, "ė": 211, "Ę": 212, "ę": 213, "Ġ": 220, "Ě": 214, "ě": 215, "Ĝ": 216, "ĝ": 217, "Ğ": 218, "ğ": 219, "ġ": 221, "Ģ": 222, "ģ": 223, "Ĥ": 224, "ĥ": 225, "Ħ": 226, "ħ": 227, "Ĩ": 228, "ĩ": 229, "Ī": 230, "ī": 231, "Ĭ": 232, "ĭ": 233, "Į": 234, "į": 235, "İ": 236, "ı": 237, "IJ": 238, "ij": 239, "Ĵ": 240, "ĵ": 241, "Ķ": 242, "ķ": 243, "ĸ": 244, "Ĺ": 245, "ĺ": 246, "Ļ": 247, "ļ": 248, "Ľ": 249, "ĠĠ": 256, "ľ": 250, "Ŀ": 251, "ŀ": 252, "Ł": 253, "rer": 38149, "ĠĠĠ": 262, "ł": 254, "Ń": 255, "'m": 2846, "'re": 2351, "can": 4919, "func": 2900, "()": 368, "Ġworld": 1917, "Ġmain": 1925, "00": 410, "123": 4513, "000": 931, "ca": 936, "'t": 956, "é": 978, "hello": 15339, "Ġw": 289, "orld": 1410, "Ġwor": 4191, "ld": 509, "main": 3902, "Ġm": 296, "ain": 467, "Ġma": 7643, "in": 258, "Ġmai": 17154, "re": 265, "'r": 97670, "unc": 1371, "fun": 12158, "fu": 33721, "nc": 1031, "ma": 1764, "mai": 77585, "wor": 50810, "or": 269, "Ġwo": 24670, "23": 1419, "12": 717, "{}": 6390, "Ġ{": 314, "an": 276, "ello": 4896, "Hel": 33813, "lo": 385, "Hell": 81394, "un": 359, "hel": 50222, "hell": 57195, "ai": 2192, "wo": 1146, "Ġh": 305, "Ġhel": 11591, "Ġhell": 15123, "el": 301, "He": 1548, "er": 261, "he": 383, "ell": 616, "ll": 657}, "merges": ["Ġ Ġ", "Ġ ĠĠ", "ĠĠ Ġ", "( )", "0 0", "0 00", "00 0", "c a", "' t", "à ©", "Ġ world", "Ġw orld", "Ġwor ld", "Ġ main", "Ġm ain", "Ġma in", "Ġmai n", "' re", "'r e", "' m", "f unc", "fun c", "fu nc", "m ain", "ma in", "mai n", "Ġ wor", "Ġw or", "Ġwo r", "1 23", "12 3", "Ġ {}", "Ġ{ }", "c an", "ca n", "{ }", "Ġ ma", "Ġm a", "H ello", "Hel lo", "Hell o", "W orld", "f un", "fu n", "w orld", "wor ld", "h ello", "hel lo", "hell o", "Ġ mai", "Ġm ai", "Ġma i", "Ġ wo", "Ġw o", "Ġ hello", "Ġh ello", "Ġhel lo", "Ġhell o", "f u", "H el", "He l", "r er", "re r", "h el", "he l", "w or", "wo r", "h ell", "he ll", "hel l", "f é", "m ai", "ma i", "H ell", "He ll", "Hel l", "' r"]}, "pre_tokenizer": {"type": "Sequence", "pretokenizers": [{"type": "Split", "pattern": {"Regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"}, "behavior": "Isolated", "invert": false}, {"type": "ByteLevel", "add_prefix_space": false, "trim_offsets": true, "use_regex": false}]}, "decoder": {"type": "ByteLevel", "add_prefix_space": true, "trim_offsets": true, "use_regex": true}, "added_tokens": [{"id": 128000, "content": "<|begin_of_text|>", "special": true}, {"id": 128001, "content": "<|end_of_text|>", "special": true}]}

x/imagegen/tokenizer/tokenizer_test.go

@@ -0,0 +1,783 @@
+package tokenizer
+
+import (
+	"bytes"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"regexp"
+	"strings"
+	"sync"
+	"testing"
+)
+
+// TestPatternCompilation validates that HuggingFace pretokenizer patterns
+// can be rewritten for Go's RE2 regexp engine and compiled successfully.
+func TestPatternCompilation(t *testing.T) {
+	patterns := []struct {
+		name    string
+		pattern string
+	}{
+		{"llama3", `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`},
+		{"qwen2", `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`},
+		{"gpt4o", `[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+`},
+		{"gpt2", `'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+`},
+		{"deepseek_cjk", `[一-龥\x{3040}-ゟ゠-ヿ]+`},
+	}
+
+	for _, p := range patterns {
+		t.Run(p.name, func(t *testing.T) {
+			rewritten := rewritePatternForRE2(p.pattern)
+			if _, err := regexp.Compile(rewritten); err != nil {
+				t.Errorf("failed to compile pattern: %v\noriginal: %s\nrewritten: %s", err, p.pattern, rewritten)
+			}
+		})
+	}
+}
+
+// TestRoundtrip verifies the fundamental property: encode(text) -> decode -> text
+// This is the key invariant from tiktoken's test suite.
+func TestRoundtrip(t *testing.T) {
+	tok, err := Load("testdata/mini_llama.json")
+	if err != nil {
+		t.Fatalf("failed to load tokenizer: %v", err)
+	}
+
+	// Test cases covering key edge cases from tiktoken
+	inputs := []string{
+		// Empty and simple
+		"",
+		"a",
+		"hello",
+		"hello world",
+
+		// Whitespace edge cases
+		" ",
+		"  ",
+		"   ",
+		" hello",
+		"hello ",
+		" hello ",
+		"hello  world",
+		"hello   world",
+		"\t",
+		"\n",
+		"\r\n",
+		"hello\nworld",
+		"hello\n\nworld",
+
+		// Contractions
+		"don't",
+		"I'm",
+		"we'll",
+		"they're",
+		"it's",
+		"DON'T", // uppercase
+
+		// Numbers
+		"123",
+		"1234567890",
+		"3.14159",
+		"$100",
+		"50%",
+
+		// Unicode
+		"こんにちは",          // Japanese
+		"你好",              // Chinese
+		"مرحبا",            // Arabic (RTL)
+		"🎉",               // Emoji
+		"Hello 世界",        // Mixed
+		"café",             // Accented
+		"naïve",            // Diaeresis
+		"Ω≈ç√∫",            // Math symbols
+
+		// Code
+		"func main() {}",
+		"if (x == 0) { return; }",
+		"import \"fmt\"",
+		"x := 42",
+		"// comment",
+		"/* block */",
+
+		// Repetition (tiktoken specifically tests this)
+		"aaaa",
+		"aaaaaaaaaaaa",
+		strings.Repeat("a", 100),
+		strings.Repeat("hello ", 50),
+
+		// Punctuation
+		"...",
+		"!!!",
+		"???",
+		"hello, world!",
+		"(parentheses)",
+		"[brackets]",
+		"{braces}",
+
+		// Mixed complexity
+		"The quick brown fox jumps over the lazy dog.",
+		"Lorem ipsum dolor sit amet, consectetur adipiscing elit.",
+		"func TestRoundtrip(t *testing.T) { t.Run(\"test\", func(t *testing.T) {}) }",
+	}
+
+	for _, input := range inputs {
+		name := input
+		if len(name) > 30 {
+			name = name[:30] + "..."
+		}
+		if name == "" {
+			name = "<empty>"
+		}
+		name = strings.ReplaceAll(name, "\n", "\\n")
+		name = strings.ReplaceAll(name, "\t", "\\t")
+
+		t.Run(name, func(t *testing.T) {
+			tokens := tok.Encode(input, false)
+			decoded := tok.Decode(tokens)
+			if decoded != input {
+				t.Errorf("roundtrip failed:\n  input:   %q\n  tokens:  %v\n  decoded: %q", input, tokens, decoded)
+			}
+		})
+	}
+}
+
+// TestSpecialTokens verifies that special tokens are handled correctly
+func TestSpecialTokens(t *testing.T) {
+	tok, err := Load("testdata/mini_llama.json")
+	if err != nil {
+		t.Fatalf("failed to load tokenizer: %v", err)
+	}
+
+	// Special tokens should be preserved through encode/decode
+	t.Run("bos_preserved", func(t *testing.T) {
+		if tok.BOS() < 0 {
+			t.Skip("no BOS token")
+		}
+		tokens := tok.Encode("hello", true)
+		if len(tokens) == 0 || tokens[0] != tok.BOS() {
+			t.Errorf("BOS not prepended: got %v, want first token to be %d", tokens, tok.BOS())
+		}
+	})
+
+	t.Run("special_token_split", func(t *testing.T) {
+		// If we have special tokens, verify they're split correctly
+		for tokenStr, tokenID := range tok.specialTokens {
+			input := "before" + tokenStr + "after"
+			tokens := tok.Encode(input, false)
+
+			found := false
+			for _, id := range tokens {
+				if id == tokenID {
+					found = true
+					break
+				}
+			}
+			if !found {
+				t.Errorf("special token %q (id=%d) not found in encoding of %q: %v",
+					tokenStr, tokenID, input, tokens)
+			}
+		}
+	})
+}
+
+// TestConcurrency verifies thread-safe encoding
+func TestConcurrency(t *testing.T) {
+	tok, err := Load("testdata/mini_llama.json")
+	if err != nil {
+		t.Fatalf("failed to load tokenizer: %v", err)
+	}
+
+	input := "The quick brown fox jumps over the lazy dog."
+	expected := tok.Encode(input, false)
+
+	var wg sync.WaitGroup
+	errors := make(chan error, 100)
+
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			got := tok.Encode(input, false)
+			if len(got) != len(expected) {
+				errors <- nil // just signal error
+				return
+			}
+			for j := range got {
+				if got[j] != expected[j] {
+					errors <- nil
+					return
+				}
+			}
+		}()
+	}
+
+	wg.Wait()
+	close(errors)
+
+	if len(errors) > 0 {
+		t.Errorf("concurrent encoding produced inconsistent results")
+	}
+}
+
+// TestIntegration runs against real model directories, comparing with Python transformers.
+// Skips if model weights are not available.
+func TestIntegration(t *testing.T) {
+	models := []string{
+		"../weights/Llama-3.2-1B",
+		"../weights/gemma-3-1b-it",
+		"../weights/gpt-oss-20b",
+	}
+
+	// Test inputs covering various edge cases
+	inputs := []string{
+		"Hello, world!",
+		"The quick brown fox jumps over the lazy dog.",
+		"こんにちは世界",
+		"def fibonacci(n):\n    if n <= 1:\n        return n\n    return fibonacci(n-1) + fibonacci(n-2)",
+		"1234567890",
+		"   spaces   ",
+		"don't won't can't",
+	}
+
+	for _, modelPath := range models {
+		modelName := filepath.Base(modelPath)
+
+		t.Run(modelName, func(t *testing.T) {
+			tokenizerPath := filepath.Join(modelPath, "tokenizer.json")
+			if _, err := os.Stat(tokenizerPath); err != nil {
+				t.Skipf("skipping: %s not found", tokenizerPath)
+			}
+
+			tok, err := Load(tokenizerPath)
+			if err != nil {
+				t.Fatalf("failed to load tokenizer: %v", err)
+			}
+
+			for _, input := range inputs {
+				t.Run(truncate(input, 20), func(t *testing.T) {
+					// Test roundtrip
+					tokens := tok.Encode(input, false)
+					decoded := tok.Decode(tokens)
+					if decoded != input {
+						t.Errorf("roundtrip failed:\n  input:   %q\n  decoded: %q", input, decoded)
+					}
+
+					// Compare with Python if available
+					if pythonTokens, err := pythonEncode(modelPath, input); err == nil {
+						if !equalInt32Slice(tokens, pythonTokens) {
+							t.Errorf("mismatch with Python:\n  go:     %v\n  python: %v", tokens, pythonTokens)
+						}
+					}
+				})
+			}
+		})
+	}
+}
+
+// pythonEncode calls Python transformers to encode text, for comparison
+func pythonEncode(modelPath, text string) ([]int32, error) {
+	script := `
+import sys, json
+from transformers import AutoTokenizer
+tok = AutoTokenizer.from_pretrained(sys.argv[1])
+tokens = tok.encode(sys.argv[2], add_special_tokens=False)
+print(json.dumps(tokens))
+`
+	cmd := exec.Command("python3", "-c", script, modelPath, text)
+	var out bytes.Buffer
+	cmd.Stdout = &out
+	cmd.Stderr = nil
+
+	if err := cmd.Run(); err != nil {
+		return nil, err
+	}
+
+	// Parse JSON array
+	var tokens []int32
+	output := strings.TrimSpace(out.String())
+	if output == "" || output == "[]" {
+		return []int32{}, nil
+	}
+
+	// Simple parsing for [1, 2, 3] format
+	output = strings.Trim(output, "[]")
+	if output == "" {
+		return []int32{}, nil
+	}
+
+	for _, s := range strings.Split(output, ",") {
+		s = strings.TrimSpace(s)
+		var v int32
+		if _, err := parseIntSimple(s, &v); err == nil {
+			tokens = append(tokens, v)
+		}
+	}
+
+	return tokens, nil
+}
+
+func parseIntSimple(s string, v *int32) (bool, error) {
+	var n int64
+	for _, c := range s {
+		if c >= '0' && c <= '9' {
+			n = n*10 + int64(c-'0')
+		}
+	}
+	*v = int32(n)
+	return true, nil
+}
+
+func equalInt32Slice(a, b []int32) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	for i := range a {
+		if a[i] != b[i] {
+			return false
+		}
+	}
+	return true
+}
+
+func truncate(s string, n int) string {
+	if len(s) <= n {
+		return s
+	}
+	return s[:n] + "..."
+}
+
+// TestBPEPretokenizer verifies BPE pretokenizer splits text correctly
+// using the GPT-2 style regex pattern (no dependency on tokenizer files)
+func TestBPEPretokenizer(t *testing.T) {
+	pattern := `'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+`
+	re := regexp.MustCompile(rewritePatternForRE2(pattern))
+
+	tests := []struct {
+		input    string
+		expected []string
+	}{
+		{"Hello", []string{"Hello"}},
+		{"Hello world", []string{"Hello", " world"}},
+		{"Hello, world!", []string{"Hello", ",", " world", "!"}},
+		{"don't", []string{"don", "'t"}},
+		{"I'm", []string{"I", "'m"}},
+		{"123", []string{"123"}},
+		{"12345", []string{"12345"}}, // GPT-2 pattern matches any digit sequence
+		{"a  b", []string{"a", " ", " b"}}, // whitespace boundary: last space prepends to word
+		{"   ", []string{"   "}},           // pure whitespace stays together
+		{"\n\n", []string{"\n\n"}},         // newlines stay together
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.input, func(t *testing.T) {
+			// Get regex matches
+			matches := re.FindAllStringIndex(tt.input, -1)
+			var chunks []string
+			for _, m := range matches {
+				chunks = append(chunks, tt.input[m[0]:m[1]])
+			}
+
+			// Apply whitespace boundary fix (same logic as Encode)
+			for i := 0; i < len(chunks)-1; i++ {
+				if isNonNewlineWhitespace(chunks[i]) && len(chunks[i+1]) > 0 {
+					r, _ := []rune(chunks[i+1])[0], 0
+					if r >= 'A' && r <= 'z' { // simplified letter check
+						// Move last space to next chunk
+						if len(chunks[i]) > 0 {
+							lastSpace := chunks[i][len(chunks[i])-1:]
+							chunks[i] = chunks[i][:len(chunks[i])-1]
+							chunks[i+1] = lastSpace + chunks[i+1]
+						}
+					}
+				}
+			}
+
+			// Filter empty chunks
+			var result []string
+			for _, c := range chunks {
+				if c != "" {
+					result = append(result, c)
+				}
+			}
+
+			if len(result) != len(tt.expected) {
+				t.Errorf("got %v, want %v", result, tt.expected)
+				return
+			}
+			for i := range result {
+				if result[i] != tt.expected[i] {
+					t.Errorf("chunk %d: got %q, want %q", i, result[i], tt.expected[i])
+				}
+			}
+		})
+	}
+}
+
+// TestSentencePiecePretokenizer verifies SentencePiece doesn't use pretokenizer
+// and correctly replaces spaces with ▁ (no dependency on tokenizer files)
+func TestSentencePiecePretokenizer(t *testing.T) {
+	// SentencePiece has no pretokenizer - whole text is one chunk
+	// Spaces are replaced with ▁ during encoding
+
+	tests := []struct {
+		input    string
+		expected string // after space replacement
+	}{
+		{"Hello", "Hello"},
+		{"Hello world", "Hello▁world"},
+		{"Hello, world!", "Hello,▁world!"},
+		{"   spaces   ", "▁▁▁spaces▁▁▁"},
+		{" Hello", "▁Hello"},
+		{"Hello ", "Hello▁"},
+		{"a b c", "a▁b▁c"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.input, func(t *testing.T) {
+			// SentencePiece encoding: replace space with ▁
+			result := strings.ReplaceAll(tt.input, " ", "▁")
+			if result != tt.expected {
+				t.Errorf("got %q, want %q", result, tt.expected)
+			}
+		})
+	}
+}
+
+// TestWordPiecePretokenizer verifies WordPiece (BERT) pretokenizer splits correctly
+// BertPreTokenizer splits on whitespace and punctuation
+func TestWordPiecePretokenizer(t *testing.T) {
+	// BertPreTokenizer behavior: split on whitespace and punctuation
+	// Whitespace is stripped, punctuation becomes separate tokens
+
+	tests := []struct {
+		input    string
+		expected []string
+	}{
+		{"Hello", []string{"Hello"}},
+		{"Hello world", []string{"Hello", "world"}},           // whitespace stripped
+		{"Hello, world!", []string{"Hello", ",", "world", "!"}}, // punct separate
+		{"don't", []string{"don", "'", "t"}},                   // apostrophe separate (unlike BPE)
+		{"   spaces   ", []string{"spaces"}},                   // whitespace stripped
+		{"Hello.World", []string{"Hello", ".", "World"}},       // punct splits
+		{"test@email.com", []string{"test", "@", "email", ".", "com"}},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.input, func(t *testing.T) {
+			result := splitBertStyle(tt.input)
+			if len(result) != len(tt.expected) {
+				t.Errorf("got %v, want %v", result, tt.expected)
+				return
+			}
+			for i := range result {
+				if result[i] != tt.expected[i] {
+					t.Errorf("token %d: got %q, want %q", i, result[i], tt.expected[i])
+				}
+			}
+		})
+	}
+}
+
+// splitBertStyle mimics BertPreTokenizer: split on whitespace and punctuation
+func splitBertStyle(s string) []string {
+	var result []string
+	var current strings.Builder
+
+	for _, r := range s {
+		if r == ' ' || r == '\t' || r == '\n' || r == '\r' {
+			// Whitespace: flush current token, don't add whitespace
+			if current.Len() > 0 {
+				result = append(result, current.String())
+				current.Reset()
+			}
+		} else if isPunct(r) {
+			// Punctuation: flush current, add punct as separate token
+			if current.Len() > 0 {
+				result = append(result, current.String())
+				current.Reset()
+			}
+			result = append(result, string(r))
+		} else {
+			current.WriteRune(r)
+		}
+	}
+	if current.Len() > 0 {
+		result = append(result, current.String())
+	}
+	return result
+}
+
+func isPunct(r rune) bool {
+	// Common ASCII punctuation
+	return (r >= '!' && r <= '/') || (r >= ':' && r <= '@') ||
+		(r >= '[' && r <= '`') || (r >= '{' && r <= '~')
+}
+
+// TestRepeatedDigits verifies correct tokenization of repeated digit sequences.
+// Llama-style tokenizers split digits in groups of 1-3 due to the \p{N}{1,3} pattern.
+func TestRepeatedDigits(t *testing.T) {
+	tok, err := Load("./testdata/mini_llama.json")
+	if err != nil {
+		t.Skipf("mini_llama.json not available: %v", err)
+	}
+
+	// Pattern: 1 digit, 2 digits, 3 digits, then repeats
+	// "0" -> [single], "00" -> [double], "000" -> [triple]
+	// "0000" -> [triple, single], etc.
+	tests := []struct {
+		input string
+		count int // expected token count
+	}{
+		{"0", 1},
+		{"00", 1},
+		{"000", 1},
+		{"0000", 2},   // 3 + 1
+		{"00000", 2},  // 3 + 2
+		{"000000", 2}, // 3 + 3
+		{"0000000", 3},
+		{"00000000", 3},
+		{"000000000", 3},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.input, func(t *testing.T) {
+			ids := tok.Encode(tt.input, false)
+			if len(ids) != tt.count {
+				t.Errorf("Encode(%q) = %d tokens, want %d", tt.input, len(ids), tt.count)
+			}
+			// Verify roundtrip
+			decoded := tok.Decode(ids)
+			if decoded != tt.input {
+				t.Errorf("Decode(Encode(%q)) = %q", tt.input, decoded)
+			}
+		})
+	}
+}
+
+// TestNullByte verifies that null bytes roundtrip correctly
+func TestNullByte(t *testing.T) {
+	tok, err := Load("./testdata/mini_llama.json")
+	if err != nil {
+		t.Skipf("mini_llama.json not available: %v", err)
+	}
+
+	ids := tok.Encode("\x00", false)
+	decoded := tok.Decode(ids)
+	if decoded != "\x00" {
+		t.Errorf("null byte roundtrip failed: got %q, want %q", decoded, "\x00")
+	}
+}
+
+// TestTokenizerTypeDetection verifies correct detection of tokenizer types
+func TestTokenizerTypeDetection(t *testing.T) {
+	tests := []struct {
+		name     string
+		decoder  string
+		expected TokenizerType
+	}{
+		{
+			name:     "ByteLevel decoder (BPE)",
+			decoder:  `{"type": "ByteLevel"}`,
+			expected: TokenizerBPE,
+		},
+		{
+			name: "Sequence with Replace ▁ (SentencePiece)",
+			decoder: `{
+				"type": "Sequence",
+				"decoders": [
+					{"type": "Replace", "pattern": {"String": "▁"}, "content": " "}
+				]
+			}`,
+			expected: TokenizerSentencePiece,
+		},
+		{
+			name:     "null decoder (BPE default)",
+			decoder:  `null`,
+			expected: TokenizerBPE,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			isSPM := detectSentencePiece([]byte(tt.decoder))
+			var got TokenizerType
+			if isSPM {
+				got = TokenizerSentencePiece
+			} else {
+				got = TokenizerBPE
+			}
+			if got != tt.expected {
+				t.Errorf("got %v, want %v", got, tt.expected)
+			}
+		})
+	}
+}
+
+// TestPADTokenDefault verifies PAD() returns -1 when not configured
+func TestPADTokenDefault(t *testing.T) {
+	tok, err := Load("testdata/mini_llama.json")
+	if err != nil {
+		t.Fatalf("failed to load tokenizer: %v", err)
+	}
+
+	// mini_llama.json has no PAD token configured, should return -1
+	if got := tok.PAD(); got != -1 {
+		t.Errorf("PAD() = %d, want -1 (not configured)", got)
+	}
+}
+
+// TestPADTokenFromConfig verifies PAD token is loaded from tokenizer_config.json
+func TestPADTokenFromConfig(t *testing.T) {
+	// Create temp directory with tokenizer files
+	dir := t.TempDir()
+
+	// Write minimal tokenizer.json
+	tokenizerJSON := `{
+		"model": {
+			"type": "BPE",
+			"vocab": {"<|endoftext|>": 0, "hello": 1, "world": 2},
+			"merges": []
+		},
+		"added_tokens": [
+			{"id": 0, "content": "<|endoftext|>", "special": true}
+		]
+	}`
+	if err := os.WriteFile(filepath.Join(dir, "tokenizer.json"), []byte(tokenizerJSON), 0o644); err != nil {
+		t.Fatalf("failed to write tokenizer.json: %v", err)
+	}
+
+	// Write tokenizer_config.json with pad_token
+	configJSON := `{
+		"pad_token": "<|endoftext|>"
+	}`
+	if err := os.WriteFile(filepath.Join(dir, "tokenizer_config.json"), []byte(configJSON), 0o644); err != nil {
+		t.Fatalf("failed to write tokenizer_config.json: %v", err)
+	}
+
+	tok, err := Load(dir)
+	if err != nil {
+		t.Fatalf("failed to load tokenizer: %v", err)
+	}
+
+	if got := tok.PAD(); got != 0 {
+		t.Errorf("PAD() = %d, want 0 (<|endoftext|>)", got)
+	}
+}
+
+// TestPADTokenFromSpecialTokensMap verifies PAD falls back to special_tokens_map.json
+func TestPADTokenFromSpecialTokensMap(t *testing.T) {
+	dir := t.TempDir()
+
+	// Write minimal tokenizer.json
+	tokenizerJSON := `{
+		"model": {
+			"type": "BPE",
+			"vocab": {"<pad>": 0, "hello": 1, "world": 2},
+			"merges": []
+		},
+		"added_tokens": [
+			{"id": 0, "content": "<pad>", "special": true}
+		]
+	}`
+	if err := os.WriteFile(filepath.Join(dir, "tokenizer.json"), []byte(tokenizerJSON), 0o644); err != nil {
+		t.Fatalf("failed to write tokenizer.json: %v", err)
+	}
+
+	// Write special_tokens_map.json with pad_token
+	mapJSON := `{
+		"pad_token": "<pad>"
+	}`
+	if err := os.WriteFile(filepath.Join(dir, "special_tokens_map.json"), []byte(mapJSON), 0o644); err != nil {
+		t.Fatalf("failed to write special_tokens_map.json: %v", err)
+	}
+
+	tok, err := Load(dir)
+	if err != nil {
+		t.Fatalf("failed to load tokenizer: %v", err)
+	}
+
+	if got := tok.PAD(); got != 0 {
+		t.Errorf("PAD() = %d, want 0 (<pad>)", got)
+	}
+}
+
+// TestPADTokenWithContentObject verifies PAD token works with {"content": "..."} format
+func TestPADTokenWithContentObject(t *testing.T) {
+	dir := t.TempDir()
+
+	// Write minimal tokenizer.json
+	tokenizerJSON := `{
+		"model": {
+			"type": "BPE",
+			"vocab": {"[PAD]": 0, "hello": 1},
+			"merges": []
+		},
+		"added_tokens": [
+			{"id": 0, "content": "[PAD]", "special": true}
+		]
+	}`
+	if err := os.WriteFile(filepath.Join(dir, "tokenizer.json"), []byte(tokenizerJSON), 0o644); err != nil {
+		t.Fatalf("failed to write tokenizer.json: %v", err)
+	}
+
+	// Write tokenizer_config.json with pad_token as object (HuggingFace format)
+	configJSON := `{
+		"pad_token": {"content": "[PAD]", "lstrip": false, "normalized": false}
+	}`
+	if err := os.WriteFile(filepath.Join(dir, "tokenizer_config.json"), []byte(configJSON), 0o644); err != nil {
+		t.Fatalf("failed to write tokenizer_config.json: %v", err)
+	}
+
+	tok, err := Load(dir)
+	if err != nil {
+		t.Fatalf("failed to load tokenizer: %v", err)
+	}
+
+	if got := tok.PAD(); got != 0 {
+		t.Errorf("PAD() = %d, want 0 ([PAD])", got)
+	}
+}
+
+// Benchmarks
+
+func BenchmarkEncode(b *testing.B) {
+	tok, err := Load("testdata/mini_llama.json")
+	if err != nil {
+		b.Fatalf("failed to load tokenizer: %v", err)
+	}
+
+	inputs := []struct {
+		name string
+		text string
+	}{
+		{"short", "Hello, world!"},
+		{"medium", "The quick brown fox jumps over the lazy dog. " + strings.Repeat("This is a test. ", 10)},
+		{"long", strings.Repeat("The quick brown fox jumps over the lazy dog. ", 100)},
+	}
+
+	for _, input := range inputs {
+		b.Run(input.name, func(b *testing.B) {
+			b.SetBytes(int64(len(input.text)))
+			for i := 0; i < b.N; i++ {
+				tok.Encode(input.text, false)
+			}
+		})
+	}
+}
+
+func BenchmarkDecode(b *testing.B) {
+	tok, err := Load("testdata/mini_llama.json")
+	if err != nil {
+		b.Fatalf("failed to load tokenizer: %v", err)
+	}
+
+	text := strings.Repeat("The quick brown fox jumps over the lazy dog. ", 100)
+	tokens := tok.Encode(text, false)
+
+	b.SetBytes(int64(len(text)))
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		tok.Decode(tokens)
+	}
+}

x/imagegen/types.go

@@ -0,0 +1,81 @@
+// Package imagegen provides a unified MLX runner for both LLM and image generation models.
+//
+// This package handles safetensors models created with `ollama create --experimental`,
+// supporting both text generation (LLM) and image generation (diffusion) models
+// through a single unified interface.
+package imagegen
+
+// Request is the request format for completion requests.
+type Request struct {
+	Prompt string `json:"prompt"`
+
+	// LLM-specific fields
+	Options *RequestOptions `json:"options,omitempty"`
+
+	// Image generation fields
+	Width  int32    `json:"width,omitempty"`
+	Height int32    `json:"height,omitempty"`
+	Steps  int      `json:"steps,omitempty"`
+	Seed   int64    `json:"seed,omitempty"`
+	Images [][]byte `json:"images,omitempty"` // Input images for image editing/conditioning
+}
+
+// RequestOptions contains LLM-specific generation options.
+type RequestOptions struct {
+	NumPredict  int      `json:"num_predict,omitempty"`
+	Temperature float64  `json:"temperature,omitempty"`
+	TopP        float64  `json:"top_p,omitempty"`
+	TopK        int      `json:"top_k,omitempty"`
+	Stop        []string `json:"stop,omitempty"`
+}
+
+// Response is streamed back for each progress update.
+type Response struct {
+	// Text generation response
+	Content string `json:"content,omitempty"`
+
+	// Image generation response
+	Image string `json:"image,omitempty"` // Base64-encoded PNG
+
+	// Common fields
+	Done       bool   `json:"done"`
+	DoneReason int    `json:"done_reason,omitempty"`
+	StopReason string `json:"stop_reason,omitempty"` // Debug: why generation stopped
+
+	// Progress fields
+	Step  int `json:"step,omitempty"`
+	Total int `json:"total,omitempty"`
+
+	// Statistics
+	PromptEvalCount    int `json:"prompt_eval_count,omitempty"`
+	PromptEvalDuration int `json:"prompt_eval_duration,omitempty"`
+	EvalCount          int `json:"eval_count,omitempty"`
+	EvalDuration       int `json:"eval_duration,omitempty"`
+}
+
+// HealthResponse is returned by the health endpoint.
+type HealthResponse struct {
+	Status   string  `json:"status"`
+	Progress float32 `json:"progress,omitempty"`
+}
+
+// ModelMode represents the type of model being run.
+type ModelMode int
+
+const (
+	// ModeLLM indicates a text generation model.
+	ModeLLM ModelMode = iota
+	// ModeImageGen indicates an image generation model.
+	ModeImageGen
+)
+
+func (m ModelMode) String() string {
+	switch m {
+	case ModeLLM:
+		return "llm"
+	case ModeImageGen:
+		return "imagegen"
+	default:
+		return "unknown"
+	}
+}

x/imagegen/vae/tiling.go

@@ -0,0 +1,213 @@
+// Package vae provides shared utilities for VAE (Variational Autoencoder) operations.
+package vae
+
+import (
+	"github.com/ollama/ollama/x/imagegen/mlx"
+)
+
+// TilingConfig holds configuration for tiled VAE decoding.
+// This is a general technique to reduce memory usage when decoding large latents.
+type TilingConfig struct {
+	TileSize int32 // Tile size in latent space (e.g., 64 latent → 512 pixels for 8x VAE)
+	Overlap  int32 // Overlap in latent space (e.g., 16 latent = 25% of 64)
+}
+
+// DefaultTilingConfig returns reasonable defaults matching diffusers.
+// tile_latent_min_size=64, tile_overlap_factor=0.25
+func DefaultTilingConfig() *TilingConfig {
+	return &TilingConfig{
+		TileSize: 64, // 64 latent pixels
+		Overlap:  16, // 25% overlap
+	}
+}
+
+// decodedTile holds a decoded tile's pixel data and dimensions
+type decodedTile struct {
+	data   []float32
+	height int32
+	width  int32
+}
+
+// DecodeTiled decodes latents using tiled processing with overlap blending.
+// This reduces memory usage for large images by processing in overlapping tiles.
+//
+// Parameters:
+//   - latents: [1, H, W, C] latent tensor in NHWC format
+//   - cfg: tiling configuration (tile size and overlap)
+//   - decoder: function to decode a single tile [1, H, W, C] -> [1, H*scale, W*scale, 3]
+//
+// Returns: [1, 3, H*scale, W*scale] decoded image in NCHW format
+func DecodeTiled(latents *mlx.Array, cfg *TilingConfig, decoder func(*mlx.Array) *mlx.Array) *mlx.Array {
+	shape := latents.Shape()
+	H := shape[1] // latent height
+	W := shape[2] // latent width
+	C := shape[3]
+
+	tileLatentSize := cfg.TileSize
+	overlapLatent := cfg.Overlap
+
+	// If image is small enough, just decode normally
+	if H <= tileLatentSize && W <= tileLatentSize {
+		decoded := decoder(latents)
+		decoded = mlx.AsType(decoded, mlx.DtypeFloat32)
+		decoded = mlx.ClipScalar(decoded, 0.0, 1.0, true, true)
+		decoded = mlx.Transpose(decoded, 0, 3, 1, 2) // NHWC -> NCHW
+		return decoded
+	}
+
+	// Calculate tiling parameters (matching diffusers)
+	overlapSize := tileLatentSize - overlapLatent // stride in latent space
+
+	// Blend extent in pixel space (assumes 8x upscale, adjust if needed)
+	// For other scale factors, this could be made configurable
+	tileSampleSize := tileLatentSize * 8     // tile size in pixels after 8x upscale
+	blendExtent := overlapLatent * 8         // blend region in pixels
+	rowLimit := tileSampleSize - blendExtent // non-overlapping region per tile
+
+	// Phase 1: Decode all tiles and store in 2D grid
+	var rows [][]decodedTile
+
+	for i := int32(0); i < H; i += overlapSize {
+		var row []decodedTile
+		for j := int32(0); j < W; j += overlapSize {
+			// Extract tile (may be smaller at edges)
+			i2 := min(i+tileLatentSize, H)
+			j2 := min(j+tileLatentSize, W)
+
+			tile := mlx.Slice(latents, []int32{0, i, j, 0}, []int32{1, i2, j2, C})
+			decoded := decoder(tile)
+			decoded = mlx.AsType(decoded, mlx.DtypeFloat32)
+			mlx.Eval(decoded)
+
+			decodedShape := decoded.Shape()
+			tileH := decodedShape[1]
+			tileW := decodedShape[2]
+			tileData := decoded.Data()
+			decoded.Free()
+
+			row = append(row, decodedTile{data: tileData, height: tileH, width: tileW})
+		}
+		rows = append(rows, row)
+	}
+
+	// Phase 2: Blend adjacent tiles (modifies in place)
+	for i := range rows {
+		for j := range rows[i] {
+			tile := &rows[i][j]
+
+			// Blend with tile above
+			if i > 0 {
+				above := &rows[i-1][j]
+				blendV(above, tile, blendExtent)
+			}
+
+			// Blend with tile to the left
+			if j > 0 {
+				left := &rows[i][j-1]
+				blendH(left, tile, blendExtent)
+			}
+		}
+	}
+
+	// Phase 3: Calculate crop dimensions for each tile
+	colWidths := make([]int32, len(rows[0]))
+	for j := range rows[0] {
+		keepW := rowLimit
+		if int32(j+1)*overlapSize >= W {
+			keepW = rows[0][j].width
+		}
+		colWidths[j] = keepW
+	}
+
+	rowHeights := make([]int32, len(rows))
+	for i := range rows {
+		keepH := rowLimit
+		if int32(i+1)*overlapSize >= H {
+			keepH = rows[i][0].height
+		}
+		rowHeights[i] = keepH
+	}
+
+	// Calculate total dimensions
+	var totalW, totalH int32
+	for _, w := range colWidths {
+		totalW += w
+	}
+	for _, h := range rowHeights {
+		totalH += h
+	}
+
+	// Phase 4: Assemble final image by interleaving tiles row-by-row
+	finalData := make([]float32, totalH*totalW*3)
+
+	dstY := int32(0)
+	for i, row := range rows {
+		keepH := rowHeights[i]
+
+		for y := int32(0); y < keepH; y++ {
+			dstX := int32(0)
+			for j, tile := range row {
+				keepW := colWidths[j]
+
+				for x := int32(0); x < keepW; x++ {
+					for c := int32(0); c < 3; c++ {
+						srcIdx := (y*tile.width + x) * 3 + c
+						dstIdx := ((dstY + y) * totalW + (dstX + x)) * 3 + c
+						finalData[dstIdx] = tile.data[srcIdx]
+					}
+				}
+				dstX += keepW
+			}
+		}
+		dstY += keepH
+	}
+
+	// Create mlx array [1, H, W, 3] then transpose to NCHW [1, 3, H, W]
+	result := mlx.NewArray(finalData, []int32{1, totalH, totalW, 3})
+	result = mlx.Transpose(result, 0, 3, 1, 2)
+	result = mlx.ClipScalar(result, 0.0, 1.0, true, true)
+
+	return result
+}
+
+// blendV blends the bottom of 'above' tile into top of 'current' tile (vertical blend)
+// Matches diffusers blend_v formula
+func blendV(above, current *decodedTile, blendExtent int32) {
+	blend := min(blendExtent, min(above.height, current.height))
+	if blend <= 0 {
+		return
+	}
+
+	w := min(above.width, current.width)
+	for y := int32(0); y < blend; y++ {
+		alpha := float32(y) / float32(blend)
+		for x := int32(0); x < w; x++ {
+			for c := int32(0); c < 3; c++ {
+				aboveIdx := ((above.height - blend + y) * above.width + x) * 3 + c
+				currIdx := (y * current.width + x) * 3 + c
+				current.data[currIdx] = above.data[aboveIdx]*(1-alpha) + current.data[currIdx]*alpha
+			}
+		}
+	}
+}
+
+// blendH blends the right of 'left' tile into left of 'current' tile (horizontal blend)
+// Matches diffusers blend_h formula
+func blendH(left, current *decodedTile, blendExtent int32) {
+	blend := min(blendExtent, min(left.width, current.width))
+	if blend <= 0 {
+		return
+	}
+
+	h := min(left.height, current.height)
+	for y := int32(0); y < h; y++ {
+		for x := int32(0); x < blend; x++ {
+			alpha := float32(x) / float32(blend)
+			for c := int32(0); c < 3; c++ {
+				leftIdx := (y * left.width + (left.width - blend + x)) * 3 + c
+				currIdx := (y * current.width + x) * 3 + c
+				current.data[currIdx] = left.data[leftIdx]*(1-alpha) + current.data[currIdx]*alpha
+			}
+		}
+	}
+}