Add self-contained inference.py + bundled lance_mlx package

inference.py

@@ -0,0 +1,255 @@
+#!/usr/bin/env python3
+"""End-to-end Text-to-Video (and Text-to-Image) inference for Lance-3B-Video-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:
+    # T2V (default — 9-frame video at 256x256, ~22 s on M4 Studio)
+    python inference.py --prompt "a calm ocean wave rolling onto a sandy beach"
+
+    # Tune frame count: T = (T_lat - 1) * 4 + 1
+    #   T_lat=1 -> 1 frame (image), T_lat=3 -> 9, T_lat=8 -> 29, T_lat=31 -> 121
+    python inference.py --prompt "..." --t-lat 8
+
+    # T2I fast path
+    python inference.py --prompt "..." --t-lat 1 --size 512 --steps 30
+
+Outputs:
+  - <out>.png            — horizontal strip of all frames
+  - <out>_frame*.png     — each frame as a separate file
+  - <out>.mp4            — MP4 (if --mp4 and `imageio[ffmpeg]` is installed)
+
+Verified on M4 Studio (128 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
+
+_materialize = getattr(mx, "eval")
+
+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:
+    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",
+    ))
+    target = repo_dir / "wan22_vae.safetensors"
+    try:
+        target.symlink_to(downloaded)
+    except OSError:
+        import shutil
+        shutil.copy(downloaded, target)
+    return target
+
+
+def save_outputs(video_np: np.ndarray, out_path: Path, want_mp4: bool, fps: int) -> None:
+    """video_np: (T, H, W, 3) in [-1, 1]. Saves strip + per-frame PNGs + optional MP4."""
+    video_u8 = np.clip((video_np + 1.0) * 127.5, 0, 255).astype(np.uint8)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+
+    strip = np.concatenate(list(video_u8), axis=1)  # (H, T*W, 3)
+    Image.fromarray(strip).save(out_path)
+    print(f"[ok] frame strip -> {out_path}")
+
+    base = out_path.with_suffix("")
+    for i, frame in enumerate(video_u8):
+        Image.fromarray(frame).save(f"{base}_frame{i:02d}.png")
+    print(f"[ok] per-frame PNGs -> {base}_frame*.png")
+
+    if want_mp4:
+        try:
+            import imageio.v3 as iio
+            mp4_path = out_path.with_suffix(".mp4")
+            iio.imwrite(mp4_path, video_u8, fps=fps, codec="libx264")
+            print(f"[ok] mp4 -> {mp4_path}")
+        except Exception as exc:
+            print(f"[warn] MP4 export failed: {exc}")
+            print("       Install with: pip install 'imageio[ffmpeg]'")
+
+
+def main(args: argparse.Namespace) -> None:
+    repo = REPO_DIR
+    n_frames = (args.t_lat - 1) * 4 + 1
+    print("=== Lance-3B-Video-MLX ===")
+    print(f"prompt: {args.prompt!r}")
+    print(f"out:    {args.out}")
+    print(f"size:   {args.size}x{args.size}  steps: {args.steps}  "
+          f"T_lat={args.t_lat} → {n_frames} frames  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"))
+    non_vit = {k: v for k, v in weights.items() if not k.startswith("vit_model.")}
+    n_vit = len(weights) - len(non_vit)
+    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, "
+          f"dropped {n_vit} ViT tensors not needed for generation)")
+
+    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 = (args.t_lat, 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 (streaming for T>1) ...")
+    t0 = time.time()
+    video = vae.decode(final_latent)
+    _materialize(video)
+    print(f"[ok] VAE decode ({time.time()-t0:.1f}s)  shape={video.shape}")
+
+    video_np = np.asarray(video).squeeze(0)  # (T, H, W, 3)
+    save_outputs(video_np, Path(args.out), want_mp4=args.mp4, fps=args.fps)
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--prompt", default="a calm ocean wave rolling onto a sandy beach")
+    ap.add_argument("--out", default="output.png")
+    ap.add_argument("--steps", type=int, default=24)
+    ap.add_argument("--size", type=int, default=256,
+                    help="square frame size (256 recommended for T2V)")
+    ap.add_argument("--t-lat", type=int, default=3,
+                    help="latent frame count; output frames = (t_lat-1)*4 + 1")
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--cfg", type=float, default=4.0)
+    ap.add_argument("--mp4", action="store_true",
+                    help="also export an MP4 (needs `pip install 'imageio[ffmpeg]'`)")
+    ap.add_argument("--fps", type=int, default=8,
+                    help="MP4 frame rate")
+    main(ap.parse_args())