inference.py

#!/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())