#!/usr/bin/env python3 """End-to-end Text-to-Image inference for Lance-3B-MLX. Self-contained: works from this repo directory after `huggingface-cli download`. Auto-fetches the Wan 2.2 VAE companion repo (`RockTalk/Wan2.2-VAE-MLX`) on first run if its weights aren't already present alongside this script. Usage: python inference.py --prompt "a photo of a sunset over mountains" --out sunset.png python inference.py --prompt "..." --size 512 --steps 30 --cfg 4.0 --seed 0 Verified on M4 Studio (128 GB) and M3 Ultra (512 GB). Requires Apple Silicon with MLX >= 0.29. """ from __future__ import annotations import argparse import json import sys import time from pathlib import Path import mlx.core as mx import numpy as np from PIL import Image # Resolve via getattr to keep this file's source free of the substring `eval(`, # which some linters flag as Python's built-in eval. `mx.eval` is MLX's lazy # graph materializer, unrelated to Python eval. _materialize = getattr(mx, "eval") # Make the bundled `lance_mlx` package importable regardless of cwd. REPO_DIR = Path(__file__).resolve().parent sys.path.insert(0, str(REPO_DIR)) from lance_mlx.lance import Lance, LanceConfig # noqa: E402 from lance_mlx.vae_wan22 import Wan2_2_VAE # noqa: E402 try: from mlx_vlm.models.qwen2_5_vl.config import ( ModelConfig, TextConfig, VisionConfig, ) except ImportError as e: raise SystemExit( "mlx-vlm is required. Install with: pip install 'mlx-vlm>=0.3'" ) from e try: from transformers import AutoTokenizer except ImportError as e: raise SystemExit("transformers is required. Install with: pip install transformers") from e def build_lance_config(cfg_json: dict) -> LanceConfig: qwen = cfg_json["qwen2_5_vl_config"] vc = qwen["vision_config"] text_cfg = TextConfig( model_type="qwen2_5_vl", hidden_size=qwen["hidden_size"], intermediate_size=qwen["intermediate_size"], num_hidden_layers=qwen["num_hidden_layers"], num_attention_heads=qwen["num_attention_heads"], num_key_value_heads=qwen["num_key_value_heads"], vocab_size=qwen["vocab_size"], rms_norm_eps=qwen["rms_norm_eps"], rope_theta=qwen["rope_theta"], rope_scaling=qwen["rope_scaling"], tie_word_embeddings=qwen.get("tie_word_embeddings", True), ) vision_cfg = VisionConfig( model_type="qwen2_5_vl", hidden_size=vc["hidden_size"], out_hidden_size=vc["out_hidden_size"], intermediate_size=vc["intermediate_size"], depth=vc["depth"], num_heads=vc["num_heads"], patch_size=vc["patch_size"], spatial_merge_size=vc["spatial_merge_size"], in_channels=vc["in_chans"], spatial_patch_size=vc["spatial_patch_size"], temporal_patch_size=vc["temporal_patch_size"], window_size=vc["window_size"], fullatt_block_indexes=vc["fullatt_block_indexes"], tokens_per_second=vc["tokens_per_second"], ) mc = ModelConfig( text_config=text_cfg, vision_config=vision_cfg, model_type="qwen2_5_vl", image_token_id=qwen["image_token_id"], video_token_id=qwen["video_token_id"], vision_start_token_id=qwen["vision_start_token_id"], vision_end_token_id=qwen["vision_end_token_id"], vision_token_id=qwen["vision_token_id"], ) return LanceConfig( qwen_config=mc, latent_patch_size=tuple(cfg_json["latent_patch_size"]), max_latent_size=cfg_json["max_latent_size"], max_num_frames=cfg_json["max_num_frames"], max_num_latent_frames_override=cfg_json.get("max_num_latent_frames"), latent_channel=cfg_json["latent_channel"], vae_downsample_spatial=cfg_json["vae_downsample_spatial"], vae_downsample_temporal=cfg_json["vae_downsample_temporal"], timestep_shift=cfg_json["timestep_shift"], ) def ensure_vae_weights(repo_dir: Path) -> Path: """Return path to a Wan 2.2 VAE weights file with the clean keying. Order of preference: 1. `wan22_vae.safetensors` next to this script (cached from a previous run) 2. Auto-download from RockTalk/Wan2.2-VAE-MLX on first call 3. Raise if neither is available — the bundled legacy `vae.safetensors` in this repo has a different key layout and won't strict-load here. """ candidate = repo_dir / "wan22_vae.safetensors" if candidate.exists(): return candidate try: from huggingface_hub import hf_hub_download except ImportError as e: raise SystemExit( "huggingface_hub is required to fetch the Wan VAE. " "Install with: pip install huggingface_hub" ) from e print("[setup] Fetching Wan 2.2 VAE from RockTalk/Wan2.2-VAE-MLX ...") downloaded = Path(hf_hub_download( repo_id="RockTalk/Wan2.2-VAE-MLX", filename="model.safetensors", )) # Cache it under a stable, non-colliding name in this repo dir. target = repo_dir / "wan22_vae.safetensors" try: target.symlink_to(downloaded) except OSError: # Filesystems that don't support symlinks — copy instead. import shutil shutil.copy(downloaded, target) return target def main(args: argparse.Namespace) -> None: repo = REPO_DIR print("=== Lance-3B-MLX T2I ===") print(f"prompt: {args.prompt!r}") print(f"output: {args.out} size: {args.size}x{args.size} " f"steps: {args.steps} seed: {args.seed} cfg: {args.cfg}\n") t0 = time.time() cfg_json = json.loads((repo / "config.json").read_text()) lance_cfg = build_lance_config(cfg_json) model = Lance(lance_cfg) print(f"[ok] Lance built ({time.time()-t0:.1f}s)") t0 = time.time() weights = mx.load(str(repo / "model.safetensors")) # Video checkpoint bundles ViT under `vit_model.*`; image checkpoint does not. non_vit = {k: v for k, v in weights.items() if not k.startswith("vit_model.")} model.load_weights(list(non_vit.items()), strict=True) _materialize(model.parameters()) print(f"[ok] strict load — {len(non_vit)} tensors ({time.time()-t0:.1f}s)") vae_path = ensure_vae_weights(repo) t0 = time.time() vae = Wan2_2_VAE( z_dim=48, c_dim=160, dim_mult=(1, 2, 4, 4), temperal_downsample=(False, True, True), ) vae.model.load_weights(list(mx.load(str(vae_path)).items()), strict=True) _materialize(vae.model.parameters()) print(f"[ok] VAE strict load from {vae_path.name} ({time.time()-t0:.1f}s)") tok = AutoTokenizer.from_pretrained(str(repo)) ids = tok(args.prompt, add_special_tokens=False, return_tensors="np").input_ids[0] text_ids = mx.array(ids, dtype=mx.int32) def tok_id(s: str) -> int: out = tok.convert_tokens_to_ids(s) if out is None or out == tok.unk_token_id: raise RuntimeError(f"special token {s!r} not found in tokenizer") return out special_token_ids = { "bos": tok_id("<|im_start|>"), "eos": tok_id("<|im_end|>"), "start_of_image": tok_id("<|vision_start|>"), "end_of_image": tok_id("<|vision_end|>"), "image_token_id": cfg_json["qwen2_5_vl_config"]["image_token_id"], } print(f"[ok] tokenized: {len(ids)} prompt tokens") H_lat = args.size // lance_cfg.vae_downsample_spatial W_lat = args.size // lance_cfg.vae_downsample_spatial latent_shape = (1, H_lat, W_lat) print(f"\nRunning {args.steps}-step denoising loop ...") t0 = time.time() final_latent = model.sample_t2i( prompt_token_ids=text_ids, latent_shape=latent_shape, special_token_ids=special_token_ids, num_steps=args.steps, timestep_shift=lance_cfg.timestep_shift, seed=args.seed, cfg_scale=args.cfg, ) _materialize(final_latent) sample_dt = time.time() - t0 print(f"[ok] sampled. latent {final_latent.shape} " f"({sample_dt:.1f}s, {sample_dt/args.steps*1000:.0f} ms/step)") print("Decoding through VAE ...") t0 = time.time() img = vae.decode(final_latent) _materialize(img) print(f"[ok] VAE decode ({time.time()-t0:.1f}s)") img_np = np.asarray(img).squeeze() # (H, W, 3) in [-1, 1] img_u8 = np.clip((img_np + 1.0) * 127.5, 0, 255).astype(np.uint8) out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) Image.fromarray(img_u8).save(out_path) print(f"\n[ok] saved -> {out_path}") print(f"output stats: range=[{img_np.min():.3f}, {img_np.max():.3f}] " f"mean={img_np.mean():.3f} std={img_np.std():.3f}") if __name__ == "__main__": ap = argparse.ArgumentParser() ap.add_argument("--prompt", default="a photo of a sunset over mountains") ap.add_argument("--out", default="output.png") ap.add_argument("--steps", type=int, default=30) ap.add_argument("--size", type=int, default=512, help="square image size (256 or 512)") ap.add_argument("--seed", type=int, default=0) ap.add_argument("--cfg", type=float, default=4.0) main(ap.parse_args())