mlboydaisuke commited on
Commit
3e1f4f7
·
verified ·
1 Parent(s): 45432df

TripoSplat -> Core AI: image->3D Gaussian splats; 6-net fp16/fp32 pipeline + recipe

Browse files
.gitattributes CHANGED
@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
 
 
 
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ decode_fp32.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text
37
+ dinov3_fp16.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text
38
+ dit_fp16.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text
39
+ gs_fp16.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text
40
+ octree_fp32.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text
41
+ vae_fp16.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # TripoSplat → Core AI (zoo's first 3D)
2
+
3
+ [VAST-AI/TripoSplat](https://github.com/VAST-AI-Research/TripoSplat) — **single image → 3D Gaussian
4
+ splats** (`.ply`/`.splat`), MIT. The zoo's first 3D model: outputs drop straight into a Gaussian-splat
5
+ viewer (e.g. Apple RealityKit on visionOS, or MetalSplatter on iOS/macOS).
6
+
7
+ Pure-PyTorch pipeline (no diffusers/CUDA kernels): bg-removal → DINOv3 ViT-H encode + Flux2-VAE encode
8
+ → 20-step flow-matching DiT denoiser → octree probability sampler → Gaussian decoder → splats.
9
+
10
+ ## What runs on Core AI
11
+
12
+ 5 neural nets converted (each gated converted-vs-eager **cos = 1.000000**, fp16):
13
+
14
+ | net | shape | bundle |
15
+ |---|---|---|
16
+ | DINOv3 ViT-H encoder | (1,3,1024,1024)→(1,4101,1280) | `dinov3` |
17
+ | Flux2-VAE encoder | (1,3,1024,1024)→(1,4096,128) | `vae` |
18
+ | DiT denoiser (one step) | latent(1,8192,16)+cam(1,1,5)+t+feat1(1,4101,1280)+feat2(1,4101,128)→latent,cam | `dit` |
19
+ | Octree probability decoder | x(1,8192,3)+l(1,)+cond(1,8192,16)→logits(1,8192,8) | `octree` |
20
+ | Decode (gs + build_gaussians + .ply activations, baked) | points(1,8192,3)+cond(1,8192,16)→(262144,14) | `decode` |
21
+
22
+ The flow-matching sampler (`FlowEulerCfgSampler`) and the octree `sample_probs` systematic resampling
23
+ stay **host-side** (data-dependent control flow). Scripts: `_conv_*.py` convert+gate each net;
24
+ `_conv_fp16.py` makes the half-size fp16 bundles; `_conv_decode.py` bakes build_gaussians + the
25
+ Gaussian `.ply`-activation math into one net so the runner just writes raw floats.
26
+
27
+ ## model.py patches (the reusable contribution — see ../../knowledge/conversion-guide.md)
28
+
29
+ coreai-torch 0.4.0 needed six edits to VAST's `model.py`; all are general gotchas:
30
+
31
+ 1. **float-arg `aten.arange`** → `bad_optional_access` C++ abort. Use int-arg arange (DINOv3 RoPE).
32
+ 2. **fx `got multiple values for 'mod'`** — submodule called with `mod=` kwarg. Pass positionally.
33
+ 3. **No complex ops** — rewrote the DiT's complex RoPE (`torch.polar`/`view_as_complex`) as real
34
+ cos/sin math (`apply_rotary_emb`, `RePo3DRotaryEmbedding.forward`).
35
+ 4. **Constant-folded `sin/cos` of huge args is low-precision** (cos→0.5) — the DiT positional embed
36
+ computed from the fixed Sobol constant was folded wrong; precompute it into a `register_buffer`.
37
+ 5. **`F.normalize` drops the eps clamp** → near-zero vectors blow up ~1e13; rewrote `MultiHeadRMSNorm`
38
+ as explicit `x*rsqrt(mean(x²)+eps)`. (Emergent only at large seq len — gate by VISUAL/true-scale.)
39
+ 6. **`prog.optimize()` hangs** on the 24-block/12k-token DiT graph (>90 min) — skip it
40
+ (`convert(optimize=False)`), AOT `coreai-build` optimizes for the device anyway.
41
+
42
+ Plus: int8 desaturates this model (per-net cos 0.9998 but colors collapse → use **fp16**, which is
43
+ GPU-identical to fp32 — gate fp16 on GPU/visual, its CPU cos looks bad but that's a CPU-compute
44
+ artifact). Octree decoder: int64 `l` (resolution) input → CoreAIError 3 at runtime, pass it as float32.
45
+
46
+ ## Running it
47
+
48
+ - **Mac**: `_run_coreai.py` (or `app_backend.py --input <img>`) loads the bundles via coreai.runtime
49
+ (`SpecializationOptions.default()` = GPU; ~1 min/gen at 20 steps, full quality). End-to-end latent
50
+ gate vs torch-DiT: **cos 0.999999**.
51
+ - **Mac app / iPhone client**: `apps/` — `TripoSplatMac` (standalone) and `TripoSplatPhone`
52
+ (capture on iPhone → Mac server `server.py` → view splats in MetalSplatter / RealityKit).
53
+
54
+ ## On-device note
55
+
56
+ Full on-device (iPhone) was **verified infeasible** with this model: DINOv3 ViT-H AOT `.aimodelc`
57
+ is ~3.1 GB and the DiT's 12294-token full-attention score matrix alone is ~4.8 GB, both over the
58
+ ~3.3 GB iOS app memory budget (weight precision doesn't fix the attention working set). Needs
59
+ flash-attention conversion / weight streaming. The Mac-link client is the shipped path.
decode_fp32.aimodel/main.hash ADDED
@@ -0,0 +1 @@
 
 
1
+ C��{_���4NS�S*�evZ>�{���&B
decode_fp32.aimodel/main.mlirb ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:43f3e57b5ff3e37fca344e5303d0532a08e8bd65765a3ea77bc8e0d71c261742
3
+ size 949402894
decode_fp32.aimodel/metadata.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ {
2
+ "assetVersion" : "2.0"
3
+ }
dinov3_fp16.aimodel/main.hash ADDED
@@ -0,0 +1 @@
 
 
1
+ inXP�ܩJ�\�������J��qT��6 �Q�
dinov3_fp16.aimodel/main.mlirb ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:696e58509edca94afe5cb5ac841ae413f0f58b4ae8fe107154a9b73609a751bd
3
+ size 1681455018
dinov3_fp16.aimodel/metadata.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ {
2
+ "assetVersion" : "2.0"
3
+ }
dit_fp16.aimodel/main.hash ADDED
@@ -0,0 +1 @@
 
 
1
+ �q �^o8^YH�n�g��g�V�����\��G��
dit_fp16.aimodel/main.mlirb ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fc710bb65e6f1e385e59488e6ea2678bbe67a656f4adf7dddc5c91b647b9e79c
3
+ size 758927980
dit_fp16.aimodel/metadata.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ {
2
+ "assetVersion" : "2.0"
3
+ }
gs_fp16.aimodel/main.hash ADDED
@@ -0,0 +1 @@
 
 
1
+ &��贜�+%>|� o�*w�H�B����z
gs_fp16.aimodel/main.mlirb ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:26ffaae8b49ca31d2b253e7c9f206ff1a42a778e1c4807ed42b0860fc6cb1a7a
3
+ size 475102303
gs_fp16.aimodel/metadata.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ {
2
+ "assetVersion" : "2.0"
3
+ }
octree_fp32.aimodel/main.hash ADDED
@@ -0,0 +1 @@
 
 
1
+ e}ž��$v����4M��;��BC�V��a��
octree_fp32.aimodel/main.mlirb ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:657dc29ebbb92476e2d6e6cf344d8aa53bcce291424319c056c0f09161eb15a6
3
+ size 203236017
octree_fp32.aimodel/metadata.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ {
2
+ "assetVersion" : "2.0"
3
+ }
vae_fp16.aimodel/main.hash ADDED
@@ -0,0 +1 @@
 
 
1
+ � K$ k�6,TA-�kh�y(@a�衉&N�I
vae_fp16.aimodel/main.mlirb ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fc094b24206beb8f36042c54412dd66b68de7928401961a6e8a18916264ee649
3
+ size 68919930
vae_fp16.aimodel/metadata.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ {
2
+ "assetVersion" : "2.0"
3
+ }