File size: 3,376 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 | """Optional fixed-tile Core ML decoder with the native MLX tile assembly.
Core ML uses FP16 internally. This is an explicitly approximate alternative to
the native BF16 decoder; use the saved-latent image comparisons before choosing.
"""
from __future__ import annotations
from pathlib import Path
import mlx.core as mx
import numpy as np
from .vae import VAE
class CoreMLVAE(VAE):
@staticmethod
def validate_metadata(meta, source_fingerprint, tile_batch):
if meta.get("h3.decoder.format") != "h3-coreml-tile-v1":
raise ValueError("This ML Package lacks H3 decoder provenance metadata")
if meta.get("h3.decoder.source") != source_fingerprint:
raise ValueError(
"Core ML decoder does not match this source VAE checkpoint"
)
fixed_batch = int(meta["h3.decoder.batch"])
if tile_batch != fixed_batch:
raise ValueError(f"ML Package was exported for --vae-batch {fixed_batch}")
return fixed_batch
def __init__(
self,
path,
*,
package,
compute_units="CPU_AND_GPU",
tile_batch=8,
dtype="bfloat16",
compile_blocks=False,
):
if dtype != "bfloat16":
raise ValueError(
"Core ML tile assembly currently requires the native BF16 policy"
)
import coremltools as ct
super().__init__(path, tile_batch=tile_batch, dtype=dtype)
self.package = Path(package)
self.model = ct.models.MLModel(
str(self.package),
compute_units=getattr(ct.ComputeUnit, compute_units),
optimization_hints={
"specializationStrategy": ct.SpecializationStrategy.FastPrediction
},
)
meta = self.model.user_defined_metadata
self.fixed_batch = self.validate_metadata(
meta, self.store.fingerprint(), tile_batch
)
def decode_tiles(self, raw):
if raw.shape[1:] != (24, 16, 16) or raw.shape[0] > self.fixed_batch:
raise ValueError(
"Core ML export requires native 256px tiles and its fixed batch size"
)
count = raw.shape[0]
# BF16 source values are exactly representable by FP16 in the ordinary
# decoder range. Preserve the native input rounding before conversion.
x = raw.astype(mx.bfloat16).astype(mx.float16)
if count < self.fixed_batch:
x = mx.pad(x, ((0, self.fixed_batch - count), (0, 0), (0, 0), (0, 0)))
pixels = self.model.predict({"raw": np.ascontiguousarray(np.asarray(x))})[
"pixels"
]
if not np.isfinite(pixels).all():
raise FloatingPointError("Core ML decoder produced non-finite pixels")
return mx.array(np.ascontiguousarray(pixels[:count])).astype(mx.bfloat16)
def clear(self):
self.model = None
super().clear()
def source_metadata(path, batch):
"""Metadata shared with the exporter, without importing the converter."""
from .weights import WeightStore
return {
"h3.decoder.format": "h3-coreml-tile-v1",
"h3.decoder.source": WeightStore(path).fingerprint(),
"h3.decoder.batch": str(batch),
"h3.decoder.precision": "Source BF16 weights and inputs, FP16 internal arithmetic, BF16 tile assembly",
}
|