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