Release v0.4.0: selective MLX model downloads

3da88e3 verified 2 days ago

8.92 kB

from __future__ import annotations

import gc
import json
import logging
import math
from pathlib import Path

import mlx.core as mx
import numpy as np
from PIL import Image

from .ops import attention, layer_norm, rms_norm, rotate, swiglu
from .weights import Linear, WeightStore

log = logging.getLogger(__name__)


def split_tiles(length, tile=256, overlap=64, ratio=16):
    if tile <= overlap or tile % ratio or overlap % ratio:
        raise ValueError(
            "Tile size/overlap must be multiples of 16, with tile > overlap"
        )
    if length <= tile:
        return [0], [length], []
    n = math.ceil(length / tile)
    while tile * n - overlap * (n - 1) < length:
        n += 1
    overlaps = [overlap] * (n - 1)
    for i in range((tile * n - sum(overlaps) - length) // ratio):
        overlaps[i % (n - 1)] += ratio
    starts = [0]
    for ov in overlaps:
        starts.append(starts[-1] + tile - ov)
    return starts, [tile] * n, overlaps


def blend(a, b, extent, axis):
    extent = min(a.shape[axis], b.shape[axis], extent)
    if extent == 0:
        return b
    shape = [1] * b.ndim
    shape[axis] = extent
    weight = mx.arange(extent, dtype=b.dtype).reshape(shape) / extent
    aa, bb, rest = (
        [slice(None)] * b.ndim,
        [slice(None)] * b.ndim,
        [slice(None)] * b.ndim,
    )
    aa[axis], bb[axis], rest[axis] = (
        slice(-extent, None),
        slice(0, extent),
        slice(extent, None),
    )
    mixed = a[tuple(aa)] * (1 - weight) + b[tuple(bb)] * weight
    return mx.concatenate((mixed, b[tuple(rest)]), axis=axis)


def assemble_tiles(rows, y_overlap, x_overlap):
    assembled = []
    for i, row in enumerate(rows):
        pieces = []
        for j, tile in enumerate(row):
            if i:
                tile = blend(rows[i - 1][j], tile, y_overlap[i - 1], -2)
            if j:
                tile = blend(row[j - 1], tile, x_overlap[j - 1], -1)
            if i < len(rows) - 1:
                tile = tile[..., : -y_overlap[i], :]
            if j < len(row) - 1:
                tile = tile[..., : -x_overlap[j]]
            pieces.append(tile)
        assembled.append(mx.concatenate(pieces, axis=-1))
    return mx.concatenate(assembled, axis=-2)


class DecoderBlock:
    def __init__(self, store, index, dtype):
        p = f"decoder.transformer_blocks.{index}"
        self.n1, self.n2 = (store.get(p + f".norm{i}.weight") for i in (1, 2))
        self.s1, self.s2 = (store.get(p + f".scale{i}") for i in (1, 2))
        self.qkv = Linear.load(store, p + ".attn.to_qkv")
        self.out = Linear.load(store, p + ".attn.to_out")
        self.w1 = Linear.load(store, p + ".ff.w1")
        self.w2 = Linear.load(store, p + ".ff.w2")

    def __call__(self, x, rope):
        b, n, _ = x.shape
        h = rms_norm(x, self.n1)
        # VAE QKV is head-interleaved too, but keep its native layout here.
        qkv = self.qkv(h).reshape(b, n, 16, 3, 64)
        q, k, v = (qkv[:, :, :, i] for i in range(3))
        q, k = rms_norm(q), rms_norm(k)
        q, k = rotate(q, *rope), rotate(k, *rope)
        a = self.out(attention(q, k, v).reshape(b, n, -1))
        x = x + a * self.s1
        h = rms_norm(x, self.n2)
        return x + self.w2(swiglu(self.w1(h))) * self.s2


class VAE:
    def __init__(
        self,
        path,
        *,
        tile_batch=1,
        dtype="bfloat16",
        tile_size=256,
        compile_blocks=False,
    ):
        self.path = Path(path)
        self.config = json.loads((self.path / "config.json").read_text())
        self.dtype = {"bfloat16": mx.bfloat16, "float32": mx.float32}[dtype]
        self.store = WeightStore(self.path, cast=self.dtype)
        self.tile_batch, self.tile_size = tile_batch, tile_size
        self.blocks = None
        self.compile_blocks = compile_blocks
        self.runners = []

    def _load(self):
        if self.blocks is not None:
            return
        s = self.store
        self.post = Linear(
            s.get("post_quant_conv.weight").reshape(24, 24),
            s.get("post_quant_conv.bias"),
        )
        self.embed = Linear.load(s, "decoder.x_embedder")
        self.register = s.get("decoder.register_tokens")
        self.norm_w, self.norm_b = (
            s.get("decoder.norm_out.weight"),
            s.get("decoder.norm_out.bias"),
        )
        self.out = Linear.load(s, "decoder.proj_out")
        self.blocks = [DecoderBlock(s, i, self.dtype) for i in range(24)]
        if self.compile_blocks:
            for block in self.blocks:

                def forward(x, rope, block=block):
                    return block(x, rope)

                self.runners.append(mx.compile(forward))

    def decode_tiles(self, raw):
        self._load()
        b, c, h, w = raw.shape
        # Native image decoder is BF16, including its pointwise convolution.
        x = raw.transpose(0, 2, 3, 1).astype(self.dtype)
        x = self.embed(self.post(x)).reshape(b, h * w, 1024)
        x = mx.concatenate(
            (
                x,
                mx.broadcast_to(self.register, (b, 4, 1024)),
                mx.zeros((b, 1, 1024), self.dtype),
            ),
            axis=1,
        )
        # Coordinates start in latent dtype (BF16), as in released ViT2D.
        yy = 2 * ((mx.arange(h, dtype=self.dtype) + mx.array(0.5, self.dtype)) / h) - 1
        xx = 2 * ((mx.arange(w, dtype=self.dtype) + mx.array(0.5, self.dtype)) / w) - 1
        yy, xx = mx.meshgrid(yy, xx, indexing="ij")
        positions = mx.stack((yy, xx), axis=-1).reshape(-1, 2)
        positions = mx.concatenate((positions, mx.zeros((5, 2), self.dtype)), axis=0)
        inv = 1 / (100 ** mx.arange(0, 1, 0.125, dtype=mx.float32))
        scaled = (2 * math.pi * positions.astype(mx.float32)).astype(self.dtype)
        angles = (scaled.astype(mx.float32)[..., None] * inv).reshape(h * w + 5, 16)
        angles = mx.concatenate((angles, angles), axis=-1)
        rope = (
            mx.cos(angles).astype(self.dtype)[None, :, None],
            mx.sin(angles).astype(self.dtype)[None, :, None],
        )
        for block in self.runners if self.compile_blocks else self.blocks:
            x = block(x, rope)
            if not self.compile_blocks:
                mx.eval(x)
        x = layer_norm(x, self.norm_w, self.norm_b)
        x = self.out(x)[:, : h * w]
        # Output channels are [C, patchH, patchW], matching the VAE checkpoint.
        return (
            x.reshape(b, h, w, 3, 16, 16)
            .transpose(0, 3, 1, 4, 2, 5)
            .reshape(b, 3, h * 16, w * 16)
        )

    def decode(self, latent, *, on_batch=None):
        c = self.config
        mean = mx.array(c["latents_mean"], mx.float32)[None, :, None, None]
        std = mx.array(c["latents_std"], mx.float32)[None, :, None, None]
        raw = latent.astype(mx.float32) * std + mean
        h, w = raw.shape[-2:]
        ys, ylen, yo = split_tiles(h * 16, self.tile_size)
        xs, xlen, xo = split_tiles(w * 16, self.tile_size)
        tasks = [
            (i, j, raw[..., y // 16 : (y + yl) // 16, x // 16 : (x + xl) // 16])
            for i, (y, yl) in enumerate(zip(ys, ylen, strict=False))
            for j, (x, xl) in enumerate(zip(xs, xlen, strict=False))
        ]
        rows = [[None for _ in xs] for _ in ys]
        for start in range(0, len(tasks), self.tile_batch):
            group = tasks[start : start + self.tile_batch]
            tiles = mx.concatenate([t for _, _, t in group], axis=0)
            decoded = self.decode_tiles(tiles)
            mx.eval(decoded)
            for k, (i, j, _) in enumerate(group):
                rows[i][j] = decoded[k : k + 1]
            if on_batch:
                on_batch(start + len(group), len(tasks))
        # Decode tile overlap arithmetic stays in the decoder dtype.
        out = assemble_tiles(rows, yo, xo)
        # torchvision's inverse Normalize subtracts its inverse mean, then divides.
        inv_std = mx.array([1 / 0.229, 1 / 0.224, 1 / 0.225], self.dtype)[
            None, :, None, None
        ]
        inv_mean = mx.array(
            [-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225], self.dtype
        )[None, :, None, None]
        out = mx.clip((out - inv_mean) / inv_std, 0, 1).astype(mx.float32)
        mx.eval(out)
        return out

    def clear(self):
        self.runners.clear()
        self.blocks = None
        for name in ("post", "embed", "register", "norm_w", "norm_b", "out"):
            if hasattr(self, name):
                delattr(self, name)
        self.store.clear()
        gc.collect()
        mx.clear_cache()


def to_image(pixels, width, height):
    array = np.asarray(
        mx.clip(pixels[0].transpose(1, 2, 0) * 255, 0, 255).astype(mx.uint8)
    )
    image = Image.fromarray(array)
    if image.size != (width, height):
        image = image.resize((width, height), Image.Resampling.LANCZOS)
    return image

	from __future__ import annotations

	import gc
	import json
	import logging
	import math
	from pathlib import Path

	import mlx.core as mx
	import numpy as np
	from PIL import Image

	from .ops import attention, layer_norm, rms_norm, rotate, swiglu
	from .weights import Linear, WeightStore

	log = logging.getLogger(__name__)


	def split_tiles(length, tile=256, overlap=64, ratio=16):
	if tile <= overlap or tile % ratio or overlap % ratio:
	raise ValueError(
	"Tile size/overlap must be multiples of 16, with tile > overlap"
	)
	if length <= tile:
	return [0], [length], []
	n = math.ceil(length / tile)
	while tile * n - overlap * (n - 1) < length:
	n += 1
	overlaps = [overlap] * (n - 1)
	for i in range((tile * n - sum(overlaps) - length) // ratio):
	overlaps[i % (n - 1)] += ratio
	starts = [0]
	for ov in overlaps:
	starts.append(starts[-1] + tile - ov)
	return starts, [tile] * n, overlaps


	def blend(a, b, extent, axis):
	extent = min(a.shape[axis], b.shape[axis], extent)
	if extent == 0:
	return b
	shape = [1] * b.ndim
	shape[axis] = extent
	weight = mx.arange(extent, dtype=b.dtype).reshape(shape) / extent
	aa, bb, rest = (
	[slice(None)] * b.ndim,
	[slice(None)] * b.ndim,
	[slice(None)] * b.ndim,
	)
	aa[axis], bb[axis], rest[axis] = (
	slice(-extent, None),
	slice(0, extent),
	slice(extent, None),
	)
	mixed = a[tuple(aa)] * (1 - weight) + b[tuple(bb)] * weight
	return mx.concatenate((mixed, b[tuple(rest)]), axis=axis)


	def assemble_tiles(rows, y_overlap, x_overlap):
	assembled = []
	for i, row in enumerate(rows):
	pieces = []
	for j, tile in enumerate(row):
	if i:
	tile = blend(rows[i - 1][j], tile, y_overlap[i - 1], -2)
	if j:
	tile = blend(row[j - 1], tile, x_overlap[j - 1], -1)
	if i < len(rows) - 1:
	tile = tile[..., : -y_overlap[i], :]
	if j < len(row) - 1:
	tile = tile[..., : -x_overlap[j]]
	pieces.append(tile)
	assembled.append(mx.concatenate(pieces, axis=-1))
	return mx.concatenate(assembled, axis=-2)


	class DecoderBlock:
	def __init__(self, store, index, dtype):
	p = f"decoder.transformer_blocks.{index}"
	self.n1, self.n2 = (store.get(p + f".norm{i}.weight") for i in (1, 2))
	self.s1, self.s2 = (store.get(p + f".scale{i}") for i in (1, 2))
	self.qkv = Linear.load(store, p + ".attn.to_qkv")
	self.out = Linear.load(store, p + ".attn.to_out")
	self.w1 = Linear.load(store, p + ".ff.w1")
	self.w2 = Linear.load(store, p + ".ff.w2")

	def __call__(self, x, rope):
	b, n, _ = x.shape
	h = rms_norm(x, self.n1)
	# VAE QKV is head-interleaved too, but keep its native layout here.
	qkv = self.qkv(h).reshape(b, n, 16, 3, 64)
	q, k, v = (qkv[:, :, :, i] for i in range(3))
	q, k = rms_norm(q), rms_norm(k)
	q, k = rotate(q, rope), rotate(k, rope)
	a = self.out(attention(q, k, v).reshape(b, n, -1))
	x = x + a * self.s1
	h = rms_norm(x, self.n2)
	return x + self.w2(swiglu(self.w1(h))) * self.s2


	class VAE:
	def __init__(
	self,
	path,
	*,
	tile_batch=1,
	dtype="bfloat16",
	tile_size=256,
	compile_blocks=False,
	):
	self.path = Path(path)
	self.config = json.loads((self.path / "config.json").read_text())
	self.dtype = {"bfloat16": mx.bfloat16, "float32": mx.float32}[dtype]
	self.store = WeightStore(self.path, cast=self.dtype)
	self.tile_batch, self.tile_size = tile_batch, tile_size
	self.blocks = None
	self.compile_blocks = compile_blocks
	self.runners = []

	def _load(self):
	if self.blocks is not None:
	return
	s = self.store
	self.post = Linear(
	s.get("post_quant_conv.weight").reshape(24, 24),
	s.get("post_quant_conv.bias"),
	)
	self.embed = Linear.load(s, "decoder.x_embedder")
	self.register = s.get("decoder.register_tokens")
	self.norm_w, self.norm_b = (
	s.get("decoder.norm_out.weight"),
	s.get("decoder.norm_out.bias"),
	)
	self.out = Linear.load(s, "decoder.proj_out")
	self.blocks = [DecoderBlock(s, i, self.dtype) for i in range(24)]
	if self.compile_blocks:
	for block in self.blocks:

	def forward(x, rope, block=block):
	return block(x, rope)

	self.runners.append(mx.compile(forward))

	def decode_tiles(self, raw):
	self._load()
	b, c, h, w = raw.shape
	# Native image decoder is BF16, including its pointwise convolution.
	x = raw.transpose(0, 2, 3, 1).astype(self.dtype)
	x = self.embed(self.post(x)).reshape(b, h * w, 1024)
	x = mx.concatenate(
	(
	x,
	mx.broadcast_to(self.register, (b, 4, 1024)),
	mx.zeros((b, 1, 1024), self.dtype),
	),
	axis=1,
	)
	# Coordinates start in latent dtype (BF16), as in released ViT2D.
	yy = 2 * ((mx.arange(h, dtype=self.dtype) + mx.array(0.5, self.dtype)) / h) - 1
	xx = 2 * ((mx.arange(w, dtype=self.dtype) + mx.array(0.5, self.dtype)) / w) - 1
	yy, xx = mx.meshgrid(yy, xx, indexing="ij")
	positions = mx.stack((yy, xx), axis=-1).reshape(-1, 2)
	positions = mx.concatenate((positions, mx.zeros((5, 2), self.dtype)), axis=0)
	inv = 1 / (100 ** mx.arange(0, 1, 0.125, dtype=mx.float32))
	scaled = (2 * math.pi * positions.astype(mx.float32)).astype(self.dtype)
	angles = (scaled.astype(mx.float32)[..., None] * inv).reshape(h * w + 5, 16)
	angles = mx.concatenate((angles, angles), axis=-1)
	rope = (
	mx.cos(angles).astype(self.dtype)[None, :, None],
	mx.sin(angles).astype(self.dtype)[None, :, None],
	)
	for block in self.runners if self.compile_blocks else self.blocks:
	x = block(x, rope)
	if not self.compile_blocks:
	mx.eval(x)
	x = layer_norm(x, self.norm_w, self.norm_b)
	x = self.out(x)[:, : h * w]
	# Output channels are [C, patchH, patchW], matching the VAE checkpoint.
	return (
	x.reshape(b, h, w, 3, 16, 16)
	.transpose(0, 3, 1, 4, 2, 5)
	.reshape(b, 3, h * 16, w * 16)
	)

	def decode(self, latent, *, on_batch=None):
	c = self.config
	mean = mx.array(c["latents_mean"], mx.float32)[None, :, None, None]
	std = mx.array(c["latents_std"], mx.float32)[None, :, None, None]
	raw = latent.astype(mx.float32) * std + mean
	h, w = raw.shape[-2:]
	ys, ylen, yo = split_tiles(h * 16, self.tile_size)
	xs, xlen, xo = split_tiles(w * 16, self.tile_size)
	tasks = [
	(i, j, raw[..., y // 16 : (y + yl) // 16, x // 16 : (x + xl) // 16])
	for i, (y, yl) in enumerate(zip(ys, ylen, strict=False))
	for j, (x, xl) in enumerate(zip(xs, xlen, strict=False))
	]
	rows = [[None for _ in xs] for _ in ys]
	for start in range(0, len(tasks), self.tile_batch):
	group = tasks[start : start + self.tile_batch]
	tiles = mx.concatenate([t for _, _, t in group], axis=0)
	decoded = self.decode_tiles(tiles)
	mx.eval(decoded)
	for k, (i, j, _) in enumerate(group):
	rows[i][j] = decoded[k : k + 1]
	if on_batch:
	on_batch(start + len(group), len(tasks))
	# Decode tile overlap arithmetic stays in the decoder dtype.
	out = assemble_tiles(rows, yo, xo)
	# torchvision's inverse Normalize subtracts its inverse mean, then divides.
	inv_std = mx.array([1 / 0.229, 1 / 0.224, 1 / 0.225], self.dtype)[
	None, :, None, None
	]
	inv_mean = mx.array(
	[-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225], self.dtype
	)[None, :, None, None]
	out = mx.clip((out - inv_mean) / inv_std, 0, 1).astype(mx.float32)
	mx.eval(out)
	return out

	def clear(self):
	self.runners.clear()
	self.blocks = None
	for name in ("post", "embed", "register", "norm_w", "norm_b", "out"):
	if hasattr(self, name):
	delattr(self, name)
	self.store.clear()
	gc.collect()
	mx.clear_cache()


	def to_image(pixels, width, height):
	array = np.asarray(
	mx.clip(pixels[0].transpose(1, 2, 0) * 255, 0, 255).astype(mx.uint8)
	)
	image = Image.fromarray(array)
	if image.size != (width, height):
	image = image.resize((width, height), Image.Resampling.LANCZOS)
	return image