| """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] |
| |
| |
| 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", |
| } |
|
|
