File size: 4,530 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 | """Compare native, compiled and Core ML decoders on the exact same final latent."""
from __future__ import annotations
import argparse
import gc
import hashlib
import json
from pathlib import Path
import mlx.core as mx
from h3_mlx.benchmark import measure
from h3_mlx.compare import image_metrics, make_comparison
from h3_mlx.coreml_vae import CoreMLVAE
from h3_mlx.runtime import configure_memory, environment
from h3_mlx.vae import VAE, to_image
def main(a):
configure_memory(wired_gb=5)
root = Path(a.output)
root.mkdir(parents=True, exist_ok=True)
data = mx.load(a.trace)
latent = data["latent"]
mx.eval(latent)
del data
source = json.loads(Path(a.record).read_text())
width, height = source["size"]
results = []
images = []
labels = []
reference = None
variants = [
("native-b1", VAE, {"tile_batch": 1}),
("native-b4", VAE, {"tile_batch": 4}),
("compiled-b4", VAE, {"tile_batch": 4, "compile_blocks": True}),
("native-b8", VAE, {"tile_batch": 8}),
("compiled-b8", VAE, {"tile_batch": 8, "compile_blocks": True}),
]
if a.include_fp32:
variants.insert(
0,
(
"fp32-b4-compiled",
VAE,
{"tile_batch": 4, "compile_blocks": True, "dtype": "float32"},
),
)
if a.package:
variants.extend(
(
"coreml-" + unit,
CoreMLVAE,
{
"tile_batch": a.coreml_batch,
"package": a.package,
"compute_units": unit,
},
)
for unit in a.units
)
for name, cls, options in variants:
model = cls(a.vae, **options)
timer = measure(
lambda model=model: model.decode(latent), warmup=1, repeats=a.repeats
)
pixels = model.decode(latent)
image = to_image(pixels, width, height)
if reference is None:
reference = image
path = root / f"{name}.png"
image.save(path)
record = {
k: source[k]
for k in (
"prompt",
"seed",
"size",
"canvas",
"rng",
"reference_images",
"reference_sha256",
"sigmas",
"steps",
)
if k in source
}
record.update(
vae=options,
vae_backend="coreml" if cls is CoreMLVAE else "mlx",
environment=environment(),
validation_scope="Same final latent; decoder only",
timing_scope="Warm decoder only; model loading excluded",
timings={"decode_s": timer["median_ms"] / 1000},
source_trace=a.trace,
source_trace_sha256=hashlib.sha256(Path(a.trace).read_bytes()).hexdigest(),
output=str(path),
png_sha256=hashlib.sha256(path.read_bytes()).hexdigest(),
)
path.with_suffix(".json").write_text(json.dumps(record, indent=2) + "\n")
row = {"variant": name, **timer, **image_metrics(reference, image)}
results.append(row)
images.append(path)
labels.append(name)
print(json.dumps(row), flush=True)
(root / "trials.json").write_text(
json.dumps(
{
"source_trace": a.trace,
"scope": "Same final latent, decoder-only hot latency",
"trials": results,
},
indent=2,
)
+ "\n"
)
model.clear()
del model
gc.collect()
mx.clear_cache()
make_comparison(images[0], images[1:], root / "comparison", labels=labels)
if __name__ == "__main__":
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--vae", required=True)
p.add_argument("--trace", required=True)
p.add_argument("--record", required=True)
p.add_argument("--package")
p.add_argument(
"--include-fp32",
action="store_true",
help="Use the original FP32 decoder as the same-latent reference",
)
p.add_argument("--coreml-batch", type=int, default=8)
p.add_argument("--units", nargs="+", default=["CPU_AND_GPU"])
p.add_argument("--repeats", type=int, default=3)
p.add_argument("--output", default="comparisons/optimization-v2/decoder-product")
main(p.parse_args())
|