| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import logging |
| from pathlib import Path |
|
|
|
|
| def model_arguments(parser): |
| parser.add_argument( |
| "--model", |
| required=True, |
| help="Downloaded variant or original image model directory", |
| ) |
| parser.add_argument( |
| "--text-encoder", |
| required=True, |
| help="Original MiniMax-H3/text_encoder directory", |
| ) |
| parser.add_argument("--cache-dir", default="cache") |
|
|
|
|
| def main(argv=None): |
| parser = argparse.ArgumentParser( |
| prog="h3-mlx", description="MiniMax H3 Image on MLX/Metal" |
| ) |
| parser.add_argument("--verbose", action="store_true") |
| sub = parser.add_subparsers(dest="command", required=True) |
| sub.add_parser( |
| "doctor", help="Report Metal availability, memory and installed versions" |
| ) |
| download = sub.add_parser( |
| "download", help="Download one variant and shared components" |
| ) |
| download.add_argument( |
| "--variant", choices=("bf16", "packed12", "q8", "q4"), default="packed12" |
| ) |
| download.add_argument("--model-dir", required=True) |
| download.add_argument( |
| "--text-encoder-path", help="Reuse an existing text encoder directory" |
| ) |
| download.add_argument("--skip-text-encoder", action="store_true") |
| download.add_argument("--local-files-only", action="store_true") |
| download.add_argument("--catalog", help="Use an alternate release catalog JSON") |
| download.add_argument( |
| "--verify", action="store_true", help="Rehash all downloaded model files" |
| ) |
| generate = sub.add_parser("generate", help="Run six native PDD steps locally") |
| model_arguments(generate) |
| generate.add_argument("prompt") |
| generate.add_argument("--width", type=int, default=2048) |
| generate.add_argument("--height", type=int, default=2048) |
| generate.add_argument("--seed", type=int, default=731) |
| generate.add_argument("--min-generation-area", type=int, default=4194304) |
| generate.add_argument( |
| "--rng", |
| choices=("torch", "mlx"), |
| default="torch", |
| help="torch CPU matches original native noise; mlx is a different seed stream", |
| ) |
| generate.add_argument("--output", default="outputs/image.png") |
| generate.add_argument( |
| "--image", |
| action="append", |
| default=[], |
| help="Ordered edit reference; repeat up to sixteen times", |
| ) |
| generate.add_argument("--dit-cache-gb", type=float, default=38) |
| generate.add_argument( |
| "--quantize", |
| type=int, |
| choices=(4, 8), |
| help="Experimental weight-only DiT quantization; AdaLN/FP32 islands stay native", |
| ) |
| generate.add_argument("--encoder-quantize", type=int, choices=(4, 8)) |
| generate.add_argument("--group-size", type=int, choices=(32, 64, 128), default=64) |
| generate.add_argument("--mlp-chunk", type=int, default=0) |
| generate.add_argument( |
| "--attention-chunk", |
| type=int, |
| default=0, |
| help="0 uses one memory-efficient Metal SDPA call", |
| ) |
| generate.add_argument( |
| "--kernel", |
| choices=("native", "metal", "metal-qk", "metal-exact", "metal-qk-exact"), |
| default="native", |
| help="metal-exact fuses M5 norms and Q/K/RoPE; metal-qk-exact keeps native full-width norms. Both retain native fallback", |
| ) |
| generate.add_argument( |
| "--tune-matmul", |
| action="store_true", |
| help="Use measured M5 BF16 FC2 tiles for 4096–8192 joint tokens; native fallback otherwise", |
| ) |
| generate.add_argument( |
| "--norm-threads", type=int, choices=(128, 256, 512), default=256 |
| ) |
| generate.add_argument("--vae-batch", type=int, default=1) |
| generate.add_argument("--compile-vae", action="store_true") |
| generate.add_argument("--vae-backend", choices=("mlx", "coreml"), default="mlx") |
| generate.add_argument( |
| "--coreml-package", |
| help="Fixed-tile decoder exported by experiments.v2_coreml_vae", |
| ) |
| generate.add_argument( |
| "--coreml-compute-units", |
| choices=("CPU_ONLY", "CPU_AND_GPU", "CPU_AND_NE", "ALL"), |
| default="CPU_AND_GPU", |
| ) |
| generate.add_argument( |
| "--vae-dtype", choices=("bfloat16", "float32"), default="bfloat16" |
| ) |
| generate.add_argument( |
| "--experimental-steps", help="Original head indices, e.g. 0,2,4; changes images" |
| ) |
| generate.add_argument("--save-trace", action="store_true") |
| generate.add_argument( |
| "--no-refiner-cache", |
| action="store_true", |
| help="Recompute refined conditioning instead of using its exact disk cache", |
| ) |
| generate.add_argument( |
| "--prefetch-blocks", |
| type=int, |
| default=0, |
| help="Overlap weight loading with GPU work, with at most N blocks in flight", |
| ) |
| generate.add_argument( |
| "--async-blocks", |
| action="store_true", |
| help="Enqueue each block asynchronously; use --eval-every to bound in-flight work", |
| ) |
| generate.add_argument( |
| "--compile-blocks", |
| action="store_true", |
| help="Compile each DiT block and fuse elementwise operators", |
| ) |
| generate.add_argument( |
| "--eval-every", |
| type=int, |
| default=1, |
| help="Synchronize every N DiT blocks; 0 evaluates the whole step", |
| ) |
| generate.add_argument( |
| "--keep-refiners", |
| action="store_true", |
| help="Retain token-refiner weights across resident requests", |
| ) |
| generate.add_argument( |
| "--wired-gb", |
| type=float, |
| help="Process-local MLX wired memory; capped at Metal's recommended set", |
| ) |
| generate.add_argument( |
| "--weight-storage", |
| choices=("bf16", "packed12"), |
| default="bf16", |
| help="packed12 caches exact BF16 weights in a smaller representation", |
| ) |
| generate.add_argument( |
| "--quantized-core", help="Preconverted DiT core from python -m h3_mlx.convert" |
| ) |
| generate.add_argument( |
| "--offload-after-denoise", |
| action="store_true", |
| help="Release DiT before VAE for smaller-memory Macs", |
| ) |
| args = parser.parse_args(argv) |
| logging.basicConfig( |
| level=logging.DEBUG if args.verbose else logging.INFO, |
| format="%(asctime)s %(message)s", |
| datefmt="%H:%M:%S", |
| ) |
| if args.command == "download": |
| from .hub import download_variant |
|
|
| result = download_variant( |
| args.variant, |
| args.model_dir, |
| encoder_path=args.text_encoder_path, |
| skip_encoder=args.skip_text_encoder, |
| local_files_only=args.local_files_only, |
| catalog_path=args.catalog, |
| verify=args.verify, |
| ) |
| print(json.dumps(result, indent=2)) |
| return |
| if args.command == "doctor": |
| from .runtime import configure_memory, environment |
|
|
| configure_memory() |
| print(json.dumps(environment(), indent=2)) |
| return |
| if args.command == "generate": |
| for p in (args.model, args.text_encoder): |
| if not Path(p).is_dir(): |
| parser.error(f"Model directory does not exist: {p}") |
| if ( |
| args.mlp_chunk < 0 |
| or args.attention_chunk < 0 |
| or args.vae_batch < 1 |
| or args.dit_cache_gb < 0 |
| or args.eval_every < 0 |
| or args.prefetch_blocks < 0 |
| ): |
| parser.error( |
| "Chunk sizes/cache must be nonnegative and VAE batch must be positive" |
| ) |
| if len(args.image) > 16: |
| parser.error("MiniMax-H3-Image accepts at most sixteen reference images") |
| steps = ( |
| None |
| if args.experimental_steps is None |
| else [int(s) for s in args.experimental_steps.split(",")] |
| ) |
| from .runtime import Pipeline |
|
|
| pipe = Pipeline( |
| args.model, |
| args.text_encoder, |
| cache_dir=args.cache_dir, |
| dit_cache_gb=args.dit_cache_gb, |
| bits=args.quantize, |
| encoder_bits=args.encoder_quantize, |
| group_size=args.group_size, |
| mlp_chunk=args.mlp_chunk, |
| query_chunk=args.attention_chunk, |
| kernel=args.kernel, |
| threads=args.norm_threads, |
| vae_batch=args.vae_batch, |
| vae_dtype=args.vae_dtype, |
| keep_resident=not args.offload_after_denoise, |
| quantized_path=args.quantized_core, |
| compile_blocks=args.compile_blocks, |
| eval_every=args.eval_every, |
| keep_refiners=args.keep_refiners, |
| wired_gb=args.wired_gb, |
| weight_storage=args.weight_storage, |
| vae_compile=args.compile_vae, |
| vae_backend=args.vae_backend, |
| coreml_package=args.coreml_package, |
| coreml_compute_units=args.coreml_compute_units, |
| async_blocks=args.async_blocks, |
| cache_refined_text=not args.no_refiner_cache, |
| prefetch_blocks=args.prefetch_blocks, |
| tune_matmul=args.tune_matmul, |
| ) |
| _, record = pipe.generate( |
| args.prompt, |
| width=args.width, |
| height=args.height, |
| seed=args.seed, |
| min_area=args.min_generation_area, |
| rng=args.rng, |
| output=args.output, |
| images=args.image, |
| steps=steps, |
| save_trace=args.save_trace, |
| ) |
| print( |
| json.dumps( |
| { |
| "output": record["output"], |
| "timings": record["timings"], |
| "memory": record["memory"], |
| }, |
| indent=2, |
| ) |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
