Release v0.4.0: selective MLX model downloads

3da88e3 verified 2 days ago

6.35 kB

"""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)

	"""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)