Add Moebius ONNX exports (unet + VAE enc/dec) + model card + lab notes

LICENSE

@@ -0,0 +1,199 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to the Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by the Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding any notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. Please also get an
+      "Alarm or alarm" page at http://www.apache.org/
+
+   Copyright 2024 Moebius Authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,3 +1,76 @@
 ---
 license: apache-2.0
+library_name: onnx
+pipeline_tag: image-to-image
+base_model: hustvl/Moebius
+tags:
+  - image-inpainting
+  - inpainting
+  - diffusion
+  - onnx
+  - onnxruntime-web
+  - webgpu
+  - in-browser
+language:
+  - en
 ---
+
+# Moebius — ONNX (browser / WebGPU)
+
+ONNX exports of the [**Moebius**](https://huggingface.co/hustvl/Moebius) 0.22B lightweight
+image-inpainting model ([hustvl/Moebius](https://github.com/hustvl/Moebius), ECCV'26), packaged
+to run **fully client-side in a web browser** via
+[ONNX Runtime Web](https://onnxruntime.ai/docs/tutorials/web/) on the **WebGPU** backend.
+
+No text encoder, no tokenizer — Moebius conditions on a small learned embedding table, so the
+whole pipeline is three graphs plus a short DDIM loop you drive in JS.
+
+## Files
+
+| File | Graph | Input → Output | Size (fp32) |
+|------|-------|----------------|-------------|
+| `unet.onnx` | student denoiser (`RemovalModel`: embedding + lambda-DWConv UNet) | `latent (B,9,64,64)`, `timesteps (B,)`, `input_ids (B,10)` → `noise (B,4,64,64)` | ~907 MB |
+| `vae_encoder.onnx` | SD VAE encoder | `image (B,3,512,512)` → `moments (B,8,64,64)` | ~137 MB |
+| `vae_decoder.onnx` | SD VAE decoder | `latent (B,4,64,64)` → `image (B,3,512,512)` | ~198 MB |
+
+- Exported at a **static 512×512** resolution (64×64 latent). The model's cross-attention uses a
+  relative-position embedding tied to the trained resolution, so spatial size is fixed.
+- The learned-embedding "prompt" conditioning stays inside `unet.onnx` as an `nn.Embedding(20, 3072)`
+  gather. For classifier-free guidance: `input_ids` rows `[0..9]` = conditional, `[10..19]` = unconditional.
+
+## Pipeline notes (must match for correct output)
+
+- **VAE `scaling_factor = 0.13025`** (this is a custom VAE — *not* the usual SD `0.18215`).
+  Encode: `latent = mean(moments[:, :4]) * 0.13025`. Decode: feed `latent / 0.13025`.
+- 9-channel UNet input = `concat([noisy_latent(4), mask(1), masked_image_latent(4)], dim=1)`.
+- Scheduler: DDIM, `beta_start=0.00085`, `beta_end=0.012`, `scaled_linear`, 1000 train steps,
+  `clip_sample=false`. 20 steps with `strength≈0.99` ⇒ 19 actual steps.
+- VAE encoder source: [`hustvl/PixelHacker`](https://huggingface.co/hustvl/PixelHacker) `vae/`.
+
+A reference TypeScript implementation (DDIM loop, CFG, 9-channel assembly, pre/post-processing)
+that loads these files lives in the accompanying web demo.
+
+## Precision
+
+These are **fp32** exports (chosen for guaranteed numeric parity with the reference pipeline).
+Parity vs PyTorch on CPU EP: decoder `max|Δ|≈5.7e-5`, unet `≈3.6e-6`. A full-pipeline check vs the
+Torch reference (identical noise) gives a decoded-image `mean|Δ|≈0.0022`. fp16 halves the download
+but is a quality gamble in the lambda layers and unstable for the VAE — validate before using.
+
+## License & attribution
+
+Licensed under **Apache 2.0**, inherited from the upstream
+[hustvl/Moebius](https://huggingface.co/hustvl/Moebius). These artifacts are a format conversion
+(PyTorch → ONNX) of the original weights; all model credit belongs to the original authors.
+
+```bibtex
+@misc{DuanAndXu2026Moebius,
+  title  = {Moebius: 0.2B Lightweight Image Inpainting Framework with 10B-Level Performance},
+  author = {Kangsheng Duan and Ziyang Xu and Wenyu Liu and Xiaohu Ruan and Xiaoxin Chen and Xinggang Wang},
+  year   = {2026},
+  eprint = {2606.19195},
+  archivePrefix = {arXiv},
+  primaryClass  = {cs.CV},
+  url    = {https://arxiv.org/abs/2606.19195}
+}
+```

notes.md

@@ -0,0 +1,132 @@
+# Notes / lab log
+
+Running log of what Claude Opus 4.8 in Claude Code figured out. Newest at the bottom of each section.
+
+## Environment
+- macOS (darwin arm64), Apple Silicon. No CUDA → torch CPU/MPS build.
+- System python 3.9.6; using `uv` to manage an isolated env.
+- git-lfs 3.7.1 present. Weights cloned to `/tmp/Moebius/Moebius-weights`
+  (pretrained, ft_celebahq, ft_ffhq, ft_places2 — each ~450MB fp32 `.bin`).
+- Code repo at `/tmp/Moebius/Moebius`.
+
+## Code map (what matters for the port)
+- Entry: `infer/infer_moebius.py` → `infer/utils.py:build_pipeline`.
+- Pipeline: `removal/v1_2/pipeline.py` (`RemovalSDXLPipeline_BatchMode`).
+- Model wrapper: `removal/v1_2/removal_model.py` (`RemovalModel` = embedding + diff UNet).
+- Core helpers: `utils_infer.py` (`encode_clean_latents`, `predict_noise`).
+- UNet impl: `model_lib/nets/unet_lambda_prune_lite.py` (+ lambda layers under
+  `model_lib/nets/layers/λ/vanillaλ.py`).
+- Config: `config/model_cfg/moebius.yaml`.
+
+## Key findings
+- The CUDA/Triton `fla` dependency is ONLY imported in `model_lib/nets/layers/gla/gla.py`
+  (the GLA teacher variant). Moebius's student UNet (lambda-DWConv) does not need it —
+  must avoid importing `unet_gla` to keep the graph clean for export.
+- "Prompt" conditioning is a plain `nn.Embedding(20, 3072)`. CFG uses fixed ids:
+  cond=[0..9], uncond=[10..19]. So encoder_hidden_states is a constant per branch →
+  can be precomputed and baked into the ONNX UNet as a constant, OR passed as input.
+- 9-channel UNet input = cat([noisy_latents(4), resized_mask(1), masked_latents(4)], dim=1).
+- CFG batches uncond+cond into one forward (batch dim ×2), then splits.
+- einsum is used in the λ layers (linear attention). Supported in ONNX; need to check
+  ORT-Web WebGPU coverage.
+
+## Phase 1 results (reference inference — DONE)
+- Got the real pipeline running end-to-end on CPU (macOS, torch 2.7.1).
+- Patches needed to load student on CPU/mac:
+  - `model_lib/__init__.py`: wrapped teacher `unet_gla` import in try/except (needs `fla`).
+  - Don't import `utils_train` (drags in orjson/library); `build_vae` is just
+    `AutoencoderKL.from_pretrained(vae_dir)`.
+- **VAE scaling_factor = 0.13025** (NOT the usual SD 0.18215!). Custom VAE.
+  block_out_channels = [128,256,512,512] → vae_scale_factor 8. This MUST be hardcoded
+  correctly in the JS port or colors/contrast will be wrong.
+- removal_model params = 226.04M confirmed. load_state_dict: all keys matched.
+- Perf: ~8.9 s/step on CPU (×19 steps + CFG ×2 = 38 UNet passes ≈ 2:48 total). WebGPU
+  expected far faster. Confirms CPU/WASM is unusable; WebGPU is the whole game.
+- Output saved to reference_out/reference_result.png — plausible inpaint. Mask convention:
+  white(255) → 1 → region to inpaint (zeroed in masked_image); black → keep.
+- num_inference_steps=20 with strength=0.99 → DDIM uses 19 steps (drops first).
+
+## Parity strategy
+- Won't try to reproduce torch RNG in JS. For PyTorch↔ONNX parity: dump identical input
+  tensors and compare outputs. For the web app: generate noise with a seedable JS RNG;
+  diffusion is robust to the particular noise draw, so visual results will be valid even
+  if not bit-identical to the torch reference.
+- DDIM `scale_model_input` is identity → skip in TS. Need to reproduce DDIM alphas/betas
+  (scaled_linear, beta 0.00085→0.012, 1000 steps) and the DDIM step update in JS.
+
+## Architecture: spatial size is FIXED (important!)
+- Self-attn (attn1): MQSλ with `r=15` → local-context path (Conv3d pos_conv). Spatially
+  dynamic, fine at any size.
+- Cross-attn (attn2): MQCλ with NO `r` → global path → `rel_pos_emb` is an
+  `nn.Parameter(n*n, m, dim_k, dim_u)` where n = per-block sample_size, m = 10. This is
+  TIED to the trained spatial resolution. Different spatial size → wrong/oob indexing.
+- ⇒ Export at STATIC 512×512 image (64×64 latent). Web app resizes user input to 512×512,
+  inpaints, resizes result back + pastes. Square only. This is the benchmark resolution.
+
+## ONNX export plan
+- Three graphs, spatial static, batch dynamic where cheap:
+  - vae_encoder: (B,3,512,512) → moments (B,8,64,64); JS uses mean=moments[:,:4]*sf.
+  - unet (RemovalModel): (B,9,64,64), timesteps(B,), input_ids(B,10) → noise(B,4,64,64).
+    Embedding (nn.Embedding 20×3072) stays IN the graph (cheap gather). CFG batches B=2.
+  - vae_decoder: (B,4,64,64) → (B,3,512,512).
+- scaling_factor = 0.13025 applied in JS (encode: latent*sf; decode: latent/sf).
+
+## Phase 2 results (ONNX export — DONE)
+- torch.onnx.export (legacy tracer, opset 18) traced all 3 graphs cleanly. No op-coverage
+  failures. The einsum/lambda/Conv3d ops all exported.
+- Parity vs PyTorch (CPU EP): decoder 5.7e-5, unet 3.6e-6, encoder mean ch ~2e-2.
+- FULL pipeline parity test (python/onnx_pipeline.py): reimplemented DDIM+CFG+9ch+scaling
+  in numpy on the ONNX sessions, vs torch models with identical noise:
+    final latents max|Δ| 0.149, decoded image mean|Δ| 0.0022, max 0.090 → visually identical.
+  This validates the ENTIRE orchestration I'll port to TS.
+- numpy DDIM vs diffusers DDIMScheduler: step max|Δ| 5e-7, timesteps identical. ✓
+
+## DDIM constants for JS (validated)
+- betas = linspace(sqrt(0.00085), sqrt(0.012), 1000)^2 ; alphas_cumprod = cumprod(1-betas)
+- timesteps(20 steps) = [950,900,...,50,0]; strength 0.99 ⇒ drop first ⇒ [900,...,0] (19).
+- ddim_step (eta=0, clip_sample=False):
+    pred_x0 = (sample - sqrt(1-ac_t)*eps) / sqrt(ac_t)
+    prev    = sqrt(ac_prev)*pred_x0 + sqrt(1-ac_prev)*eps
+    ac_prev = alphas_cumprod[prev_t], or final_alpha_cumprod=1.0 when prev_t<0 (last step).
+- noise_offset 0.0357: noise += 0.0357 * randn(B,4,1,1). (optional; small)
+
+## Web pipeline recipe (numpy → TS)
+1. resize image+mask to 512×512 (mask NEAREST, binarize ≥128).
+2. img→[-1,1] CHW; masked = img*(1-mask).
+3. encode img & masked → moments; take mean[:4] * 0.13025.
+4. mask→64×64 NEAREST, 1ch.
+5. latents = randn(1,4,64,64) [+ noise_offset].
+6. loop t in timesteps: nine=cat([latents,mask64,maskedLat]); batch×2; unet;
+   cfg = u + g*(c-u); latents = ddim_step.
+7. decode(latents/0.13025); (x+1)/2; clip; → image.
+8. paste: out*blur(mask) + (1-blur(mask))*orig.
+
+## Phase 3 (web app) — in progress
+- Vite + TS + onnxruntime-web (1.27.0). Default `onnxruntime-web/webgpu` import resolves
+  to the self-contained `ort.webgpu.bundle.min.mjs`.
+- Models served LOCALLY: web/public/models -> ../models symlink, at /models/*.onnx.
+  Total ~1.24GB fetched over localhost (no internet). ORT runtime served from
+  web/ort-dist at /ort/* via a custom static middleware (see vite.config.ts).
+- BUG FIXED: ORT glue .mjs must NOT be in /public (Vite tries to module-transform it).
+  Fix = serve /ort/* as raw static files via configureServer middleware.
+- COOP/COEP headers set (needed for threaded WASM / SharedArrayBuffer).
+
+## WebGPU op coverage (confirmed from ORT source at /tmp/Moebius/onnxruntime)
+- js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts registers: Einsum ✓, Conv ✓
+  (conv.ts has computeConv3DInfo / createConv3DNaiveProgramInfo → Conv3d pos_conv works,
+  naive kernel so possibly slow), InstanceNormalization ✓, MatMul/Gemm ✓, Softmax ✓,
+  Reduce* ✓, Transpose/Concat/Gather/Pad/Resize/Where ✓.
+- GroupNorm: not registered by that name, BUT torch.onnx exports nn.GroupNorm as a
+  Reshape→InstanceNormalization→Reshape→Mul→Add decomposition → covered. (VAE decoder
+  CPU-EP parity was 5.7e-5, so the graph is decomposed, not a single GroupNorm op.)
+- ⇒ No expected silent CPU fallback for the heavy ops. Confirm empirically in console.
+
+## Verification without a GPU browser (sandbox can't drive Chrome — user's live Chrome
+## holds the playwright profile)
+- web/test/fixture/*.bin: dumped inputs + reference final latents from the validated
+  numpy/ONNX pipeline, to check the TS port (ddim.ts + 9ch assembly) in Node.
+
+## TODO / unknowns
+- fp16 export to ~450MB UNet for real deployment (VAE fp16 unstable: decoder Cast issue;
+  keep VAE fp32). Quality risk in λ layers — validate before shipping fp16.
+- Confirm in-browser: WebGPU selected, no CPU fallback, end-to-end correctness + timing.

unet.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:adfb872fead9f7fa8750ecfb6a7acea836c8f96e0cf1d01d9786dcec63d70f4e
+size 906555486

vae_decoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d90ef0b7f6c8c8b7234459c8b449d70be0033bf1576c842e8b9991baf3934280
+size 198078671

vae_encoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8b81d41e757222a0707665ba9d826703987855e5bed056036b90b988968042f
+size 136757093