| 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 |
| |
| 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 |
|
|
