File size: 6,351 Bytes
3da88e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 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 169 | """Image-reference CNN VAE encoder (single-frame causal convolutions in MLX).
For a single frame, causal temporal zero padding selects the last temporal
kernel plane. This is the exact 2D reduction of the released 3D CNN, including
its spatial reflection padding, group norms and sampled posterior.
"""
from __future__ import annotations
import json
import logging
from pathlib import Path
import mlx.core as mx
import numpy as np
from .layout import patchify
from .ops import silu
from .vae import assemble_tiles, split_tiles
from .vision import bound_reference
from .weights import WeightStore
log = logging.getLogger(__name__)
def reflect_pad(x, left, right, top, bottom):
# NHWC; explicit slices avoid NumPy/CPU roundtrips and support Metal.
if top or bottom:
parts = []
if top:
parts.append(x[:, 1 : top + 1][:, ::-1])
parts.append(x)
if bottom:
parts.append(x[:, -bottom - 1 : -1][:, ::-1])
x = mx.concatenate(parts, axis=1)
if left or right:
parts = []
if left:
parts.append(x[:, :, 1 : left + 1][:, :, ::-1])
parts.append(x)
if right:
parts.append(x[:, :, -right - 1 : -1][:, :, ::-1])
x = mx.concatenate(parts, axis=2)
return x
class ReferenceEncoder:
def __init__(self, path):
self.store = WeightStore(path, cache_gb=2)
self.config = json.loads((Path(path) / "config.json").read_text())
def conv(self, x, prefix, *, stride=1, pad=1, down=False):
weight = self.store.get(prefix + ".weight")
if weight.ndim == 5:
weight = weight[:, :, -1]
weight = weight.transpose(0, 2, 3, 1)
if down:
x = reflect_pad(x, 0, 1, 0, 1)
elif pad:
x = reflect_pad(x, pad, pad, pad, pad)
out = mx.conv2d(x, weight, stride=stride)
return out + self.store.get(prefix + ".bias")
def norm(self, x, prefix):
b, h, w, c = x.shape
# GroupNorm reduces spatial dimensions and channels within each group.
f = x.astype(mx.float32).reshape(b, h, w, 32, c // 32)
mean = mx.mean(f, axis=(1, 2, 4), keepdims=True)
var = mx.mean((f - mean) ** 2, axis=(1, 2, 4), keepdims=True)
f = ((f - mean) * mx.rsqrt(var + 1e-6)).reshape(x.shape)
return f * self.store.get(prefix + ".weight") + self.store.get(prefix + ".bias")
def resblock(self, x, prefix):
h = self.conv(silu(self.norm(x, prefix + ".norm1")), prefix + ".conv1")
h = self.conv(silu(self.norm(h, prefix + ".norm2")), prefix + ".conv2")
if prefix + ".nin_shortcut.weight" in self.store:
x = self.conv(x, prefix + ".nin_shortcut", pad=0)
return h + x
def encode_tile(self, x):
x = self.conv(x, "encoder.conv_in")
for i in range(6):
for j in range(2):
x = self.resblock(x, f"encoder.down.{i}.block.{j}")
mx.eval(x)
if i < 4:
x = self.conv(
x, f"encoder.down.{i}.downsample.conv", stride=2, pad=0, down=True
)
mx.eval(x)
x = self.conv(silu(self.norm(x, "encoder.norm_out")), "encoder.conv_out")
x = self.conv(x, "quant_conv", pad=0)
mx.eval(x)
return x.transpose(0, 3, 1, 2)
def encode(self, image):
# Native conditioning CNN runs FP32, independently from decoder BF16.
image = bound_reference(image, 1024, 2048)
w, h = image.size
w, h = w // 16 * 16, h // 16 * 16
array = np.asarray(image)[:h, :w].astype(np.float32) / 255
x = mx.array(
(array - np.array([0.485, 0.456, 0.406], np.float32))
/ np.array([0.229, 0.224, 0.225], np.float32)
)[None]
ys, yl, yo = split_tiles(h)
xs, xl, xo = split_tiles(w)
rows = []
for y, hh in zip(ys, yl, strict=False):
row = []
for xx, ww in zip(xs, xl, strict=False):
row.append(self.encode_tile(x[:, y : y + hh, xx : xx + ww]))
rows.append(row)
moments = assemble_tiles(rows, [n // 16 for n in yo], [n // 16 for n in xo])
mx.eval(moments)
return moments
def encode_references(path, images, *, generator, seed):
encoder = ReferenceEncoder(path)
bounded = [bound_reference(im, 1024, 2048) for im in images]
grids = [(im.height // 32 * 2, im.width // 32 * 2) for im in bounded]
# Consume every anchor draw before any posterior draw, as in H3 native.
if generator is not None:
import torch
anchor_noise = [
mx.array(
torch.randn(
1, 24, h, w, generator=generator, dtype=torch.float32
).numpy()
)
for h, w in grids
]
else:
anchor_noise = [
mx.random.normal((1, 24, h, w), key=mx.random.key(seed + 1 + i))
for i, (h, w) in enumerate(grids)
]
anchors = []
cfg = encoder.config
mean = mx.array(cfg["latents_mean"], mx.float32)[None, :, None, None]
std = mx.array(cfg["latents_std"], mx.float32)[None, :, None, None]
for i, (image, (h, w), noise) in enumerate(
zip(bounded, grids, anchor_noise, strict=False)
):
moments = encoder.encode(image)
mu, logvar = mx.split(moments, 2, axis=1)
logvar = mx.clip(logvar, -30, 20)
if generator is not None:
# Native posterior uses CPU torch noise then moves to execution device.
draw = mx.array(
torch.randn(
tuple(mu.shape), generator=generator, dtype=torch.float32
).numpy()
)
else:
draw = mx.random.normal(mu.shape, key=mx.random.key(seed + 100 + i))
z = mu + mx.exp(0.5 * logvar) * draw
# The upstream encode_images contract rounds its sampled latent to FP16.
z = z.astype(mx.float16).astype(mx.float32)[:, :, :h, :w]
clean = patchify((z - mean) / std)
t = mx.array(0.999, mx.float32)
anchors.append(t * clean + (1 - t) * patchify(noise))
mx.eval(anchors[-1])
log.info("Encoded edit reference %d/%d", i + 1, len(images))
encoder.store.clear()
return tuple(grids), mx.concatenate(anchors)
|