Blame for h3_mlx/runtime.py

from __future__ import annotations

import gc
import hashlib
import importlib.metadata
import json
import logging
import os
import platform
import time
from datetime import datetime, timezone
from pathlib import Path

import mlx.core as mx

log = logging.getLogger(__name__)


def configure_memory(*, cache_gb=1, wired_gb=None):
    if not mx.metal.is_available():
        raise RuntimeError("MLX Metal acceleration requires an Apple silicon Mac")
    mx.set_default_device(mx.gpu)
    mx.set_cache_limit(int(cache_gb * 1024**3))
    if wired_gb is not None:
        requested = int(wired_gb * 1024**3)
        limit = mx.device_info()["max_recommended_working_set_size"]
        if not 0 <= requested <= limit:
            raise ValueError("wired_gb must be within Metal's recommended working set")
        mx.set_wired_limit(requested)


def environment():
    versions = {}
    for package in (
        "mlx",
        "mlx-metal",
        "numpy",
        "safetensors",
        "transformers",
        "torch",
        "Pillow",
        "scikit-image",
    ):
        try:
            versions[package] = importlib.metadata.version(package)
        except importlib.metadata.PackageNotFoundError:
            pass
    return {
        "platform": platform.platform(),
        "machine": platform.machine(),
        "metal": mx.device_info(),
        "versions": versions,
        "precision_environment": {
            k: os.environ.get(k) for k in ("MLX_ENABLE_TF32", "MLX_CONV_WINOGRAD")
        },
    }


def memory_stats():
    result = {
        "active_gb": mx.get_active_memory() / 1024**3,
        "peak_gb": mx.get_peak_memory() / 1024**3,
        "cache_gb": mx.get_cache_memory() / 1024**3,
    }
    if platform.system() == "Darwin":
        from .memory import footprint

        physical = footprint()
        result.update(
            process_physical_gb=physical["physical_bytes"] / 1024**3,
            process_lifetime_peak_physical_gb=physical["lifetime_peak_physical_bytes"]
            / 1024**3,
        )
    return result


def initial_noise(seed, h, w, *, rng="torch"):
    if rng == "torch":
        try:
            import torch
        except ImportError as error:
            raise RuntimeError(
                "Install the audit extra for native CPU noise, or use --rng mlx"
            ) from error
        generator = torch.Generator(device="cpu").manual_seed(seed)
        return mx.array(
            torch.randn(1, 24, h, w, generator=generator, dtype=torch.float32).numpy()
        ), generator
    if rng == "mlx":
        return mx.random.normal((1, 24, h, w), key=mx.random.key(seed)), None
    raise ValueError("rng must be torch or mlx")


