Initial release: Lance-3B-MLX first MLX port, T2I verified end-to-end

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+samples/orange_cat_chair.png filter=lfs diff=lfs merge=lfs -text
+samples/snowy_peak.png filter=lfs diff=lfs merge=lfs -text
+samples/sunset_mountains.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -3,111 +3,163 @@ license: apache-2.0
 base_model:
 - bytedance-research/Lance
 - Qwen/Qwen2.5-VL-3B-Instruct
-pipeline_tag: any-to-any
+pipeline_tag: text-to-image
 library_name: mlx
 tags:
 - multimodal
 - mlx
 - apple-silicon
+- text-to-image
 - image-generation
 - video-generation
-- image-editing
-- video-understanding
-- any-to-any
+- diffusion
+- flow-matching
+- moe
+- qwen2_5_vl
+- wan
 - port
 ---
 
 # Lance-3B-MLX
 
-A native [MLX](https://github.com/ml-explore/mlx) port of [ByteDance's Lance](https://huggingface.co/bytedance-research/Lance) — a 3B-parameter unified multimodal model for image and video generation, editing, and understanding.
+First native [MLX](https://github.com/ml-explore/mlx) port of [ByteDance Research's Lance](https://huggingface.co/bytedance-research/Lance) — a 3B-parameter unified multimodal model for image/video generation, editing, and understanding. Runs natively on Apple Silicon, no CUDA required.
 
-Built on top of the Qwen2.5-VL-3B-Instruct backbone, with Lance's custom multi-task adapters and a Wan 2.2 VAE.
+The architecture is **Qwen2.5-VL-3B + parallel MoE-gen experts + Wan 2.2 VAE**. Lance uses a "Mixture-of-Tokens" routing: every attention block and MLP has a parallel `*_moe_gen` branch. Text tokens go through normal weights; VAE-latent (generation) tokens go through the `_moe_gen` weights, in the same forward pass.
 
-## Status
+## What works
 
-**Weight conversion is complete** — all tensors from the upstream PyTorch
-checkpoint are present in MLX safetensors layout, verified bit-exact on
-sampled tensors. **Inference wrapper is not yet runnable.**
-
-| Component | Status |
+| Capability | Status |
+|---|---|
+| Text-to-image (T2I), single image, CFG | ✅ Working, verified |
+| Strict load of all 1021 LLM/adapter tensors | ✅ Working |
+| Wan 2.2 VAE encode/decode (T=1) | ✅ Working (uses [RockTalk/Wan2.2-VAE-MLX](https://huggingface.co/RockTalk/Wan2.2-VAE-MLX)) |
+| Flow-matching denoising loop | ✅ Working |
+| Classifier-free guidance | ✅ Working |
+| 3D mrope position embeddings | ✅ Working |
+| MoE-gen routing (per-token attention + MLP + layernorm) | ✅ Working |
+| Text-to-video (T2V) | ⏳ Needs VAE T>1 streaming cache + larger pos embed |
+| Image/video editing (TI2I, TIV2V) | ⏳ Phase 2 — needs ViT integration |
+| X→T (image/video understanding) | ⏳ Phase 2 — needs AR sampling loop + KV cache |
+
+## Sample generations
+
+Verified on M4 Studio (128 GB). 30 steps, CFG=4, 512×512:
+
+| Prompt | Output |
 |---|---|
-| Weight conversion (PT → MLX safetensors, layout + name remaps) | ✅ DONE — bit-exact spot check (40/40 sampled tensors) |
-| `modeling_utils` (TimestepEmbedder, PositionEmbedding3D, MLP, sincos tables) | ✅ DONE |
-| `vae_wan22` — image-mode encode/decode | ✅ DONE |
-| `vae_wan22` — video streaming feat-cache | ⏳ PENDING |
-| `lance.py` adapters (vae2llm, llm2vae, time/pos embed) | ⚠ PARTIAL — primitives present; patchify path doesn't match source, needs debug |
-| Lance's MoE-gen attention (`q_proj_moe_gen` etc — 505 tensors) | ⏳ NOT YET WRAPPED — weights bundled, wrapper class still uses bare Qwen2.5-VL |
-| Lance's QK-norm extension (`q_norm`/`k_norm` in attention — 73 tensors) | ⏳ NOT YET WRAPPED — weights bundled |
-| Flow-matching sampler (`validation_gen` — T2I/T2V/edit) | ⏳ STUB |
-| X→T (understanding) autoregressive loop | ⏳ Phase 2 |
-| NaViT variable-resolution image packing | ⏳ Phase 2 |
-
-**What this repo gives you today:** correctly-converted weights you can load
-in MLX, with the bundle keys matching mlx-vlm's Qwen2.5-VL convention plus
-Lance's own adapter top-levels (`vae2llm`, `llm2vae`, `time_embedder`,
-`latent_pos_embed`, `vae_model.*`). The MoE-gen generation-path weights are
-included verbatim so any future MoE-aware wrapper can pick them up.
-
-**What's missing for end-to-end inference:** the Python wrapper class needs
-(1) MoE-routing attention layers with parallel und/gen weight sets, (2)
-QK-norm, (3) a corrected patchify path before `vae2llm`, (4) the
-flow-matching denoising loop + CFG. The source PT implementation in
-`Lance/modeling/lance/qwen2_navit.py` (~1300 lines) and `lance.py` (~1900
-lines) is the reference.
+| *"a photo of a sunset over mountains"* | ![sunset](samples/sunset_mountains.png) |
+| *"a fluffy orange cat sitting on a wooden chair, photorealistic"* | ![cat](samples/orange_cat_chair.png) |
+| *"a majestic snowy mountain peak with a dramatic blue sky and clouds"* | ![mountain](samples/snowy_peak.png) |
+
+## Performance
+
+Measured on M4 Studio (128 GB) at CFG=4 (one conditional + one unconditional forward per step):
+
+| Resolution | Steps | Per-step | Total sample | VAE decode |
+|---|---|---|---|---|
+| 256×256 | 24 | ~400 ms | ~9.6 s | ~0.1 s |
+| 512×512 | 30 | ~1.2 s  | ~36 s | ~0.5 s |
+
+First-call kernel-compile penalty: ~few seconds per new resolution.
 
 ## Files
 
-- `model.safetensors` — Lance 3B image variant LLM + adapters (Qwen2.5-VL language model, vae2llm/llm2vae, time_embedder, latent_pos_embed), MLX-layout, **22.9 GB** (1021 tensors, ~6.19B params incl. embed table)
-- `vit.safetensors` — Qwen2.5-VL ViT visual encoder, MLX-layout (NTHWC conv weights), **1.25 GB** (390 tensors, ~668M params, fp16 — bundled here for offline use; the source ships it as a separate shard)
-- `vae.safetensors` — Wan 2.2 VAE, MLX-layout, **2.62 GB** (196 tensors, ~705M params). Converted from the upstream `Wan2.2_VAE.pth` pickle.
-- `config.json` — distilled architecture config + embedded Qwen2.5-VL sub-config
-- `vit_config.json` — Qwen2.5-VL ViT sub-config
-- `tokenizer.json`, `tokenizer_config.json`, `vocab.json`, `merges.txt`, `generation_config.json` — copied verbatim from upstream
+| File | Size | Description |
+|---|---|---|
+| `model.safetensors` | 23 GB | LLM (Qwen2.5-VL with MoE-gen) + Lance adapters, 1021 tensors |
+| `vit.safetensors` | 1.25 GB | Qwen2.5-VL ViT (for understanding mode — Phase 2) |
+| `vae.safetensors` | 2.62 GB | Wan 2.2 VAE (older keying — for compatibility; the standalone [RockTalk/Wan2.2-VAE-MLX](https://huggingface.co/RockTalk/Wan2.2-VAE-MLX) uses cleaner keys and is recommended) |
+| `config.json` | — | Distilled architecture config |
+| `tokenizer.json`, `vocab.json`, `merges.txt` | — | Qwen2.5-VL tokenizer, verbatim |
+| `samples/*.png` | — | Verified T2I outputs from this checkpoint |
 
-## Hardware
+## Usage
 
-Targets Apple Silicon with unified memory. Verified on M3 Ultra (512 GB). Lower-RAM Macs may need to run the LLM forward only (no joint backbone + VAE).
+Requires `mlx >= 0.29`, `mlx-vlm >= 0.3`, `numpy`, `einops`, `transformers`, `pillow`, and the [`lance-mlx`](https://github.com/RockTalk/Lance-MLX) companion repo for the `Lance` Python class.
 
-## Loading
+```bash
+pip install mlx mlx-vlm numpy einops transformers pillow
+```
 
 ```python
 import mlx.core as mx
-import mlx.nn as nn
 from lance_mlx.lance import Lance, LanceConfig
-from mlx_vlm.models.qwen2_5_vl.config import ModelConfig as Qwen25VLConfig
-import json
+from lance_mlx.vae_wan22 import Wan2_2_VAE
+
+# Build + strict-load (see tools/lance_t2i.py in the companion repo for the
+# full builder; LanceConfig takes a Qwen2.5-VL ModelConfig built from
+# config.json).
+model = Lance(lance_cfg)
+model.load_weights(list(mx.load("model.safetensors").items()), strict=True)
+
+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("vae.safetensors").items()), strict=True)
+
+# Sample
+latent = model.sample_t2i(
+    prompt_token_ids=text_ids,                  # (P,) int32 from tokenizer (no specials)
+    latent_shape=(1, 32, 32),                   # (T_lat, H_lat, W_lat) for 512×512 image
+    special_token_ids={"bos": 151644, "eos": 151645,
+                       "start_of_image": 151652, "end_of_image": 151653,
+                       "image_token_id": 151655},
+    num_steps=30, timestep_shift=3.5, cfg_scale=4.0, seed=0,
+)
+img = vae.decode(latent)                        # (1, 1, 512, 512, 3) in [-1, 1]
+```
+
+End-to-end script: `tools/lance_t2i.py` in the [companion repo](https://github.com/RockTalk/Lance-MLX).
 
-from lance_mlx import load_lance
+## How the MoE-gen routing is implemented in MLX
 
-model, cfg = load_lance("./")
-# cfg['_loaded_into_model'] reports how many tensors landed in the wrapper.
-# cfg['_vae_weights'] holds the Wan VAE keys (load into a separate vae module).
-# cfg['_moe_gen_weights'] holds Lance's generation-path weights, parked until
-# a MoE-aware wrapper is available.
+Lance's checkpoint contains *two* sets of weights per Qwen2 block:
+
+```
+self_attn.{q,k,v,o}_proj         self_attn.{q,k,v,o}_proj_moe_gen
+self_attn.{q,k}_norm             self_attn.{q,k}_norm_moe_gen
+mlp.{gate,down,up}_proj          mlp_moe_gen.{gate,down,up}_proj
+input_layernorm                  input_layernorm_moe_gen
+post_attention_layernorm         post_attention_layernorm_moe_gen
 ```
 
-## Citation
+For T2I/T2V the sequence layout is:
 
-```bibtex
-@article{lance2026,
-  title   = {Lance: Unified Multimodal Modeling by Multi-Task Synergy},
-  author  = {Fu, Fengyi and Huang, Mengqi and Wu, Shaojin and Jiang, Yunsheng and Huo, Yufei and Guo, Jianzhu and others},
-  journal = {arXiv preprint arXiv:2605.18678},
-  year    = {2026},
-  url     = {http://arxiv.org/abs/2605.18678}
-}
 ```
+<|im_start|> [prompt tokens] <|im_end|> <|vision_start|> [N latent placeholders] <|vision_end|>
+                                                          └──── routed through moe_gen ────┘
+                                                          ↑ everything else: normal weights
+```
+
+The MLX port (`qwen2_navit_mlx.py`) routes by slicing the sequence into the latent slab vs the surrounding text, applying the appropriate expert to each slab, and concatenating. mrope position ids continue to flow normally across both slabs (with axis-T/H/W coordinates only varying inside the latent slab).
+
+## Conversion source
+
+Converted from `bytedance-research/Lance/Lance_3B/*` using the open-source pipeline at https://github.com/RockTalk/Lance-MLX (`tools/convert_weights.py`). Layout transforms:
+
+- Conv weights: PT `(O, I, [T,] H, W)` → MLX `(O, [T,] H, W, I)`
+- Embedding weights: shape preserved
+- `lm_head.weight` tied to `embed_tokens.weight` (Qwen default)
+- All `*_moe_gen.*` keys copied verbatim under the same names
 
 ## License
 
-Apache-2.0, inherited from upstream `bytedance-research/Lance`.
+Apache 2.0, inherited from upstream `bytedance-research/Lance`. The Wan 2.2 VAE component is also Apache 2.0 from Alibaba's Wan team.
 
-## Acknowledgments
+## Acknowledgements
 
-- ByteDance Research for the original Lance training and PyTorch release
-- The `mlx` and `mlx-vlm` teams at Apple
-- Qwen team for Qwen2.5-VL-3B-Instruct
+- **ByteDance Research** — original Lance training + PT release
+- **Qwen team** — Qwen2.5-VL-3B-Instruct backbone
+- **Alibaba Wan team** — Wan 2.2 VAE training
+- **Apple `mlx` and `mlx-vlm` teams** — the underlying frameworks
+- **This MLX port** — RockTalk
 
----
+## Citation
 
-**Port status reporting honestly:** this repo currently provides MLX-format weights with verified-loading scaffolding. Inference sampling (T2I/T2V) is a follow-up release; the building blocks are in place but the diffusion loop has not been parity-validated end-to-end yet. Pull requests welcome.
+```bibtex
+@misc{lance_mlx,
+  title  = {Lance-3B-MLX — First MLX port of ByteDance's Lance},
+  author = {RockTalk},
+  year   = {2026},
+  url    = {https://huggingface.co/RockTalk/Lance-3B-MLX}
+}
+```

config.json

@@ -67,14 +67,15 @@
   },
   "latent_patch_size": [
     1,
-    2,
-    2
+    1,
+    1
   ],
-  "max_latent_size": 32,
-  "max_num_frames": 25,
+  "max_latent_size": 64,
+  "max_num_frames": 0,
   "latent_channel": 48,
   "vae_downsample_spatial": 16,
   "vae_downsample_temporal": 4,
   "connector_act": "gelu_pytorch_tanh",
-  "timestep_shift": 3.5
+  "timestep_shift": 3.5,
+  "max_num_latent_frames": 1
 }