class Pipeline:
    def __init__(
        self,
        model_path,
        encoder_path,
        *,
        cache_dir="cache",
        dit_cache_gb=38,
        bits=None,
        encoder_bits=None,
        group_size=64,
        mlp_chunk=0,
        query_chunk=0,
        kernel="native",
        threads=256,
        vae_batch=1,
        vae_dtype="bfloat16",
        keep_resident=True,
        quantized_path=None,
        compile_blocks=False,
        eval_every=1,
        keep_refiners=False,
        wired_gb=None,
        weight_storage="bf16",
        vae_compile=False,
        vae_backend="mlx",
        coreml_package=None,
        coreml_compute_units="CPU_AND_GPU",
        async_blocks=False,
        cache_refined_text=True,
        prefetch_blocks=0,
        allow_nonfinite=False,
        tune_matmul=False,
    ):
        self.model_path, self.encoder_path = Path(model_path), Path(encoder_path)
        self.cache_dir = Path(cache_dir)
        self.dit_options = dict(
            cache_dir=cache_dir,
            cache_gb=dit_cache_gb,
            bits=bits,
            group_size=group_size,
            mlp_chunk=mlp_chunk,
            query_chunk=query_chunk,
            kernel=kernel,
            threads=threads,
            quantized_path=quantized_path,
            compile_blocks=compile_blocks,
            eval_every=eval_every,
            keep_refiners=keep_refiners,
            weight_storage=weight_storage,
            async_blocks=async_blocks,
            cache_refined_text=cache_refined_text,
            prefetch_blocks=prefetch_blocks,
            tune_matmul=tune_matmul,
        )
        self.encoder_options = dict(
            cache_dir=cache_dir,
            bits=encoder_bits,
            group_size=group_size,
            query_chunk=query_chunk,
        )
        if vae_backend not in ("mlx", "coreml"):
            raise ValueError("vae_backend must be mlx or coreml")
        if vae_backend == "coreml" and (
            coreml_package is None or not Path(coreml_package).is_dir()
        ):
            raise ValueError("Core ML decoding requires an exported --coreml-package")
        self.vae_backend = vae_backend
        self.vae_options = dict(
            tile_batch=vae_batch, dtype=vae_dtype, compile_blocks=vae_compile
        )
        if vae_backend == "coreml":
            if vae_dtype != "bfloat16":
                raise ValueError("Core ML tile assembly requires --vae-dtype bfloat16")
            import coremltools as ct

            from .coreml_vae import CoreMLVAE
            from .weights import WeightStore

            metadata = ct.models.MLModel(
                str(coreml_package), skip_model_load=True
            ).user_defined_metadata
            CoreMLVAE.validate_metadata(
                metadata, WeightStore(self.model_path / "vae").fingerprint(), vae_batch
            )
            self.vae_options["compile_blocks"] = False
            self.vae_options.update(
                package=str(coreml_package), compute_units=coreml_compute_units
            )
        self.keep_resident = keep_resident
        self.wired_gb = wired_gb
        self.allow_nonfinite = allow_nonfinite
        self.model = None
        self.vae = None

    def clear(self):
        for component in (self.model, self.vae):
            if component is not None:
                component.clear()
        self.model = self.vae = None
        gc.collect()
        mx.clear_cache()

    def generate(
        self,
        prompt,
        *,
        width=2048,
        height=2048,
        seed=731,
        min_area=4194304,
        rng="torch",
        output="outputs/image.png",
        images=(),
        steps=None,
        save_trace=False,
    ):
        from .config import canvas_size, load_sigmas
        from .dit import DiT
        from .encoder import TextEncoder
        from .layout import patchify, unpatchify
        from .vae import VAE, to_image

        configure_memory(wired_gb=self.wired_gb)
        mx.reset_peak_memory()
        output = Path(output)
        output.parent.mkdir(parents=True, exist_ok=True)
        cw, ch = canvas_size(width, height, min_area)
        h, w = ch // 16, cw // 16
        sigmas = load_sigmas(self.model_path)
        selected = list(range(6)) if steps is None else list(steps)
        if (
            not selected
            or selected[0] != 0
            or sorted(set(selected)) != selected
            or selected[-1] > 5
        ):
            raise ValueError(
                "Experimental steps must be ascending unique indices in [0,5], starting at 0"
            )
        record = {
            "started_utc": datetime.now(timezone.utc).isoformat(),
            "prompt": prompt,
            "seed": seed,
            "size": [width, height],
            "canvas": [cw, ch],
            "rng": rng,
            "sigmas": sigmas,
            "steps": selected,
            "dit": self.dit_options,
            "encoder": self.encoder_options,
            "vae": self.vae_options,
            "vae_backend": self.vae_backend,
            "environment": environment(),
            "wired_gb": self.wired_gb,
            "timings": {},
            "nonfinite_steps": [],
            "allow_nonfinite_diagnostic": self.allow_nonfinite,
            "reference_images": [str(i) for i in images],
            "reference_sha256": [
                hashlib.sha256(Path(i).read_bytes()).hexdigest() for i in images
            ],
        }
        start = time.perf_counter()
        noise, generator = initial_noise(seed, h, w, rng=rng)
        rows = patchify(noise)
        mx.eval(rows)
        image_objects = []
        if images:
            from PIL import Image

            image_objects = [Image.open(p).convert("RGB") for p in images]
        enc_start = time.perf_counter()
        encoder = TextEncoder(self.encoder_path, **self.encoder_options)
        hidden, tags = encoder.encode(prompt, images=image_objects)
        record["conditioning_cache_hit"] = encoder.cache_hit
        record["timings"]["encode_s"] = time.perf_counter() - enc_start
        encoder.store.clear()
        del encoder
        gc.collect()
        mx.clear_cache()
        references, anchors = (), None
        if image_objects:
            from .reference import encode_references

            references, anchors = encode_references(
                self.model_path / "vae", image_objects, generator=generator, seed=seed
            )
            rows = mx.concatenate((anchors, rows), axis=0)
            mx.eval(rows)
        was_resident = self.model is not None
        if self.model is None:
            self.model = DiT(self.model_path, **self.dit_options)
        model = self.model
        record["dit_effective"] = {
            "bits": model.bits,
            "group_size": model.group_size,
            "cache_limit_gb": model.block_limit / 1024**3,
            "checkpoint_fingerprint": model.store.fingerprint(),
        }
        metrics_before = model.metrics.copy()
        record["dit_was_resident"] = was_resident
        setup = time.perf_counter()
        model.prepare_modulation(sigmas, edit=bool(images))
        text = model.refine(hidden)
        record["refiner_cache_hit"] = model.refiner_cache_hit
        record["timings"]["dit_setup_s"] = time.perf_counter() - setup
        nref = 0 if anchors is None else anchors.shape[0]
        trace = {"initial_rows": rows, "conditioning": hidden, "text_tags": tags}
        for ordinal, step in enumerate(selected):
            tick = time.perf_counter()
            velocity = model.predict(
                rows, text, step, h, w, references=references, text_tags=tags
            )
            next_step = selected[ordinal + 1] if ordinal + 1 < len(selected) else 6
            # Subtract float32 schedule values, then multiply and add separately.
            delta = mx.array(sigmas[step], mx.float32) - mx.array(
                sigmas[next_step], mx.float32
            )
            target = rows[nref:] + delta * velocity[nref:]
            rows = (
                target if anchors is None else mx.concatenate((anchors, target), axis=0)
            )
            mx.eval(rows)
            if not mx.all(mx.isfinite(rows)).item():
                record["nonfinite_steps"].append(step)
                if not self.allow_nonfinite:
                    raise FloatingPointError(
                        f"Non-finite denoising state at original head {step}; "
                        "use the native BF16 projection policy and verify the checkpoint"
                    )
                log.warning(
                    "Head %d produced non-finite state; saving an invalid diagnostic output",
                    step,
                )
            if save_trace:
                trace[f"rows.{step}"] = rows
                trace[f"velocity.{step}"] = velocity
            elapsed = time.perf_counter() - tick
            record["timings"][f"step_{step}_s"] = elapsed
            log.info(
                "DiT step %d/6: %.3fs, %.2f GiB active",
                step + 1,
                elapsed,
                mx.get_active_memory() / 1024**3,
            )
        record["dit_metrics"] = {
            k: v - metrics_before[k] for k, v in model.metrics.items()
        }
        if model.async_blocks or model.prefetch_blocks:
            record["dit_metrics"].pop("block_compute_s", None)
            record["dit_metrics"]["timing_note"] = (
                "Weight loading can overlap GPU execution. block_wall_s and step timings measure elapsed time; block_load_s is CPU load time and is not additive."
            )
        latent = unpatchify(rows[nref:], h, w)
        mx.eval(latent)
        if save_trace:
            trace["latent"] = latent
            mx.save_safetensors(str(output.with_suffix(".trace.safetensors")), trace)
        record["memory_before_decode"] = memory_stats()
        record["valid_denoising"] = not record["nonfinite_steps"]
        if not self.keep_resident:
            model.clear()
            self.model = None
        del model, velocity, text, hidden, rows, trace
        gc.collect()
        mx.clear_cache()
        decode_start = time.perf_counter()
        if self.vae is None:
            if self.vae_backend == "coreml":
                from .coreml_vae import CoreMLVAE

                self.vae = CoreMLVAE(self.model_path / "vae", **self.vae_options)
            else:
                self.vae = VAE(self.model_path / "vae", **self.vae_options)
        vae = self.vae
        last_log = [time.perf_counter()]

        def progress(done, total):
            now = time.perf_counter()
            if now - last_log[0] > 20 or done == total:
                log.info("VAE tiles %d/%d", done, total)
                last_log[0] = now

        pixels = vae.decode(latent, on_batch=progress)
        image = to_image(pixels, width, height)
        image.save(output)
        if not self.keep_resident:
            vae.clear()
            self.vae = None
        record["timings"]["decode_s"] = time.perf_counter() - decode_start
        record["timings"]["total_s"] = time.perf_counter() - start
        record["memory"] = memory_stats()
        record["memory"]["scope"] = (
            "active/peak/cache: MLX allocator, Core ML excluded; process_*: macOS parent process including Metal/Core ML, lifetime peak includes earlier requests; overlapping counters, do not sum"
            if self.vae_backend == "coreml"
            else "active/peak/cache: MLX allocator; process_*: macOS parent process including Metal, lifetime peak includes earlier requests; overlapping counters, do not sum"
        )
        record["output"] = str(output)
        record["png_sha256"] = hashlib.sha256(output.read_bytes()).hexdigest()
        record["finished_utc"] = datetime.now(timezone.utc).isoformat()
        output.with_suffix(".json").write_text(json.dumps(record, indent=2) + "\n")
        log.info("Saved %s in %.3fs", output, record["timings"]["total_s"])
        return image, record