diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..b85d558a3db6d00f55f9ba052bb3cb3a65699fb6 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,4 @@ 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 +den_pos_emb.f32 filter=lfs diff=lfs merge=lfs -text diff --git a/LICENSES.md b/LICENSES.md new file mode 100644 index 0000000000000000000000000000000000000000..faa1f2aa03c8183b6ee3867537ea4e58ac982161 --- /dev/null +++ b/LICENSES.md @@ -0,0 +1,37 @@ +# Licences of the parts of the Idol pipeline (read 2026-11-09) + +| Part | Weights used here | Licence | Source read | +|---|---|---|---| +| Denoiser, octree decoder, Gaussian decoder | speridlabs/idol @2731c229 (`denoiser`, `decoder`, `decoder_gs` .safetensors) | Apache-2.0 | HF card front matter line 2 `license: apache-2.0`; idol-sdk LICENSE @4702a7f0 (Apache License 2.0, lines 1-2) | +| Turbo LoRA (slow / fast experts) | speridlabs/idol-turbo @45cdd41f (`idol_turbo_lora.safetensors`) | Apache-2.0 | HF card front matter `license: apache-2.0`; card §License "Apache-2.0, inherited from speridlabs/idol" | +| Background removal (shipped: `bir.tflite`, LiteRT CPU) | speridlabs/idol @2731c229 `birefnet.safetensors` (BiRefNet, Swin-L) | MIT (BiRefNet); the idol repo that packages it is Apache-2.0 | github ZhengPeng7/BiRefNet licence = MIT (GitHub API); idol HF card front matter line 2 | +| Background removal, alternative in the app (off by default) | litert-community/DIS-ISNet-LiteRT @1b966dbe `dis.tflite` (IS-Net DIS general-use, xuebinqin/DIS) | Apache-2.0 | litert-community/DIS-ISNet-LiteRT card `license: apache-2.0` (Hub API) | +| FLUX.2 VAE (encoder side used) | black-forest-labs/FLUX.2-klein-4B @e7b7dc27 `vae/` (ungated) | Apache-2.0 | black-forest-labs/flux2 README @50fe5162, line 111: "The autoencoder is released under Apache 2.0"; klein-4B card `license: apache-2.0` | +| DINOv3 ViT-H+/16 | Comfy-Org/Idol @e907d8dc `clip_vision/dino_v3_vit_h_bf16.safetensors` (bf16 repack of facebook/dinov3-vith16plus-pretrain-lvd1689m, gated) | DINOv3 License (Meta, "Last Updated: August 19, 2025") | facebookresearch/dinov3 LICENSE.md @6876159a | +| Idol acknowledges TripoSplat | — (no TripoSplat weights used) | MIT | github VAST-AI-Research/TripoSplat licence = MIT (GitHub API) | + +## DINOv3 License — the clauses that decide redistribution (dinov3_LICENSE.md) + +- Line 20: "By clicking "I Accept" below or by using or distributing any portion or element of the DINO Materials, you agree to be bound by this Agreement." → using the Comfy-Org repack binds the user of it to the Agreement as if accepted on the gated page. +- Line 24 (§1.a Grant of Rights): "non-exclusive, worldwide, non-transferable and royalty-free limited license … to use, reproduce, distribute, copy, create derivative works of, and make modifications to the DINO Materials." → a converted .tflite of the encoder is a derivative work and may be distributed. +- Line 28 (§1.b.i): "If you distribute or make the DINO Materials, or any derivative works thereof, available to a third party, you may only do so under the terms of this Agreement and you shall provide a copy of this Agreement with any such DINO Materials." → a hosted encoder file must carry the DINOv3 License text and be offered under it (not under Apache-2.0). +- Line 30 (§1.b.ii): acknowledgement in research publications. +- Line 34 (§1.b.iv): no use that involves or encourages reverse engineering / decompiling the DINO Materials. +- Line 36 (§1.b.v): trade-control limits (no ITAR / military / nuclear / espionage / weapons end uses; licensee not a target of trade controls). +- Line 52 (§5.a): derivative works made by the licensee are owned by the licensee (subject to Meta's ownership of the DINO Materials). +- No non-commercial clause anywhere in the text (66 lines read in full). + +## FLUX.2 VAE + +- flux2 README line 111 (commit 50fe5162): the FLUX.2 autoencoder is released under Apache 2.0. The FLUX.2 [dev] Non-Commercial License of the FLUX.2-dev repo covers the [dev] transformer, not the autoencoder. +- The klein-4B `vae/` dir is ungated and tagged apache-2.0; the Idol pipeline loads the VAE in bf16 (`conditioning.py` `torch_dtype=torch.bfloat16`) and klein-4B's file is bf16. +- Verified 2026-11-09 14:04: klein-4B `vae/diffusion_pytorch_model.safetensors` (sha256 ca70d220…, bf16) = the FLUX.2-dev VAE cast to bf16. The FLUX.2-dev VAE was read from Comfy-Org/Idol `vae/flux2-vae.safetensors`, whose LFS sha256 868fe7b3… equals FLUX.2-dev's `ae.safetensors` (Hub API, FLUX.2-dev @26afe3a7): the same file. All 250 float tensors match bit for bit after the bf16 cast (84,046,372 values; 8 attention projections are [512, 512] in klein and [512, 512, 1, 1] in BFL, same values). The one other tensor differs: `bn.num_batches_tracked` (int64 training counter, 0 vs 400000), unused at inference. + +## Hosting + +- litert-community/Idol (ungated, Apache-2.0, with the per-part provenance above): the denoiser, octree and Gaussian + decoder graphs, the VAE encoder graphs, BiRefNet (`bir.tflite`, MIT part named in the card), the side files, the + manifest, and the Python / Android sources. +- The DINOv3 encoder graphs and their RoPE tables: a separate litert-community repo, gated (automatic approval), carrying + the full DINOv3 License text and the provenance (base_model facebook/dinov3-vith16plus-pretrain-lvd1689m), the way + litert-community/Llama-3.2-1B carries the Llama licence. diff --git a/android/app/build.gradle.kts b/android/app/build.gradle.kts new file mode 100644 index 0000000000000000000000000000000000000000..6a62e1396899146913ebcba1defdc9b9dd50bd2f --- /dev/null +++ b/android/app/build.gradle.kts @@ -0,0 +1,67 @@ +plugins { + alias(libs.plugins.android.application) + alias(libs.plugins.kotlin.android) +} + +android { + namespace = "com.idol" + compileSdk = 35 + buildToolsVersion = "35.0.0" + + defaultConfig { + applicationId = "com.idol" + minSdk = 26 + targetSdk = 35 + versionCode = 1 + versionName = "0.1" + buildConfigField("String", "LITERT_VERSION", "\"${libs.versions.litert.get()}\"") + ndk { + abiFilters += "arm64-v8a" + } + } + + buildTypes { + release { + isMinifyEnabled = false + // sample convenience: installable without a release keystore + signingConfig = signingConfigs.getByName("debug") + } + } + + androidResources { + // graphs are never assets here (they are installed in files/models/), but keep any .tflite mmappable + noCompress += "tflite" + } + + packaging { + // The GPU accelerator library (litert-gpu) is loaded from the extracted native library directory, as in the + // zoo's other GPU samples on the same phones. + jniLibs { useLegacyPackaging = true } + } + + compileOptions { + sourceCompatibility = JavaVersion.VERSION_17 + targetCompatibility = JavaVersion.VERSION_17 + } + buildFeatures { + buildConfig = true + } +} + +kotlin { + compilerOptions { + jvmTarget.set(org.jetbrains.kotlin.gradle.dsl.JvmTarget.JVM_17) + // LiteRT 2.3.0's Kotlin API carries Kotlin 2.4 metadata; Google's own 2.3.0 samples on Kotlin 2.2.21 pass this + freeCompilerArgs.add("-Xskip-metadata-version-check") + } +} + +dependencies { + // LiteRT CompiledModel (CPU / XNNPACK for BiRefNet) and its GPU accelerator (ML Drift, OpenCL) for the rest + implementation(libs.litert) + implementation(libs.litert.gpu) + implementation(libs.androidx.core.ktx) + implementation(libs.androidx.activity.ktx) + // WebViewAssetLoader: the splat viewer and the .ply served to the WebView from the app itself + implementation(libs.androidx.webkit) +} diff --git a/android/app/src/main/AndroidManifest.xml b/android/app/src/main/AndroidManifest.xml new file mode 100644 index 0000000000000000000000000000000000000000..8d70f10ac8a39e1aba337169ff3ea4b7e0f3e225 --- /dev/null +++ b/android/app/src/main/AndroidManifest.xml @@ -0,0 +1,45 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/android/app/src/main/assets/viewer/LICENSE b/android/app/src/main/assets/viewer/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..39a8dff394d8af03f2463f709b95cf206e240c0c --- /dev/null +++ b/android/app/src/main/assets/viewer/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 Kevin Kwok + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/android/app/src/main/assets/viewer/index.html b/android/app/src/main/assets/viewer/index.html new file mode 100644 index 0000000000000000000000000000000000000000..045bd2c119c3521390a914f20dd7de0da93f426d --- /dev/null +++ b/android/app/src/main/assets/viewer/index.html @@ -0,0 +1,281 @@ + + + + WebGL Gaussian Splat Viewer + + + + + + + + +
+

WebGL 3D Gaussian Splat Viewer

+

+ + By Kevin Kwok. + Code on + Github. + +

+ +
+ Use mouse or arrow keys to navigate. + +
movement (arrow keys) +- left/right arrow keys to strafe side to side +- up/down arrow keys to move forward/back +- space to jump + +camera angle (wasd) +- a/d to turn camera left/right +- w/s to tilt camera up/down +- q/e to roll camera counterclockwise/clockwise +- i/k and j/l to orbit + +trackpad +- scroll up/down/left/right to orbit +- pinch to move forward/back +- ctrl key + scroll to move forward/back +- shift + scroll to move up/down or strafe + +mouse +- click and drag to orbit +- right click (or ctrl/cmd key) and drag up/down to move + +touch (mobile) +- one finger to orbit +- two finger pinch to move forward/back +- two finger rotate to rotate camera clockwise/counterclockwise +- two finger pan to move side-to-side and up-down + +gamepad +- if you have a game controller connected it should work + +other +- press 0-9 to switch to one of the pre-loaded camera views +- press '-' or '+'key to cycle loaded cameras +- press p to resume default animation +- drag and drop .ply file to convert to .splat +- drag and drop cameras.json to load cameras +
+ +
+ +
+ +
+ +
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ + +
+ +
+
+ +
+ + + diff --git a/android/app/src/main/assets/viewer/main.js b/android/app/src/main/assets/viewer/main.js new file mode 100644 index 0000000000000000000000000000000000000000..8235c8d1cd85ba65fb0647b6c9b4134574eb0854 --- /dev/null +++ b/android/app/src/main/assets/viewer/main.js @@ -0,0 +1,1527 @@ +let cameras = [ + { + id: 0, + img_name: "00001", + width: 1959, + height: 1090, + position: [ + -3.0089893469241797, -0.11086489695181866, -3.7527640949141428, + ], + rotation: [ + [0.876134201218856, 0.06925962026449776, 0.47706599800804744], + [-0.04747421839895102, 0.9972110940209488, -0.057586739349882114], + [-0.4797239414934443, 0.027805376500959853, 0.8769787916452908], + ], + fy: 1164.6601287484507, + fx: 1159.5880733038064, + }, + { + id: 1, + img_name: "00009", + width: 1959, + height: 1090, + position: [ + -2.5199776022057296, -0.09704735754873686, -3.6247725540304545, + ], + rotation: [ + [0.9982731285632193, -0.011928707708098955, -0.05751927260507243], + [0.0065061360949636325, 0.9955928229282383, -0.09355533724430458], + [0.058381769258182864, 0.09301955098900708, 0.9939511719154457], + ], + fy: 1164.6601287484507, + fx: 1159.5880733038064, + }, + { + id: 2, + img_name: "00017", + width: 1959, + height: 1090, + position: [ + -0.7737533667465242, -0.3364271945329695, -2.9358969417573753, + ], + rotation: [ + [0.9998813418672372, 0.013742375651625236, -0.0069605529394208224], + [-0.014268370388586709, 0.996512943252834, -0.08220929105659476], + [0.00580653013657589, 0.08229885200307129, 0.9965907801935302], + ], + fy: 1164.6601287484507, + fx: 1159.5880733038064, + }, + { + id: 3, + img_name: "00025", + width: 1959, + height: 1090, + position: [ + 1.2198221749590001, -0.2196687861401182, -2.3183162007028453, + ], + rotation: [ + [0.9208648867765482, 0.0012010625395201253, 0.389880004297208], + [-0.06298204172269357, 0.987319521752825, 0.14571693239364383], + [-0.3847611242348369, -0.1587410451475895, 0.9092635249821667], + ], + fy: 1164.6601287484507, + fx: 1159.5880733038064, + }, + { + id: 4, + img_name: "00033", + width: 1959, + height: 1090, + position: [ + 1.742387858893817, -0.13848225198886954, -2.0566370113193146, + ], + rotation: [ + [0.24669889292141334, -0.08370189346592856, -0.9654706879349405], + [0.11343747891376445, 0.9919082664242816, -0.05700815184573074], + [0.9624300466054861, -0.09545671285663988, 0.2541976029815521], + ], + fy: 1164.6601287484507, + fx: 1159.5880733038064, + }, + { + id: 5, + img_name: "00041", + width: 1959, + height: 1090, + position: [ + 3.6567309419223935, -0.16470990600750707, -1.3458085590422042, + ], + rotation: [ + [0.2341293058324528, -0.02968330457755884, -0.9717522161434825], + [0.10270823606832301, 0.99469554638321, -0.005638106875665722], + [0.9667649592295676, -0.09848690996657204, 0.2359360976431732], + ], + fy: 1164.6601287484507, + fx: 1159.5880733038064, + }, + { + id: 6, + img_name: "00049", + width: 1959, + height: 1090, + position: [ + 3.9013554243203497, -0.2597500978038105, -0.8106154188297828, + ], + rotation: [ + [0.6717235545638952, -0.015718162115524837, -0.7406351366386528], + [0.055627354673906296, 0.9980224478387622, 0.029270992841185218], + [0.7387104058127439, -0.060861588786650656, 0.6712695459756353], + ], + fy: 1164.6601287484507, + fx: 1159.5880733038064, + }, + { + id: 7, + img_name: "00057", + width: 1959, + height: 1090, + position: [4.742994605467533, -0.05591660945412069, 0.9500365976084458], + rotation: [ + [-0.17042655709210375, 0.01207080756938, -0.9852964448542146], + [0.1165090336695526, 0.9931575292530063, -0.00798543433078162], + [0.9784581921120181, -0.1161568667478904, -0.1706667764862097], + ], + fy: 1164.6601287484507, + fx: 1159.5880733038064, + }, + { + id: 8, + img_name: "00065", + width: 1959, + height: 1090, + position: [4.34676307626522, 0.08168160516967145, 1.0876221470355405], + rotation: [ + [-0.003575447631888379, -0.044792503246552894, -0.9989899137764799], + [0.10770152645126597, 0.9931680875192705, -0.04491693593046672], + [0.9941768441149182, -0.10775333677534978, 0.0012732004866391048], + ], + fy: 1164.6601287484507, + fx: 1159.5880733038064, + }, + { + id: 9, + img_name: "00073", + width: 1959, + height: 1090, + position: [3.264984351114202, 0.078974937336732, 1.0117200284114904], + rotation: [ + [-0.026919994628162257, -0.1565891128261527, -0.9872968974090509], + [0.08444552208239385, 0.983768234577625, -0.1583319754069128], + [0.9960643893290491, -0.0876350978794554, -0.013259786205163005], + ], + fy: 1164.6601287484507, + fx: 1159.5880733038064, + }, +]; + +let camera = cameras[0]; + +// idol: the camera of the contract renders (idol/scripts/render_views.py): it looks at the origin from distance 2.2, +// 20 degrees above the object, image up = -y (idol's PLY is y-down), field of view 40 degrees; the orbit turns it +// about the vertical axis. Azimuth 90 degrees faces the object's front (scripts/make_hero_app.py). +const IDOL_DISTANCE = 2.2; +const IDOL_ELEVATION = (20 * Math.PI) / 180; +const IDOL_FOV = (40 * Math.PI) / 180; +const IDOL_START_AZIMUTH = Math.PI / 2; +const IDOL_TURN_MS = 12000; + +// idol: world -> camera (rows x, y, z = right, down, forward; t = -R c), column-major like getViewMatrix's result +function idolViewMatrix(azimuth) { + const ce = Math.cos(IDOL_ELEVATION); + const se = Math.sin(IDOL_ELEVATION); + const c = [ + IDOL_DISTANCE * ce * Math.sin(azimuth), + -IDOL_DISTANCE * se, + IDOL_DISTANCE * ce * Math.cos(azimuth), + ]; + const n = Math.hypot(c[0], c[1], c[2]); + const z = [-c[0] / n, -c[1] / n, -c[2] / n]; + const up = [0, -1, 0]; + let x = [ + z[1] * up[2] - z[2] * up[1], + z[2] * up[0] - z[0] * up[2], + z[0] * up[1] - z[1] * up[0], + ]; + const xn = Math.hypot(x[0], x[1], x[2]); + x = [x[0] / xn, x[1] / xn, x[2] / xn]; + const y = [ + z[1] * x[2] - z[2] * x[1], + z[2] * x[0] - z[0] * x[2], + z[0] * x[1] - z[1] * x[0], + ]; + const t = [ + -(x[0] * c[0] + x[1] * c[1] + x[2] * c[2]), + -(y[0] * c[0] + y[1] * c[1] + y[2] * c[2]), + -(z[0] * c[0] + z[1] * c[1] + z[2] * c[2]), + ]; + return [x[0], y[0], z[0], 0, x[1], y[1], z[1], 0, x[2], y[2], z[2], 0, t[0], t[1], t[2], 1]; +} + +function getProjectionMatrix(fx, fy, width, height) { + const znear = 0.2; + const zfar = 200; + return [ + [(2 * fx) / width, 0, 0, 0], + [0, -(2 * fy) / height, 0, 0], + [0, 0, zfar / (zfar - znear), 1], + [0, 0, -(zfar * znear) / (zfar - znear), 0], + ].flat(); +} + +function getViewMatrix(camera) { + const R = camera.rotation.flat(); + const t = camera.position; + const camToWorld = [ + [R[0], R[1], R[2], 0], + [R[3], R[4], R[5], 0], + [R[6], R[7], R[8], 0], + [ + -t[0] * R[0] - t[1] * R[3] - t[2] * R[6], + -t[0] * R[1] - t[1] * R[4] - t[2] * R[7], + -t[0] * R[2] - t[1] * R[5] - t[2] * R[8], + 1, + ], + ].flat(); + return camToWorld; +} +// function translate4(a, x, y, z) { +// return [ +// ...a.slice(0, 12), +// a[0] * x + a[4] * y + a[8] * z + a[12], +// a[1] * x + a[5] * y + a[9] * z + a[13], +// a[2] * x + a[6] * y + a[10] * z + a[14], +// a[3] * x + a[7] * y + a[11] * z + a[15], +// ]; +// } + +function multiply4(a, b) { + return [ + b[0] * a[0] + b[1] * a[4] + b[2] * a[8] + b[3] * a[12], + b[0] * a[1] + b[1] * a[5] + b[2] * a[9] + b[3] * a[13], + b[0] * a[2] + b[1] * a[6] + b[2] * a[10] + b[3] * a[14], + b[0] * a[3] + b[1] * a[7] + b[2] * a[11] + b[3] * a[15], + b[4] * a[0] + b[5] * a[4] + b[6] * a[8] + b[7] * a[12], + b[4] * a[1] + b[5] * a[5] + b[6] * a[9] + b[7] * a[13], + b[4] * a[2] + b[5] * a[6] + b[6] * a[10] + b[7] * a[14], + b[4] * a[3] + b[5] * a[7] + b[6] * a[11] + b[7] * a[15], + b[8] * a[0] + b[9] * a[4] + b[10] * a[8] + b[11] * a[12], + b[8] * a[1] + b[9] * a[5] + b[10] * a[9] + b[11] * a[13], + b[8] * a[2] + b[9] * a[6] + b[10] * a[10] + b[11] * a[14], + b[8] * a[3] + b[9] * a[7] + b[10] * a[11] + b[11] * a[15], + b[12] * a[0] + b[13] * a[4] + b[14] * a[8] + b[15] * a[12], + b[12] * a[1] + b[13] * a[5] + b[14] * a[9] + b[15] * a[13], + b[12] * a[2] + b[13] * a[6] + b[14] * a[10] + b[15] * a[14], + b[12] * a[3] + b[13] * a[7] + b[14] * a[11] + b[15] * a[15], + ]; +} + +function invert4(a) { + let b00 = a[0] * a[5] - a[1] * a[4]; + let b01 = a[0] * a[6] - a[2] * a[4]; + let b02 = a[0] * a[7] - a[3] * a[4]; + let b03 = a[1] * a[6] - a[2] * a[5]; + let b04 = a[1] * a[7] - a[3] * a[5]; + let b05 = a[2] * a[7] - a[3] * a[6]; + let b06 = a[8] * a[13] - a[9] * a[12]; + let b07 = a[8] * a[14] - a[10] * a[12]; + let b08 = a[8] * a[15] - a[11] * a[12]; + let b09 = a[9] * a[14] - a[10] * a[13]; + let b10 = a[9] * a[15] - a[11] * a[13]; + let b11 = a[10] * a[15] - a[11] * a[14]; + let det = + b00 * b11 - b01 * b10 + b02 * b09 + b03 * b08 - b04 * b07 + b05 * b06; + if (!det) return null; + return [ + (a[5] * b11 - a[6] * b10 + a[7] * b09) / det, + (a[2] * b10 - a[1] * b11 - a[3] * b09) / det, + (a[13] * b05 - a[14] * b04 + a[15] * b03) / det, + (a[10] * b04 - a[9] * b05 - a[11] * b03) / det, + (a[6] * b08 - a[4] * b11 - a[7] * b07) / det, + (a[0] * b11 - a[2] * b08 + a[3] * b07) / det, + (a[14] * b02 - a[12] * b05 - a[15] * b01) / det, + (a[8] * b05 - a[10] * b02 + a[11] * b01) / det, + (a[4] * b10 - a[5] * b08 + a[7] * b06) / det, + (a[1] * b08 - a[0] * b10 - a[3] * b06) / det, + (a[12] * b04 - a[13] * b02 + a[15] * b00) / det, + (a[9] * b02 - a[8] * b04 - a[11] * b00) / det, + (a[5] * b07 - a[4] * b09 - a[6] * b06) / det, + (a[0] * b09 - a[1] * b07 + a[2] * b06) / det, + (a[13] * b01 - a[12] * b03 - a[14] * b00) / det, + (a[8] * b03 - a[9] * b01 + a[10] * b00) / det, + ]; +} + +function rotate4(a, rad, x, y, z) { + let len = Math.hypot(x, y, z); + x /= len; + y /= len; + z /= len; + let s = Math.sin(rad); + let c = Math.cos(rad); + let t = 1 - c; + let b00 = x * x * t + c; + let b01 = y * x * t + z * s; + let b02 = z * x * t - y * s; + let b10 = x * y * t - z * s; + let b11 = y * y * t + c; + let b12 = z * y * t + x * s; + let b20 = x * z * t + y * s; + let b21 = y * z * t - x * s; + let b22 = z * z * t + c; + return [ + a[0] * b00 + a[4] * b01 + a[8] * b02, + a[1] * b00 + a[5] * b01 + a[9] * b02, + a[2] * b00 + a[6] * b01 + a[10] * b02, + a[3] * b00 + a[7] * b01 + a[11] * b02, + a[0] * b10 + a[4] * b11 + a[8] * b12, + a[1] * b10 + a[5] * b11 + a[9] * b12, + a[2] * b10 + a[6] * b11 + a[10] * b12, + a[3] * b10 + a[7] * b11 + a[11] * b12, + a[0] * b20 + a[4] * b21 + a[8] * b22, + a[1] * b20 + a[5] * b21 + a[9] * b22, + a[2] * b20 + a[6] * b21 + a[10] * b22, + a[3] * b20 + a[7] * b21 + a[11] * b22, + ...a.slice(12, 16), + ]; +} + +function translate4(a, x, y, z) { + return [ + ...a.slice(0, 12), + a[0] * x + a[4] * y + a[8] * z + a[12], + a[1] * x + a[5] * y + a[9] * z + a[13], + a[2] * x + a[6] * y + a[10] * z + a[14], + a[3] * x + a[7] * y + a[11] * z + a[15], + ]; +} + +function createWorker(self) { + let buffer; + let vertexCount = 0; + let viewProj; + // 6*4 + 4 + 4 = 8*4 + // XYZ - Position (Float32) + // XYZ - Scale (Float32) + // RGBA - colors (uint8) + // IJKL - quaternion/rot (uint8) + const rowLength = 3 * 4 + 3 * 4 + 4 + 4; + let lastProj = []; + let depthIndex = new Uint32Array(); + let lastVertexCount = 0; + + var _floatView = new Float32Array(1); + var _int32View = new Int32Array(_floatView.buffer); + + function floatToHalf(float) { + _floatView[0] = float; + var f = _int32View[0]; + + var sign = (f >> 31) & 0x0001; + var exp = (f >> 23) & 0x00ff; + var frac = f & 0x007fffff; + + var newExp; + if (exp == 0) { + newExp = 0; + } else if (exp < 113) { + newExp = 0; + frac |= 0x00800000; + frac = frac >> (113 - exp); + if (frac & 0x01000000) { + newExp = 1; + frac = 0; + } + } else if (exp < 142) { + newExp = exp - 112; + } else { + newExp = 31; + frac = 0; + } + + return (sign << 15) | (newExp << 10) | (frac >> 13); + } + + function packHalf2x16(x, y) { + return (floatToHalf(x) | (floatToHalf(y) << 16)) >>> 0; + } + + function generateTexture() { + if (!buffer) return; + const f_buffer = new Float32Array(buffer); + const u_buffer = new Uint8Array(buffer); + + var texwidth = 1024 * 2; // Set to your desired width + var texheight = Math.ceil((2 * vertexCount) / texwidth); // Set to your desired height + var texdata = new Uint32Array(texwidth * texheight * 4); // 4 components per pixel (RGBA) + var texdata_c = new Uint8Array(texdata.buffer); + var texdata_f = new Float32Array(texdata.buffer); + + // Here we convert from a .splat file buffer into a texture + // With a little bit more foresight perhaps this texture file + // should have been the native format as it'd be very easy to + // load it into webgl. + for (let i = 0; i < vertexCount; i++) { + // x, y, z + texdata_f[8 * i + 0] = f_buffer[8 * i + 0]; + texdata_f[8 * i + 1] = f_buffer[8 * i + 1]; + texdata_f[8 * i + 2] = f_buffer[8 * i + 2]; + + // r, g, b, a + texdata_c[4 * (8 * i + 7) + 0] = u_buffer[32 * i + 24 + 0]; + texdata_c[4 * (8 * i + 7) + 1] = u_buffer[32 * i + 24 + 1]; + texdata_c[4 * (8 * i + 7) + 2] = u_buffer[32 * i + 24 + 2]; + texdata_c[4 * (8 * i + 7) + 3] = u_buffer[32 * i + 24 + 3]; + + // quaternions + let scale = [ + f_buffer[8 * i + 3 + 0], + f_buffer[8 * i + 3 + 1], + f_buffer[8 * i + 3 + 2], + ]; + let rot = [ + (u_buffer[32 * i + 28 + 0] - 128) / 128, + (u_buffer[32 * i + 28 + 1] - 128) / 128, + (u_buffer[32 * i + 28 + 2] - 128) / 128, + (u_buffer[32 * i + 28 + 3] - 128) / 128, + ]; + + // Compute the matrix product of S and R (M = S * R) + const M = [ + 1.0 - 2.0 * (rot[2] * rot[2] + rot[3] * rot[3]), + 2.0 * (rot[1] * rot[2] + rot[0] * rot[3]), + 2.0 * (rot[1] * rot[3] - rot[0] * rot[2]), + + 2.0 * (rot[1] * rot[2] - rot[0] * rot[3]), + 1.0 - 2.0 * (rot[1] * rot[1] + rot[3] * rot[3]), + 2.0 * (rot[2] * rot[3] + rot[0] * rot[1]), + + 2.0 * (rot[1] * rot[3] + rot[0] * rot[2]), + 2.0 * (rot[2] * rot[3] - rot[0] * rot[1]), + 1.0 - 2.0 * (rot[1] * rot[1] + rot[2] * rot[2]), + ].map((k, i) => k * scale[Math.floor(i / 3)]); + + const sigma = [ + M[0] * M[0] + M[3] * M[3] + M[6] * M[6], + M[0] * M[1] + M[3] * M[4] + M[6] * M[7], + M[0] * M[2] + M[3] * M[5] + M[6] * M[8], + M[1] * M[1] + M[4] * M[4] + M[7] * M[7], + M[1] * M[2] + M[4] * M[5] + M[7] * M[8], + M[2] * M[2] + M[5] * M[5] + M[8] * M[8], + ]; + + texdata[8 * i + 4] = packHalf2x16(4 * sigma[0], 4 * sigma[1]); + texdata[8 * i + 5] = packHalf2x16(4 * sigma[2], 4 * sigma[3]); + texdata[8 * i + 6] = packHalf2x16(4 * sigma[4], 4 * sigma[5]); + } + + self.postMessage({ texdata, texwidth, texheight }, [texdata.buffer]); + } + + function runSort(viewProj) { + if (!buffer) return; + const f_buffer = new Float32Array(buffer); + if (lastVertexCount == vertexCount) { + let dot = + lastProj[2] * viewProj[2] + + lastProj[6] * viewProj[6] + + lastProj[10] * viewProj[10]; + if (Math.abs(dot - 1) < 0.01) { + return; + } + } else { + generateTexture(); + lastVertexCount = vertexCount; + } + + console.time("sort"); + let maxDepth = -Infinity; + let minDepth = Infinity; + let sizeList = new Int32Array(vertexCount); + for (let i = 0; i < vertexCount; i++) { + let depth = + ((viewProj[2] * f_buffer[8 * i + 0] + + viewProj[6] * f_buffer[8 * i + 1] + + viewProj[10] * f_buffer[8 * i + 2]) * + 4096) | + 0; + sizeList[i] = depth; + if (depth > maxDepth) maxDepth = depth; + if (depth < minDepth) minDepth = depth; + } + + // This is a 16 bit single-pass counting sort + let depthInv = (256 * 256 - 1) / (maxDepth - minDepth); + let counts0 = new Uint32Array(256 * 256); + for (let i = 0; i < vertexCount; i++) { + sizeList[i] = ((sizeList[i] - minDepth) * depthInv) | 0; + counts0[sizeList[i]]++; + } + let starts0 = new Uint32Array(256 * 256); + for (let i = 1; i < 256 * 256; i++) + starts0[i] = starts0[i - 1] + counts0[i - 1]; + depthIndex = new Uint32Array(vertexCount); + for (let i = 0; i < vertexCount; i++) + depthIndex[starts0[sizeList[i]]++] = i; + + console.timeEnd("sort"); + + lastProj = viewProj; + self.postMessage({ depthIndex, viewProj, vertexCount }, [ + depthIndex.buffer, + ]); + } + + function processPlyBuffer(inputBuffer) { + const ubuf = new Uint8Array(inputBuffer); + // 10KB ought to be enough for a header... + const header = new TextDecoder().decode(ubuf.slice(0, 1024 * 10)); + const header_end = "end_header\n"; + const header_end_index = header.indexOf(header_end); + if (header_end_index < 0) + throw new Error("Unable to read .ply file header"); + const vertexCount = parseInt(/element vertex (\d+)\n/.exec(header)[1]); + console.log("Vertex Count", vertexCount); + let row_offset = 0, + offsets = {}, + types = {}; + const TYPE_MAP = { + double: "getFloat64", + int: "getInt32", + uint: "getUint32", + float: "getFloat32", + short: "getInt16", + ushort: "getUint16", + uchar: "getUint8", + }; + for (let prop of header + .slice(0, header_end_index) + .split("\n") + .filter((k) => k.startsWith("property "))) { + const [p, type, name] = prop.split(" "); + const arrayType = TYPE_MAP[type] || "getInt8"; + types[name] = arrayType; + offsets[name] = row_offset; + row_offset += parseInt(arrayType.replace(/[^\d]/g, "")) / 8; + } + console.log("Bytes per row", row_offset, types, offsets); + + let dataView = new DataView( + inputBuffer, + header_end_index + header_end.length, + ); + let row = 0; + const attrs = new Proxy( + {}, + { + get(target, prop) { + if (!types[prop]) throw new Error(prop + " not found"); + return dataView[types[prop]]( + row * row_offset + offsets[prop], + true, + ); + }, + }, + ); + + console.time("calculate importance"); + let sizeList = new Float32Array(vertexCount); + let sizeIndex = new Uint32Array(vertexCount); + for (row = 0; row < vertexCount; row++) { + sizeIndex[row] = row; + if (!types["scale_0"]) continue; + const size = + Math.exp(attrs.scale_0) * + Math.exp(attrs.scale_1) * + Math.exp(attrs.scale_2); + const opacity = 1 / (1 + Math.exp(-attrs.opacity)); + sizeList[row] = size * opacity; + } + console.timeEnd("calculate importance"); + + console.time("sort"); + sizeIndex.sort((b, a) => sizeList[a] - sizeList[b]); + console.timeEnd("sort"); + + // 6*4 + 4 + 4 = 8*4 + // XYZ - Position (Float32) + // XYZ - Scale (Float32) + // RGBA - colors (uint8) + // IJKL - quaternion/rot (uint8) + const rowLength = 3 * 4 + 3 * 4 + 4 + 4; + const buffer = new ArrayBuffer(rowLength * vertexCount); + + console.time("build buffer"); + for (let j = 0; j < vertexCount; j++) { + row = sizeIndex[j]; + + const position = new Float32Array(buffer, j * rowLength, 3); + const scales = new Float32Array(buffer, j * rowLength + 4 * 3, 3); + const rgba = new Uint8ClampedArray( + buffer, + j * rowLength + 4 * 3 + 4 * 3, + 4, + ); + const rot = new Uint8ClampedArray( + buffer, + j * rowLength + 4 * 3 + 4 * 3 + 4, + 4, + ); + + if (types["scale_0"]) { + const qlen = Math.sqrt( + attrs.rot_0 ** 2 + + attrs.rot_1 ** 2 + + attrs.rot_2 ** 2 + + attrs.rot_3 ** 2, + ); + + rot[0] = (attrs.rot_0 / qlen) * 128 + 128; + rot[1] = (attrs.rot_1 / qlen) * 128 + 128; + rot[2] = (attrs.rot_2 / qlen) * 128 + 128; + rot[3] = (attrs.rot_3 / qlen) * 128 + 128; + + scales[0] = Math.exp(attrs.scale_0); + scales[1] = Math.exp(attrs.scale_1); + scales[2] = Math.exp(attrs.scale_2); + } else { + scales[0] = 0.01; + scales[1] = 0.01; + scales[2] = 0.01; + + rot[0] = 255; + rot[1] = 0; + rot[2] = 0; + rot[3] = 0; + } + + position[0] = attrs.x; + position[1] = attrs.y; + position[2] = attrs.z; + + if (types["f_dc_0"]) { + const SH_C0 = 0.28209479177387814; + rgba[0] = (0.5 + SH_C0 * attrs.f_dc_0) * 255; + rgba[1] = (0.5 + SH_C0 * attrs.f_dc_1) * 255; + rgba[2] = (0.5 + SH_C0 * attrs.f_dc_2) * 255; + } else { + rgba[0] = attrs.red; + rgba[1] = attrs.green; + rgba[2] = attrs.blue; + } + if (types["opacity"]) { + rgba[3] = (1 / (1 + Math.exp(-attrs.opacity))) * 255; + } else { + rgba[3] = 255; + } + } + console.timeEnd("build buffer"); + return buffer; + } + + const throttledSort = () => { + if (!sortRunning) { + sortRunning = true; + let lastView = viewProj; + runSort(lastView); + setTimeout(() => { + sortRunning = false; + if (lastView !== viewProj) { + throttledSort(); + } + }, 0); + } + }; + + let sortRunning; + self.onmessage = (e) => { + if (e.data.ply) { + vertexCount = 0; + runSort(viewProj); + buffer = processPlyBuffer(e.data.ply); + vertexCount = Math.floor(buffer.byteLength / rowLength); + postMessage({ buffer: buffer, save: !!e.data.save }); + } else if (e.data.buffer) { + buffer = e.data.buffer; + vertexCount = e.data.vertexCount; + } else if (e.data.vertexCount) { + vertexCount = e.data.vertexCount; + } else if (e.data.view) { + viewProj = e.data.view; + throttledSort(); + } + }; +} + +const vertexShaderSource = ` +#version 300 es +precision highp float; +precision highp int; + +uniform highp usampler2D u_texture; +uniform mat4 projection, view; +uniform vec2 focal; +uniform vec2 viewport; + +in vec2 position; +in int index; + +out vec4 vColor; +out vec2 vPosition; + +void main () { + uvec4 cen = texelFetch(u_texture, ivec2((uint(index) & 0x3ffu) << 1, uint(index) >> 10), 0); + vec4 cam = view * vec4(uintBitsToFloat(cen.xyz), 1); + vec4 pos2d = projection * cam; + + float clip = 1.2 * pos2d.w; + if (pos2d.z < -clip || pos2d.x < -clip || pos2d.x > clip || pos2d.y < -clip || pos2d.y > clip) { + gl_Position = vec4(0.0, 0.0, 2.0, 1.0); + return; + } + + uvec4 cov = texelFetch(u_texture, ivec2(((uint(index) & 0x3ffu) << 1) | 1u, uint(index) >> 10), 0); + vec2 u1 = unpackHalf2x16(cov.x), u2 = unpackHalf2x16(cov.y), u3 = unpackHalf2x16(cov.z); + mat3 Vrk = mat3(u1.x, u1.y, u2.x, u1.y, u2.y, u3.x, u2.x, u3.x, u3.y); + + mat3 J = mat3( + focal.x / cam.z, 0., -(focal.x * cam.x) / (cam.z * cam.z), + 0., -focal.y / cam.z, (focal.y * cam.y) / (cam.z * cam.z), + 0., 0., 0. + ); + + mat3 T = transpose(mat3(view)) * J; + mat3 cov2d = transpose(T) * Vrk * T; + + float mid = (cov2d[0][0] + cov2d[1][1]) / 2.0; + float radius = length(vec2((cov2d[0][0] - cov2d[1][1]) / 2.0, cov2d[0][1])); + float lambda1 = mid + radius, lambda2 = mid - radius; + + if(lambda2 < 0.0) return; + vec2 diagonalVector = normalize(vec2(cov2d[0][1], lambda1 - cov2d[0][0])); + vec2 majorAxis = min(sqrt(2.0 * lambda1), 1024.0) * diagonalVector; + vec2 minorAxis = min(sqrt(2.0 * lambda2), 1024.0) * vec2(diagonalVector.y, -diagonalVector.x); + + vColor = clamp(pos2d.z/pos2d.w+1.0, 0.0, 1.0) * vec4((cov.w) & 0xffu, (cov.w >> 8) & 0xffu, (cov.w >> 16) & 0xffu, (cov.w >> 24) & 0xffu) / 255.0; + vPosition = position; + + vec2 vCenter = vec2(pos2d) / pos2d.w; + gl_Position = vec4( + vCenter + + position.x * majorAxis / viewport + + position.y * minorAxis / viewport, 0.0, 1.0); + +} +`.trim(); + +const fragmentShaderSource = ` +#version 300 es +precision highp float; + +in vec4 vColor; +in vec2 vPosition; + +out vec4 fragColor; + +void main () { + float A = -dot(vPosition, vPosition); + if (A < -4.0) discard; + float B = exp(A) * vColor.a; + fragColor = vec4(B * vColor.rgb, B); +} + +`.trim(); + +// idol: start at the contract camera facing the object's front (was a fixed matrix for the demo scenes) +let defaultViewMatrix = idolViewMatrix(IDOL_START_AZIMUTH); +let viewMatrix = defaultViewMatrix; +async function main() { + let carousel = true; + const params = new URLSearchParams(location.search); + try { + viewMatrix = JSON.parse(decodeURIComponent(location.hash.slice(1))); + carousel = false; + } catch (err) {} + const url = new URL( + // "nike.splat", + // location.href, + params.get("url") || "train.splat", + "https://huggingface.co/cakewalk/splat-data/resolve/main/", + ); + const req = await fetch(url, { + mode: "cors", // no-cors, *cors, same-origin + credentials: "omit", // include, *same-origin, omit + }); + console.log(req); + if (req.status != 200) + throw new Error(req.status + " Unable to load " + req.url); + + const rowLength = 3 * 4 + 3 * 4 + 4 + 4; + const reader = req.body.getReader(); + let splatData = new Uint8Array(req.headers.get("content-length")); + + const downsample = + splatData.length / rowLength > 500000 ? 1 : 1 / devicePixelRatio; + console.log(splatData.length / rowLength, downsample); + + const worker = new Worker( + URL.createObjectURL( + new Blob(["(", createWorker.toString(), ")(self)"], { + type: "application/javascript", + }), + ), + ); + + const canvas = document.getElementById("canvas"); + const fps = document.getElementById("fps"); + const camid = document.getElementById("camid"); + + let projectionMatrix; + + const gl = canvas.getContext("webgl2", { + antialias: false, + }); + + const vertexShader = gl.createShader(gl.VERTEX_SHADER); + gl.shaderSource(vertexShader, vertexShaderSource); + gl.compileShader(vertexShader); + if (!gl.getShaderParameter(vertexShader, gl.COMPILE_STATUS)) + console.error(gl.getShaderInfoLog(vertexShader)); + + const fragmentShader = gl.createShader(gl.FRAGMENT_SHADER); + gl.shaderSource(fragmentShader, fragmentShaderSource); + gl.compileShader(fragmentShader); + if (!gl.getShaderParameter(fragmentShader, gl.COMPILE_STATUS)) + console.error(gl.getShaderInfoLog(fragmentShader)); + + const program = gl.createProgram(); + gl.attachShader(program, vertexShader); + gl.attachShader(program, fragmentShader); + gl.linkProgram(program); + gl.useProgram(program); + + if (!gl.getProgramParameter(program, gl.LINK_STATUS)) + console.error(gl.getProgramInfoLog(program)); + + gl.disable(gl.DEPTH_TEST); // Disable depth testing + + // Enable blending + gl.enable(gl.BLEND); + gl.blendFuncSeparate( + gl.ONE_MINUS_DST_ALPHA, + gl.ONE, + gl.ONE_MINUS_DST_ALPHA, + gl.ONE, + ); + gl.blendEquationSeparate(gl.FUNC_ADD, gl.FUNC_ADD); + + const u_projection = gl.getUniformLocation(program, "projection"); + const u_viewport = gl.getUniformLocation(program, "viewport"); + const u_focal = gl.getUniformLocation(program, "focal"); + const u_view = gl.getUniformLocation(program, "view"); + + // positions + const triangleVertices = new Float32Array([-2, -2, 2, -2, 2, 2, -2, 2]); + const vertexBuffer = gl.createBuffer(); + gl.bindBuffer(gl.ARRAY_BUFFER, vertexBuffer); + gl.bufferData(gl.ARRAY_BUFFER, triangleVertices, gl.STATIC_DRAW); + const a_position = gl.getAttribLocation(program, "position"); + gl.enableVertexAttribArray(a_position); + gl.bindBuffer(gl.ARRAY_BUFFER, vertexBuffer); + gl.vertexAttribPointer(a_position, 2, gl.FLOAT, false, 0, 0); + + var texture = gl.createTexture(); + gl.bindTexture(gl.TEXTURE_2D, texture); + + var u_textureLocation = gl.getUniformLocation(program, "u_texture"); + gl.uniform1i(u_textureLocation, 0); + + const indexBuffer = gl.createBuffer(); + const a_index = gl.getAttribLocation(program, "index"); + gl.enableVertexAttribArray(a_index); + gl.bindBuffer(gl.ARRAY_BUFFER, indexBuffer); + gl.vertexAttribIPointer(a_index, 1, gl.INT, false, 0, 0); + gl.vertexAttribDivisor(a_index, 1); + + const resize = () => { + // idol: focal length for a 40 degree field of view across the shorter side of the screen (the square + // contract render's framing, whole object in view in portrait and landscape) + camera.fx = camera.fy = + Math.min(innerWidth, innerHeight) / 2 / Math.tan(IDOL_FOV / 2); + gl.uniform2fv(u_focal, new Float32Array([camera.fx, camera.fy])); + + projectionMatrix = getProjectionMatrix( + camera.fx, + camera.fy, + innerWidth, + innerHeight, + ); + + gl.uniform2fv(u_viewport, new Float32Array([innerWidth, innerHeight])); + + gl.canvas.width = Math.round(innerWidth / downsample); + gl.canvas.height = Math.round(innerHeight / downsample); + gl.viewport(0, 0, gl.canvas.width, gl.canvas.height); + + gl.uniformMatrix4fv(u_projection, false, projectionMatrix); + }; + + window.addEventListener("resize", resize); + resize(); + + worker.onmessage = (e) => { + if (e.data.buffer) { + splatData = new Uint8Array(e.data.buffer); + if (e.data.save) { + const blob = new Blob([splatData.buffer], { + type: "application/octet-stream", + }); + const link = document.createElement("a"); + link.download = "model.splat"; + link.href = URL.createObjectURL(blob); + document.body.appendChild(link); + link.click(); + } + } else if (e.data.texdata) { + const { texdata, texwidth, texheight } = e.data; + // console.log(texdata) + gl.bindTexture(gl.TEXTURE_2D, texture); + gl.texParameteri( + gl.TEXTURE_2D, + gl.TEXTURE_WRAP_S, + gl.CLAMP_TO_EDGE, + ); + gl.texParameteri( + gl.TEXTURE_2D, + gl.TEXTURE_WRAP_T, + gl.CLAMP_TO_EDGE, + ); + gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST); + gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST); + + gl.texImage2D( + gl.TEXTURE_2D, + 0, + gl.RGBA32UI, + texwidth, + texheight, + 0, + gl.RGBA_INTEGER, + gl.UNSIGNED_INT, + texdata, + ); + gl.activeTexture(gl.TEXTURE0); + gl.bindTexture(gl.TEXTURE_2D, texture); + } else if (e.data.depthIndex) { + const { depthIndex, viewProj } = e.data; + gl.bindBuffer(gl.ARRAY_BUFFER, indexBuffer); + gl.bufferData(gl.ARRAY_BUFFER, depthIndex, gl.DYNAMIC_DRAW); + vertexCount = e.data.vertexCount; + } + }; + + let activeKeys = []; + let currentCameraIndex = 0; + + window.addEventListener("keydown", (e) => { + // if (document.activeElement != document.body) return; + carousel = false; + if (!activeKeys.includes(e.code)) activeKeys.push(e.code); + if (/\d/.test(e.key)) { + currentCameraIndex = parseInt(e.key); + camera = cameras[currentCameraIndex]; + viewMatrix = getViewMatrix(camera); + } + if (["-", "_"].includes(e.key)) { + currentCameraIndex = + (currentCameraIndex + cameras.length - 1) % cameras.length; + viewMatrix = getViewMatrix(cameras[currentCameraIndex]); + } + if (["+", "="].includes(e.key)) { + currentCameraIndex = (currentCameraIndex + 1) % cameras.length; + viewMatrix = getViewMatrix(cameras[currentCameraIndex]); + } + camid.innerText = "cam " + currentCameraIndex; + if (e.code == "KeyV") { + location.hash = + "#" + + JSON.stringify( + viewMatrix.map((k) => Math.round(k * 100) / 100), + ); + camid.innerText = ""; + } else if (e.code === "KeyP") { + carousel = true; + camid.innerText = ""; + } + }); + window.addEventListener("keyup", (e) => { + activeKeys = activeKeys.filter((k) => k !== e.code); + }); + window.addEventListener("blur", () => { + activeKeys = []; + }); + + window.addEventListener( + "wheel", + (e) => { + carousel = false; + e.preventDefault(); + const lineHeight = 10; + const scale = + e.deltaMode == 1 + ? lineHeight + : e.deltaMode == 2 + ? innerHeight + : 1; + let inv = invert4(viewMatrix); + if (e.shiftKey) { + inv = translate4( + inv, + (e.deltaX * scale) / innerWidth, + (e.deltaY * scale) / innerHeight, + 0, + ); + } else if (e.ctrlKey || e.metaKey) { + // inv = rotate4(inv, (e.deltaX * scale) / innerWidth, 0, 0, 1); + // inv = translate4(inv, 0, (e.deltaY * scale) / innerHeight, 0); + // let preY = inv[13]; + inv = translate4( + inv, + 0, + 0, + (-10 * (e.deltaY * scale)) / innerHeight, + ); + // inv[13] = preY; + } else { + let d = 4; + inv = translate4(inv, 0, 0, d); + inv = rotate4(inv, -(e.deltaX * scale) / innerWidth, 0, 1, 0); + inv = rotate4(inv, (e.deltaY * scale) / innerHeight, 1, 0, 0); + inv = translate4(inv, 0, 0, -d); + } + + viewMatrix = invert4(inv); + }, + { passive: false }, + ); + + let startX, startY, down; + canvas.addEventListener("mousedown", (e) => { + carousel = false; + e.preventDefault(); + startX = e.clientX; + startY = e.clientY; + down = e.ctrlKey || e.metaKey ? 2 : 1; + }); + canvas.addEventListener("contextmenu", (e) => { + carousel = false; + e.preventDefault(); + startX = e.clientX; + startY = e.clientY; + down = 2; + }); + + canvas.addEventListener("mousemove", (e) => { + e.preventDefault(); + if (down == 1) { + let inv = invert4(viewMatrix); + let dx = (5 * (e.clientX - startX)) / innerWidth; + let dy = (5 * (e.clientY - startY)) / innerHeight; + let d = 4; + + inv = translate4(inv, 0, 0, d); + inv = rotate4(inv, dx, 0, 1, 0); + inv = rotate4(inv, -dy, 1, 0, 0); + inv = translate4(inv, 0, 0, -d); + // let postAngle = Math.atan2(inv[0], inv[10]) + // inv = rotate4(inv, postAngle - preAngle, 0, 0, 1) + // console.log(postAngle) + viewMatrix = invert4(inv); + + startX = e.clientX; + startY = e.clientY; + } else if (down == 2) { + let inv = invert4(viewMatrix); + // inv = rotateY(inv, ); + // let preY = inv[13]; + inv = translate4( + inv, + (-10 * (e.clientX - startX)) / innerWidth, + 0, + (10 * (e.clientY - startY)) / innerHeight, + ); + // inv[13] = preY; + viewMatrix = invert4(inv); + + startX = e.clientX; + startY = e.clientY; + } + }); + canvas.addEventListener("mouseup", (e) => { + e.preventDefault(); + down = false; + startX = 0; + startY = 0; + }); + + let altX = 0, + altY = 0; + canvas.addEventListener( + "touchstart", + (e) => { + e.preventDefault(); + if (e.touches.length === 1) { + carousel = false; + startX = e.touches[0].clientX; + startY = e.touches[0].clientY; + down = 1; + } else if (e.touches.length === 2) { + // console.log('beep') + carousel = false; + startX = e.touches[0].clientX; + altX = e.touches[1].clientX; + startY = e.touches[0].clientY; + altY = e.touches[1].clientY; + down = 1; + } + }, + { passive: false }, + ); + canvas.addEventListener( + "touchmove", + (e) => { + e.preventDefault(); + if (e.touches.length === 1 && down) { + let inv = invert4(viewMatrix); + let dx = (4 * (e.touches[0].clientX - startX)) / innerWidth; + let dy = (4 * (e.touches[0].clientY - startY)) / innerHeight; + + // idol: orbit about the object at the origin (was a point 4 units ahead) + let d = Math.hypot(inv[12], inv[13], inv[14]); + inv = translate4(inv, 0, 0, d); + // inv = translate4(inv, -x, -y, -z); + // inv = translate4(inv, x, y, z); + inv = rotate4(inv, dx, 0, 1, 0); + inv = rotate4(inv, -dy, 1, 0, 0); + inv = translate4(inv, 0, 0, -d); + + viewMatrix = invert4(inv); + + startX = e.touches[0].clientX; + startY = e.touches[0].clientY; + } else if (e.touches.length === 2) { + // alert('beep') + const dtheta = + Math.atan2(startY - altY, startX - altX) - + Math.atan2( + e.touches[0].clientY - e.touches[1].clientY, + e.touches[0].clientX - e.touches[1].clientX, + ); + const dscale = + Math.hypot(startX - altX, startY - altY) / + Math.hypot( + e.touches[0].clientX - e.touches[1].clientX, + e.touches[0].clientY - e.touches[1].clientY, + ); + const dx = + (e.touches[0].clientX + + e.touches[1].clientX - + (startX + altX)) / + 2; + const dy = + (e.touches[0].clientY + + e.touches[1].clientY - + (startY + altY)) / + 2; + let inv = invert4(viewMatrix); + // inv = translate4(inv, 0, 0, d); + inv = rotate4(inv, dtheta, 0, 0, 1); + + inv = translate4(inv, -dx / innerWidth, -dy / innerHeight, 0); + + // let preY = inv[13]; + inv = translate4(inv, 0, 0, 3 * (1 - dscale)); + // inv[13] = preY; + + viewMatrix = invert4(inv); + + startX = e.touches[0].clientX; + altX = e.touches[1].clientX; + startY = e.touches[0].clientY; + altY = e.touches[1].clientY; + } + }, + { passive: false }, + ); + canvas.addEventListener( + "touchend", + (e) => { + e.preventDefault(); + down = false; + startX = 0; + startY = 0; + }, + { passive: false }, + ); + + let jumpDelta = 0; + let vertexCount = 0; + + let lastFrame = 0; + let avgFps = 0; + let start = 0; + + window.addEventListener("gamepadconnected", (e) => { + const gp = navigator.getGamepads()[e.gamepad.index]; + console.log( + `Gamepad connected at index ${gp.index}: ${gp.id}. It has ${gp.buttons.length} buttons and ${gp.axes.length} axes.`, + ); + }); + window.addEventListener("gamepaddisconnected", (e) => { + console.log("Gamepad disconnected"); + }); + + let leftGamepadTrigger, rightGamepadTrigger; + + const frame = (now) => { + let inv = invert4(viewMatrix); + let shiftKey = + activeKeys.includes("Shift") || + activeKeys.includes("ShiftLeft") || + activeKeys.includes("ShiftRight"); + + if (activeKeys.includes("ArrowUp")) { + if (shiftKey) { + inv = translate4(inv, 0, -0.03, 0); + } else { + inv = translate4(inv, 0, 0, 0.1); + } + } + if (activeKeys.includes("ArrowDown")) { + if (shiftKey) { + inv = translate4(inv, 0, 0.03, 0); + } else { + inv = translate4(inv, 0, 0, -0.1); + } + } + if (activeKeys.includes("ArrowLeft")) + inv = translate4(inv, -0.03, 0, 0); + // + if (activeKeys.includes("ArrowRight")) + inv = translate4(inv, 0.03, 0, 0); + // inv = rotate4(inv, 0.01, 0, 1, 0); + if (activeKeys.includes("KeyA")) inv = rotate4(inv, -0.01, 0, 1, 0); + if (activeKeys.includes("KeyD")) inv = rotate4(inv, 0.01, 0, 1, 0); + if (activeKeys.includes("KeyQ")) inv = rotate4(inv, 0.01, 0, 0, 1); + if (activeKeys.includes("KeyE")) inv = rotate4(inv, -0.01, 0, 0, 1); + if (activeKeys.includes("KeyW")) inv = rotate4(inv, 0.005, 1, 0, 0); + if (activeKeys.includes("KeyS")) inv = rotate4(inv, -0.005, 1, 0, 0); + + const gamepads = navigator.getGamepads ? navigator.getGamepads() : []; + let isJumping = activeKeys.includes("Space"); + for (let gamepad of gamepads) { + if (!gamepad) continue; + + const axisThreshold = 0.1; // Threshold to detect when the axis is intentionally moved + const moveSpeed = 0.06; + const rotateSpeed = 0.02; + + // Assuming the left stick controls translation (axes 0 and 1) + if (Math.abs(gamepad.axes[0]) > axisThreshold) { + inv = translate4(inv, moveSpeed * gamepad.axes[0], 0, 0); + carousel = false; + } + if (Math.abs(gamepad.axes[1]) > axisThreshold) { + inv = translate4(inv, 0, 0, -moveSpeed * gamepad.axes[1]); + carousel = false; + } + if (gamepad.buttons[12].pressed || gamepad.buttons[13].pressed) { + inv = translate4( + inv, + 0, + -moveSpeed * + (gamepad.buttons[12].pressed - + gamepad.buttons[13].pressed), + 0, + ); + carousel = false; + } + + if (gamepad.buttons[14].pressed || gamepad.buttons[15].pressed) { + inv = translate4( + inv, + -moveSpeed * + (gamepad.buttons[14].pressed - + gamepad.buttons[15].pressed), + 0, + 0, + ); + carousel = false; + } + + // Assuming the right stick controls rotation (axes 2 and 3) + if (Math.abs(gamepad.axes[2]) > axisThreshold) { + inv = rotate4(inv, rotateSpeed * gamepad.axes[2], 0, 1, 0); + carousel = false; + } + if (Math.abs(gamepad.axes[3]) > axisThreshold) { + inv = rotate4(inv, -rotateSpeed * gamepad.axes[3], 1, 0, 0); + carousel = false; + } + + let tiltAxis = gamepad.buttons[6].value - gamepad.buttons[7].value; + if (Math.abs(tiltAxis) > axisThreshold) { + inv = rotate4(inv, rotateSpeed * tiltAxis, 0, 0, 1); + carousel = false; + } + if (gamepad.buttons[4].pressed && !leftGamepadTrigger) { + camera = + cameras[(cameras.indexOf(camera) + 1) % cameras.length]; + inv = invert4(getViewMatrix(camera)); + carousel = false; + } + if (gamepad.buttons[5].pressed && !rightGamepadTrigger) { + camera = + cameras[ + (cameras.indexOf(camera) + cameras.length - 1) % + cameras.length + ]; + inv = invert4(getViewMatrix(camera)); + carousel = false; + } + leftGamepadTrigger = gamepad.buttons[4].pressed; + rightGamepadTrigger = gamepad.buttons[5].pressed; + if (gamepad.buttons[0].pressed) { + isJumping = true; + carousel = false; + } + if (gamepad.buttons[3].pressed) { + carousel = true; + } + } + + if ( + ["KeyJ", "KeyK", "KeyL", "KeyI"].some((k) => activeKeys.includes(k)) + ) { + let d = 4; + inv = translate4(inv, 0, 0, d); + inv = rotate4( + inv, + activeKeys.includes("KeyJ") + ? -0.05 + : activeKeys.includes("KeyL") + ? 0.05 + : 0, + 0, + 1, + 0, + ); + inv = rotate4( + inv, + activeKeys.includes("KeyI") + ? 0.05 + : activeKeys.includes("KeyK") + ? -0.05 + : 0, + 1, + 0, + 0, + ); + inv = translate4(inv, 0, 0, -d); + } + + viewMatrix = invert4(inv); + + if (carousel) { + // idol: a full turn about the vertical axis, about 12 s per turn, until the first touch (was a swing + // left and right for the demo scenes) + const azimuth = + IDOL_START_AZIMUTH + + (2 * Math.PI * Math.max(0, Date.now() - start)) / IDOL_TURN_MS; + viewMatrix = idolViewMatrix(azimuth); + } + + if (isJumping) { + jumpDelta = Math.min(1, jumpDelta + 0.05); + } else { + jumpDelta = Math.max(0, jumpDelta - 0.05); + } + + let inv2 = invert4(viewMatrix); + inv2 = translate4(inv2, 0, -jumpDelta, 0); + inv2 = rotate4(inv2, -0.1 * jumpDelta, 1, 0, 0); + let actualViewMatrix = invert4(inv2); + + const viewProj = multiply4(projectionMatrix, actualViewMatrix); + worker.postMessage({ view: viewProj }); + + const currentFps = 1000 / (now - lastFrame) || 0; + avgFps = avgFps * 0.9 + currentFps * 0.1; + + if (vertexCount > 0) { + document.getElementById("spinner").style.display = "none"; + gl.uniformMatrix4fv(u_view, false, actualViewMatrix); + gl.clear(gl.COLOR_BUFFER_BIT); + gl.drawArraysInstanced(gl.TRIANGLE_FAN, 0, 4, vertexCount); + } else { + gl.clear(gl.COLOR_BUFFER_BIT); + document.getElementById("spinner").style.display = ""; + start = Date.now() + 2000; + } + const progress = (100 * vertexCount) / (splatData.length / rowLength); + if (progress < 100) { + document.getElementById("progress").style.width = progress + "%"; + } else { + document.getElementById("progress").style.display = "none"; + } + fps.innerText = Math.round(avgFps) + " fps"; + if (isNaN(currentCameraIndex)) { + camid.innerText = ""; + } + lastFrame = now; + requestAnimationFrame(frame); + }; + + frame(); + + const isPly = (splatData) => + splatData[0] == 112 && + splatData[1] == 108 && + splatData[2] == 121 && + splatData[3] == 10; + + const selectFile = (file) => { + const fr = new FileReader(); + if (/\.json$/i.test(file.name)) { + fr.onload = () => { + cameras = JSON.parse(fr.result); + viewMatrix = getViewMatrix(cameras[0]); + projectionMatrix = getProjectionMatrix( + camera.fx / downsample, + camera.fy / downsample, + canvas.width, + canvas.height, + ); + gl.uniformMatrix4fv(u_projection, false, projectionMatrix); + + console.log("Loaded Cameras"); + }; + fr.readAsText(file); + } else { + stopLoading = true; + fr.onload = () => { + splatData = new Uint8Array(fr.result); + console.log("Loaded", Math.floor(splatData.length / rowLength)); + + if (isPly(splatData)) { + // ply file magic header means it should be handled differently + worker.postMessage({ ply: splatData.buffer, save: true }); + } else { + worker.postMessage({ + buffer: splatData.buffer, + vertexCount: Math.floor(splatData.length / rowLength), + }); + } + }; + fr.readAsArrayBuffer(file); + } + }; + + window.addEventListener("hashchange", (e) => { + try { + viewMatrix = JSON.parse(decodeURIComponent(location.hash.slice(1))); + carousel = false; + } catch (err) {} + }); + + const preventDefault = (e) => { + e.preventDefault(); + e.stopPropagation(); + }; + document.addEventListener("dragenter", preventDefault); + document.addEventListener("dragover", preventDefault); + document.addEventListener("dragleave", preventDefault); + document.addEventListener("drop", (e) => { + e.preventDefault(); + e.stopPropagation(); + selectFile(e.dataTransfer.files[0]); + }); + + let bytesRead = 0; + let lastVertexCount = -1; + let stopLoading = false; + + while (true) { + const { done, value } = await reader.read(); + if (done || stopLoading) break; + + splatData.set(value, bytesRead); + bytesRead += value.length; + + if (vertexCount > lastVertexCount) { + if (!isPly(splatData)) { + worker.postMessage({ + buffer: splatData.buffer, + vertexCount: Math.floor(bytesRead / rowLength), + }); + } + lastVertexCount = vertexCount; + } + } + if (!stopLoading) { + if (isPly(splatData)) { + // ply file magic header means it should be handled differently + worker.postMessage({ ply: splatData.buffer, save: false }); + } else { + worker.postMessage({ + buffer: splatData.buffer, + vertexCount: Math.floor(bytesRead / rowLength), + }); + } + } +} + +main().catch((err) => { + document.getElementById("spinner").style.display = "none"; + document.getElementById("message").innerText = err.toString(); +}); diff --git a/android/app/src/main/java/com/idol/AutoRun.kt b/android/app/src/main/java/com/idol/AutoRun.kt new file mode 100644 index 0000000000000000000000000000000000000000..f72d40fb1392555c9b166767f70edd4ea9d72817 --- /dev/null +++ b/android/app/src/main/java/com/idol/AutoRun.kt @@ -0,0 +1,312 @@ +package com.idol + +import android.os.Build +import android.os.Process +import android.os.SystemClock +import android.util.Log +import com.google.ai.edge.litert.CompiledModel +import org.json.JSONArray +import org.json.JSONObject +import java.io.File + +/** + * Debug builds only: runs a plan of rows through the app and writes what each produced, for the device lines of + * CONTRACT.md ((a) parity, (b) the .ply, (c) speed, (d) failure paths, (e) memory). + * + * adb shell am start -n com.idol/.MainActivity --es autorun plans/.json --es out runs/ + * + * Paths are under the app's files dir. Plan: {"preload": true, "reserve_mb": 2500, "mem_check": true, + * "handoff": "host" | "direct", "matting": "bir" | "dis", "cpu_threads": 6, "gpu_priority": "default", "rows": [...]} + * (every key optional; a row may override "preload"). Rows: + * {"id", "kind": "e2e", "fixture": "fx/", "runs": 4, "preload": true}: the app's whole pipeline from the + * fixture's raw image and mode with its noise arrays; ..ply + ..json per run (run 0 is the + * warm-up), .alpha.u8 once (raw uint8 matting alpha), .json for the row. + * {"id", "kind": "parity", "fixture": "fx/"}: each part on the fixture's reference input, ..f32. + * {"id", "kind": "image", "image": "", "mode": "real"}: a normal run of a picture (failure paths); + * .json with the outcome and the message on the screen. + * {"id", "kind": "snippet", "fixture": "fx/"}: the model card's Kotlin block (CardSnippet.kt) on the + * fixture's picture; its line goes to .json, DONE.json and logcat. + * Output: /header.json (device, options, the load and compile seconds of each preload), the row files, + * /DONE.json; logcat tag IdolAutoRun: START / ROW / DONE lines. A file /STOP ends the plan after the + * current run. + */ +class AutoRun(private val activity: MainActivity, private val planFile: File, private val outDir: File) { + private val header = JSONObject() + private val loads = JSONArray() + + fun start() { + Thread({ + try { + runAll() + } catch (e: Throwable) { + Log.e(TAG, "IDOL_AUTORUN FAILED", e) + runCatching { + outDir.mkdirs() + File(outDir, "DONE.json").writeText( + JSONObject().put("status", "FAILED").put("error", e.toString()).toString(1), + ) + } + Log.i(TAG, "IDOL_AUTORUN DONE status=FAILED pid=${Process.myPid()}") + } finally { + activity.autorunFinished() + } + }, "IdolAutoRun").start() + } + + private fun stopped() = File(outDir, "STOP").exists() + + private fun writeHeader() = File(outDir, "header.json").writeText(header.toString(1)) + + private fun runAll() { + val plan = JSONObject(planFile.readText()) + outDir.mkdirs() + File(outDir, "DONE.json").delete() + header.put("plan", planFile.name).put("plan_json", plan) + .put("device", DeviceInfo.displayName()).put("model", Build.MODEL).put("android_sdk", Build.VERSION.SDK_INT) + .put("fingerprint", Build.FINGERPRINT).put("litert", BuildConfig.LITERT_VERSION) + .put("app_version", BuildConfig.VERSION_NAME).put("pid", Process.myPid()) + .put("mem_total_kb", DeviceInfo.memTotalKb()).put("mem_available_kb", DeviceInfo.memAvailableKb()) + .put("thermal", DeviceInfo.thermalStatus(activity)).put("cpuset", DeviceInfo.cpuset()) + .put("start_ms", System.currentTimeMillis()).put("load", loads) + writeHeader() + Log.i(TAG, "IDOL_AUTORUN START plan=${planFile.name} out=${outDir.name} pid=${Process.myPid()}") + val rows = plan.getJSONArray("rows") + val summary = JSONArray() + val done = JSONObject().put("status", "DONE") + for (r in 0 until rows.length()) { + if (stopped()) { + done.put("stopped", true) + break + } + val row = rows.getJSONObject(r) + val id = row.getString("id") + val kind = row.optString("kind", "e2e") + Log.i(TAG, "IDOL_AUTORUN ROW_START id=$id kind=$kind") + val rec = try { + when (kind) { + "e2e" -> e2e(plan, row) + "parity" -> parity(plan, row) + "image" -> image(plan, row) + "snippet" -> snippet(row) + else -> JSONObject().put("id", id).put("kind", kind).put("status", "unknown_kind") + } + } catch (e: Throwable) { + Log.e(TAG, "row $id failed", e) + JSONObject().put("id", id).put("kind", kind).put("status", "failed").put("error", e.toString()) + } + summary.put(rec) + if (kind == "snippet") done.put("snippet", rec.optString("line")) + Log.i(TAG, "IDOL_AUTORUN ROW id=$id kind=$kind status=${rec.optString("status")} ${rec.optString("summary")}") + } + done.put("rows", summary).put("pid", Process.myPid()).put("end_ms", System.currentTimeMillis()) + .put("self_status_kb", JSONObject(DeviceInfo.selfStatusKb() as Map<*, *>)) + File(outDir, "DONE.json").writeText(done.toString(1)) + Log.i(TAG, "IDOL_AUTORUN DONE rows=${summary.length()} pid=${Process.myPid()}") + } + + /** The engine options of [row] in [plan] (the row's keys first). */ + private fun options(plan: JSONObject, row: JSONObject): EngineOptions { + fun has(k: String) = row.has(k) || plan.has(k) + fun src(k: String) = if (row.has(k)) row else plan + val defaults = EngineOptions() + val manifestMatting = ModelFiles.manifest(ModelFiles.dir(activity.filesDir))?.defaultMatting ?: "bir" + val priority = when (if (has("gpu_priority")) src("gpu_priority").getString("gpu_priority") else "default") { + "low" -> CompiledModel.GpuOptions.Priority.LOW + "normal" -> CompiledModel.GpuOptions.Priority.NORMAL + "high" -> CompiledModel.GpuOptions.Priority.HIGH + else -> null + } + return EngineOptions( + preload = if (has("preload")) src("preload").getBoolean("preload") else defaults.preload, + reserveMb = if (has("reserve_mb")) src("reserve_mb").getLong("reserve_mb") else defaults.reserveMb, + memoryCheck = if (has("mem_check")) src("mem_check").getBoolean("mem_check") else defaults.memoryCheck, + directHandoff = has("handoff") && src("handoff").getString("handoff") == "direct", + matting = if (has("matting")) src("matting").getString("matting") else manifestMatting, + cpuThreads = if (has("cpu_threads")) src("cpu_threads").getInt("cpu_threads") else defaults.cpuThreads, + gpuPriority = priority, + disFloor = if (has("dis_floor")) src("dis_floor").getInt("dis_floor") else defaults.disFloor, + ) + } + + private fun json(o: EngineOptions) = JSONObject().put("preload", o.preload).put("reserve_mb", o.reserveMb) + .put("mem_check", o.memoryCheck).put("handoff", if (o.directHandoff) "direct" else "host") + .put("matting", o.matting).put("cpu_threads", o.cpuThreads) + .put("gpu_priority", o.gpuPriority?.name ?: "default").put("dis_floor", o.disFloor) + + private fun outcomeFields(rec: JSONObject, outcome: Outcome) { + when (outcome) { + is Outcome.Done -> rec.put("status", "done") + is Outcome.NoObject -> rec.put("status", "no_object").put("reason", outcome.reason) + is Outcome.Missing -> rec.put("status", "missing_models").put("message", ModelFiles.message(outcome.problems)) + is Outcome.Failed -> rec.put("status", outcome.kind).put("message", outcome.message) + } + } + + private fun e2e(plan: JSONObject, row: JSONObject): JSONObject { + val id = row.getString("id") + val fx = Fixture(File(activity.filesDir, row.getString("fixture"))) + val runs = row.optInt("runs", 4) + val opts = options(plan, row) + val rowRec = JSONObject().put("id", id).put("kind", "e2e").put("fixture", row.getString("fixture")) + .put("mode", fx.mode).put("runs", runs).put("options", json(opts)) + val problems = activity.ensureEngine(opts) + if (problems.isNotEmpty()) { + return rowRec.put("status", "missing_models").put("message", ModelFiles.message(problems)) + } + if (opts.preload) { + val load = activity.onModel { requireNotNull(Idol.engine).preload(activity.progress) } + loads.put(JSONObject().put("row", id).put("record", load)) + writeHeader() + rowRec.put("load_total_s", load.optDouble("total_s")).put("load_loaded", load.optInt("loaded")) + .put("load_of", load.optInt("of")) + } + val photo = Photo.fromFile(fx.image) + val noise = fx.noise() + val statuses = JSONArray() + val measured = ArrayList() + var alphaWritten = false + for (k in 0 until runs) { + if (stopped()) { + rowRec.put("stopped_after", k) + break + } + val ply = File(outDir, "$id.$k.ply") + val rec = JSONObject().put("id", id).put("run", k).put("warmup", k == 0) + .put("fixture", row.getString("fixture")).put("mode", fx.mode).put("noise", noise.source) + .put("options", json(opts)) + .put("thermal_start", DeviceInfo.thermalStatus(activity)) + .put("mem_available_kb_start", DeviceInfo.memAvailableKb()) + .put("vmhwm_reset", DeviceInfo.resetPeakRss()) + .put("cpuset", DeviceInfo.cpuset()) + val sampler = GpuPageSampler().also { it.start() } + val t0 = SystemClock.elapsedRealtimeNanos() + val outcome = activity.runAndWait(photo, fx.mode, noise, ply, showResult = false) + val wall = (SystemClock.elapsedRealtimeNanos() - t0) / 1e9 + val gpuMax = sampler.stop() + val st = DeviceInfo.selfStatusKb() + rec.put("wall_s", wall).put("thermal_end", DeviceInfo.thermalStatus(activity)) + .put("vmhwm_kb", st["VmHWM"] ?: -1L).put("vmrss_kb", st["VmRSS"] ?: -1L) + .put("vmswap_kb", st["VmSwap"] ?: -1L).put("gpu_page_alloc_max_bytes", gpuMax) + .put("mem_available_kb_end", DeviceInfo.memAvailableKb()) + outcomeFields(rec, outcome) + if (outcome is Outcome.Done) { + val r = outcome.result + val stages = JSONObject() + for ((key, v) in r.times) if (key != "generate" && key != "all") stages.put(key, v) + rec.put("stages", stages).put("generate", r.times["generate"]).put("total", r.times["all"]) + .put("ply", ply.name).put("ply_bytes", ply.length()).put("gaussians", r.gaussians) + .put("box", JSONArray(r.pre.box.toList())) + .put("alpha_size", JSONArray(listOf(r.pre.alphaWidth, r.pre.alphaHeight))) + .put("own_alpha", r.pre.ownAlpha).put("handoff", r.handoff) + .put("compiles_in_run", JSONObject().put("count", r.compiles.size) + .put("seconds", r.compiles.sumOf { it.seconds }) + .put("graphs", JSONArray(r.compiles.map { it.json() }))) + .put("memory_events", JSONArray(r.events)) + if (!alphaWritten) { + File(outDir, "$id.alpha.u8").writeBytes(r.pre.alphaRaw) + rowRec.put("alpha", JSONObject().put("file", "$id.alpha.u8").put("width", r.pre.alphaWidth) + .put("height", r.pre.alphaHeight).put("box", JSONArray(r.pre.box.toList()))) + alphaWritten = true + } + if (k > 0) r.times["generate"]?.let { measured.add(it) } + } + File(outDir, "$id.$k.json").writeText(rec.toString(1)) + statuses.put(rec.optString("status")) + Log.i(TAG, "IDOL_AUTORUN RUN id=$id run=$k warmup=${k == 0} status=${rec.optString("status")} " + + "generate_s=${rec.optDouble("generate", -1.0)} vmhwm_kb=${rec.optLong("vmhwm_kb")}") + } + val allDone = (0 until statuses.length()).all { statuses.getString(it) == "done" } + rowRec.put("statuses", statuses).put("status", if (allDone && statuses.length() > 0) "done" else "incomplete") + if (measured.isNotEmpty()) { + val sorted = measured.sorted() + val median = if (sorted.size % 2 == 1) sorted[sorted.size / 2] + else (sorted[sorted.size / 2 - 1] + sorted[sorted.size / 2]) / 2 + rowRec.put("generate_measured_s", JSONArray(measured)).put("generate_median_s", median) + .put("summary", "generate median %.2f s over %d measured runs".format(median, measured.size)) + } + File(outDir, "$id.json").writeText(rowRec.toString(1)) + return rowRec + } + + private fun parity(plan: JSONObject, row: JSONObject): JSONObject { + val id = row.getString("id") + val fx = Fixture(File(activity.filesDir, row.getString("fixture"))) + val opts = options(plan, row).copy(preload = false) + val problems = activity.ensureEngine(opts) + if (problems.isNotEmpty()) { + return JSONObject().put("id", id).put("kind", "parity").put("status", "missing_models") + .put("message", ModelFiles.message(problems)) + } + activity.showRunScreen("Part agreement: ${fx.dir.name}") + val t0 = SystemClock.elapsedRealtimeNanos() + val rec = activity.onModel { requireNotNull(Idol.engine).parity(fx, outDir, id, activity.progress) } + rec.put("kind", "parity").put("status", "done").put("options", json(opts)) + .put("wall_s", (SystemClock.elapsedRealtimeNanos() - t0) / 1e9) + .put("summary", "${rec.getJSONObject("arrays").length()} arrays") + File(outDir, "$id.json").writeText(rec.toString(1)) + return rec + } + + private fun image(plan: JSONObject, row: JSONObject): JSONObject { + val id = row.getString("id") + val path = row.getString("image") + val mode = row.optString("mode", "real") + val rec = JSONObject().put("id", id).put("kind", "image").put("image", path).put("mode", mode) + activity.ensureEngine(options(plan, row)) + val photo = try { + Photo.fromFile(File(activity.filesDir, path)) + } catch (e: Throwable) { + return rec.put("status", "decode_failed").put("error", e.toString()) + } + val outcome = activity.runAndWait( + photo, mode, Noise.random(MainActivity.DEFAULT_SEED), File(outDir, "$id.ply"), + showResult = row.optBoolean("show_result", false), + ) + outcomeFields(rec, outcome) + if (outcome is Outcome.Done) rec.put("generate", outcome.result.times["generate"]) + // the outcome is on the screen once the UI thread has run the posts before this one + rec.put("message_on_screen", activity.homeMessageText()).put("pid", Process.myPid()) + rec.put("summary", rec.optString("message_on_screen").lineSequence().firstOrNull() ?: "") + File(outDir, "$id.json").writeText(rec.toString(1)) + return rec + } + + private fun snippet(row: JSONObject): JSONObject { + val id = row.getString("id") + val fx = Fixture(File(activity.filesDir, row.getString("fixture"))) + val photo = Photo.fromFile(fx.image) + val rgb = Preprocess.resized(photo.image).rgb() + val image = Matting.birefnetInput(rgb) + val dir = ModelFiles.dir(activity.filesDir) + val t0 = SystemClock.elapsedRealtimeNanos() + val line = activity.onModel { + // the card's block builds its own environment and graph; free the engine's graphs first + Idol.engine?.pool?.releaseAll() + val bir = cardSnippet(dir, image) + if (!fx.has("pixels") || !fx.has("feature1")) return@onModel bir + // the DINOv3 card's block on the fixture's preprocessed picture, against the reference's feature1 + val got = dinoCardSnippet(dir, fx.floats("pixels")) + val want = fx.floats("feature1") + var dot = 0.0 + var ng = 0.0 + var nw = 0.0 + for (i in want.indices) { + dot += got[i].toDouble() * want[i] + ng += got[i].toDouble() * got[i] + nw += want[i].toDouble() * want[i] + } + "$bir | dino card snippet: ${got.size} values, cos vs the reference feature1 %.6f" + .format(dot / Math.sqrt(ng * nw)) + } + Log.i(TAG, "IDOL_AUTORUN SNIPPET $line") + val rec = JSONObject().put("id", id).put("kind", "snippet").put("status", "done").put("line", line) + .put("picture", JSONArray(listOf(rgb.width, rgb.height))) + .put("seconds", (SystemClock.elapsedRealtimeNanos() - t0) / 1e9).put("summary", line) + File(outDir, "$id.json").writeText(rec.toString(1)) + return rec + } + + companion object { + const val TAG = "IdolAutoRun" + } +} diff --git a/android/app/src/main/java/com/idol/CardSnippet.kt b/android/app/src/main/java/com/idol/CardSnippet.kt new file mode 100644 index 0000000000000000000000000000000000000000..31331ee9d603786c1f49eb76b3d545379765e874 --- /dev/null +++ b/android/app/src/main/java/com/idol/CardSnippet.kt @@ -0,0 +1,75 @@ +package com.idol + +import com.google.ai.edge.litert.Accelerator +import com.google.ai.edge.litert.CompiledModel +import com.google.ai.edge.litert.Environment +import java.io.File + +/** + * The model card's Kotlin block, run on the phone as written (autorun plan plans/snippet.json). [dir] is the app's + * model folder; [image] is the card's `image`: Matting.birefnetInput(the picture resized to a short side of 1024), + * [1, 3, 1024, 1024]. Returns one line about the alpha it produced. + */ +fun cardSnippet(dir: File, image: FloatArray): String { + // ---- the card's block, verbatim + val env = Environment.create() + val cpu = CompiledModel.Options(Accelerator.CPU) // bir.tflite: BiRefNet, the one CPU graph + val bir = CompiledModel.create(File(dir, "bir.tflite").absolutePath, cpu, env) + val inputs = bir.createInputBuffers() + val outputs = bir.createOutputBuffers() + inputs[0].writeFloat(image) // [1, 3, 1024, 1024]: RGB / 255, ImageNet mean / std + bir.run(inputs, outputs) + val alpha = outputs[0].readFloat() // [1024 * 1024] sigmoid map, 1 = the object + bir.close() + val gpu = CompiledModel.Options(Accelerator.GPU).apply { // every other graph + gpuOptions = CompiledModel.GpuOptions(precision = CompiledModel.GpuOptions.Precision.FP16_WITH_FP32_ACCUM) + } + // ---- end of the card's block + for (b in inputs + outputs) runCatching { b.close() } + runCatching { env.close() } + var mn = Float.POSITIVE_INFINITY + var mx = Float.NEGATIVE_INFINITY + var sum = 0.0 + var above = 0 + for (v in alpha) { + if (v < mn) mn = v + if (v > mx) mx = v + sum += v + if (v > 0.5f) above++ + } + return "card snippet: alpha ${alpha.size} values, min %.4f max %.4f mean %.4f, fraction > 0.5 %.4f; gpu options built (%s)" + .format(mn, mx, sum / alpha.size, above.toDouble() / alpha.size, gpu.gpuOptions?.precision) +} + +/** + * The DINOv3 repo card's Kotlin block, run on the phone as written (autorun plan plans/snippet.json, after + * [cardSnippet]). [dir] holds dino_0..3.tflite and the RoPE tables; [pixels] is the card's `pixels`: the picture Idol + * feeds, [1, 3, 1024, 1024] RGB in [0, 1]. Returns the feature (for the check against the fixture's feature1). + */ +fun dinoCardSnippet(dir: File, pixels: FloatArray): FloatArray { + // ---- the DINOv3 card's block, verbatim + fun floats(f: File): FloatArray { + val b = java.nio.ByteBuffer.wrap(f.readBytes()).order(java.nio.ByteOrder.LITTLE_ENDIAN).asFloatBuffer() + return FloatArray(b.remaining()).also { b.get(it) } + } + val env = Environment.create() + val gpu = CompiledModel.Options(Accelerator.GPU).apply { + gpuOptions = CompiledModel.GpuOptions(precision = CompiledModel.GpuOptions.Precision.FP16_WITH_FP32_ACCUM) + } + val rope = mapOf("args_1" to floats(File(dir, "dino_rope_cos.f32")), "args_2" to floats(File(dir, "dino_rope_sin.f32"))) + var x = pixels // [1, 3, 1024, 1024] RGB in [0, 1] + for (i in 0..3) { // 8 layers per graph + val model = CompiledModel.create(File(dir, "dino_$i.tflite").absolutePath, gpu, env) + val inputs = listOf("args_0", "args_1", "args_2").associateWith { model.createInputBuffer(it, "serving_default") } + val output = model.createOutputBuffer("output_0", "serving_default") + for ((name, buffer) in inputs) buffer.writeFloat(if (name == "args_0") x else rope.getValue(name)) + model.run(inputs, mapOf("output_0" to output), "serving_default") + x = output.readFloat() // [1, 4101, 1280] + for (b in inputs.values + output) b.close() + model.close() + } + env.close() + // x: the image feature Idol's denoiser reads (1 class, 4 register and 64 x 64 patch tokens) + // ---- end of the DINOv3 card's block + return x +} diff --git a/android/app/src/main/java/com/idol/DeviceInfo.kt b/android/app/src/main/java/com/idol/DeviceInfo.kt new file mode 100644 index 0000000000000000000000000000000000000000..ec091338d1fbb6a92cb93a146fea310be0a17b40 --- /dev/null +++ b/android/app/src/main/java/com/idol/DeviceInfo.kt @@ -0,0 +1,102 @@ +package com.idol + +import android.content.Context +import android.os.Build +import android.os.PowerManager +import java.io.File + +/** Device facts for the status line and the run JSON. */ +object DeviceInfo { + /** `MemAvailable` of `/proc/meminfo` in kB (what the memory guard compares with), or -1. */ + fun memAvailableKb(): Long = meminfo("MemAvailable:") + + fun memTotalKb(): Long = meminfo("MemTotal:") + + private fun meminfo(key: String): Long = runCatching { + File("/proc/meminfo").useLines { lines -> + lines.first { it.startsWith(key) }.trim().split(Regex("\\s+"))[1].toLong() + } + }.getOrDefault(-1L) + + /** VmHWM / VmRSS / VmSwap of this process in kB (`/proc/self/status`). */ + fun selfStatusKb(): Map = runCatching { + File("/proc/self/status").readLines().mapNotNull { line -> + val key = line.substringBefore(':') + if (key in setOf("VmHWM", "VmRSS", "VmSwap")) key to line.substringAfter(':').trim().split(' ')[0].toLong() + else null + }.toMap() + }.getOrDefault(emptyMap()) + + /** + * Resets this process's VmHWM to its current RSS (`echo 5 > /proc/self/clear_refs`), so that VmHWM read after a run + * is that run's peak. Returns false when the system does not allow it (the VmHWM is then the process's peak). + */ + fun resetPeakRss(): Boolean = runCatching { + File("/proc/self/clear_refs").writeText("5") + true + }.getOrDefault(false) + + /** The system's GPU pages (`/sys/class/kgsl/kgsl/page_alloc`, bytes), or -1 when the file cannot be read. */ + fun gpuPageAllocBytes(): Long = runCatching { + File("/sys/class/kgsl/kgsl/page_alloc").readText().trim().split(Regex("\\s+"))[0].toLong() + }.getOrDefault(-1L) + + /** "Galaxy S26 (SM-S942Q)" for the phones this sample was measured on, else `Build.MODEL`. */ + fun displayName(): String { + val model = Build.MODEL + val market = MARKET_NAMES.entries.firstOrNull { model.startsWith(it.key) }?.value + return if (market == null || market == model) model else "$market ($model)" + } + + fun thermalStatus(context: Context): Int = + if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q) { + context.getSystemService(PowerManager::class.java).currentThermalStatus + } else { + -1 + } + + /** The cpuset of this process (`/top-app` while the app is in front). */ + fun cpuset(): String = runCatching { + File("/proc/self/cgroup").readLines().firstOrNull { ":cpuset:" in it } ?: "none" + }.getOrDefault("unreadable") + + private val MARKET_NAMES = mapOf("SM-S942" to "Galaxy S26") +} + +/** + * Samples the system's GPU pages every [periodMs] while a run is in progress and keeps the maximum (CONTRACT.md + * line (e)). Best effort: when the app cannot read the file, [maxBytes] stays -1. + */ +class GpuPageSampler(private val periodMs: Long = 100) { + @Volatile var maxBytes = -1L + private set + @Volatile private var running = false + private var thread: Thread? = null + + fun start() { + maxBytes = DeviceInfo.gpuPageAllocBytes() + if (maxBytes < 0) return + running = true + thread = Thread({ + while (running) { + val v = DeviceInfo.gpuPageAllocBytes() + if (v > maxBytes) maxBytes = v + try { + Thread.sleep(periodMs) + } catch (_: InterruptedException) { + break + } + } + }, "IdolGpuPages").apply { isDaemon = true; start() } + } + + fun stop(): Long { + running = false + thread?.interrupt() + thread?.join(500) + thread = null + val v = DeviceInfo.gpuPageAllocBytes() + if (v > maxBytes) maxBytes = v + return maxBytes + } +} diff --git a/android/app/src/main/java/com/idol/Fixture.kt b/android/app/src/main/java/com/idol/Fixture.kt new file mode 100644 index 0000000000000000000000000000000000000000..d36bc453e33f73384f61d9aaa318f3ec80863bc4 --- /dev/null +++ b/android/app/src/main/java/com/idol/Fixture.kt @@ -0,0 +1,66 @@ +package com.idol + +import org.json.JSONObject +import java.io.File +import java.io.FileOutputStream +import java.io.RandomAccessFile +import java.nio.ByteBuffer +import java.nio.ByteOrder +import java.nio.channels.FileChannel + +/** + * A fixture directory `files/fx//` as scripts/make_fixtures.py writes it: fixture.json + * {"tag", "mode", "seed", "image", "arrays": {name: {"file", "shape"}}}, the raw input image and little-endian + * float32 arrays (the reference run's shared randomness and each part's own reference input). + */ +class Fixture(val dir: File) { + private val json = JSONObject(File(dir, "fixture.json").readText()) + val tag: String = json.optString("tag", dir.name) + val mode: String = json.getString("mode") + val seed: Long = json.optLong("seed", 0L) + val image: File = File(dir, json.getString("image")) + private val arrays = json.getJSONObject("arrays") + + fun has(name: String) = arrays.has(name) + + fun shape(name: String): IntArray { + val a = arrays.getJSONObject(name).getJSONArray("shape") + return IntArray(a.length()) { a.getInt(it) } + } + + fun floats(name: String): FloatArray { + val e = arrays.getJSONObject(name) + val shape = shape(name) + val n = shape.fold(1) { a, b -> a * b } + return readF32(File(dir, e.getString("file")), n) + } + + /** The five noise arrays, in place of the app's random draws. */ + fun noise(): Noise = Noise( + vae = floats("vae_noise"), + latent = floats("latent_noise"), + camera = floats("camera_noise"), + pool = floats("octree_pool"), + leaf = floats("leaf_pool"), + source = "fixture ${dir.name}", + ) + + companion object { + /** [n] little-endian float32 values from [file] (its size must be n x 4 bytes). */ + fun readF32(file: File, n: Int): FloatArray { + require(file.length() == n * 4L) { "${file.name} is ${file.length()} bytes, expected ${n * 4L}" } + val out = FloatArray(n) + RandomAccessFile(file, "r").use { raf -> + val map = raf.channel.map(FileChannel.MapMode.READ_ONLY, 0, n * 4L) + map.order(ByteOrder.LITTLE_ENDIAN).asFloatBuffer().get(out) + } + return out + } + + fun writeF32(file: File, data: FloatArray) { + val buf = ByteBuffer.allocate(data.size * 4).order(ByteOrder.LITTLE_ENDIAN) + buf.asFloatBuffer().put(data) + FileOutputStream(file).use { it.write(buf.array()) } + } + } +} diff --git a/android/app/src/main/java/com/idol/Glue.kt b/android/app/src/main/java/com/idol/Glue.kt new file mode 100644 index 0000000000000000000000000000000000000000..5c65f3fe207ce2c38c5e2050882efbd768024547 --- /dev/null +++ b/android/app/src/main/java/com/idol/Glue.kt @@ -0,0 +1,472 @@ +package com.idol + +import org.json.JSONObject +import java.io.BufferedOutputStream +import java.io.File +import java.io.FileOutputStream +import java.nio.ByteBuffer +import java.nio.ByteOrder +import kotlin.math.PI +import kotlin.math.cos +import kotlin.math.exp +import kotlin.math.expm1 +import kotlin.math.ln +import kotlin.math.ln1p +import kotlin.math.pow +import kotlin.math.sin +import kotlin.math.sqrt +import kotlin.math.tanh + +/** The five arrays of shared randomness (CONTRACT.md "Shared randomness"), flat in C order. */ +class Noise( + /** [1, 32, 128, 128] normal: the VAE posterior noise. */ + val vae: FloatArray, + /** [1, 8192, 16] normal: x at t = 1. */ + val latent: FloatArray, + /** [1, 1, 5] normal: the camera at t = 1. */ + val camera: FloatArray, + /** [8, 8192] uniform: the octree sampler's per-row offsets. */ + val pool: FloatArray, + /** [1, 8192, 3] uniform: the leaf positions. */ + val leaf: FloatArray, + /** "seed " or the fixture it came from. */ + val source: String, +) { + init { + require(vae.size == 32 * 128 * 128 && latent.size == Glue.NUM_POINTS * Glue.IN_CH && camera.size == Glue.CAM_CH) + require(pool.size == Glue.MAX_LEVEL * Glue.NUM_POINTS && leaf.size == Glue.NUM_POINTS * 3) + } + + companion object { + /** + * The app's own draws for a run from a picture: same distributions as the reference (normal, uniform [0, 1)), + * from java.util.Random, not torch's generator; the fixtures carry the reference's own numbers. + */ + fun random(seed: Long): Noise { + val r = java.util.Random(seed) + fun normal(n: Int) = FloatArray(n) { r.nextGaussian().toFloat() } + fun uniform(n: Int) = FloatArray(n) { r.nextFloat() } + val vae = normal(32 * 128 * 128) + val latent = normal(Glue.NUM_POINTS * Glue.IN_CH) + val camera = normal(Glue.CAM_CH) + val pool = uniform(Glue.MAX_LEVEL * Glue.NUM_POINTS) + val leaf = uniform(Glue.NUM_POINTS * 3) + return Noise(vae, latent, camera, pool, leaf, "seed $seed") + } + } +} + +/** Activated Gaussians (z-up, as GaussianFixedlenDecoder.to_representation returns them), N = 8192 x 32. */ +class Gaussians( + val count: Int, + val xyz: FloatArray, + val featuresDc: FloatArray, + val scaling: FloatArray, + /** (w, x, y, z) */ + val rotation: FloatArray, + val opacity: FloatArray, +) + +/** gs_rep.json: the Gaussian decoder's representation constants. */ +class GaussianRep( + val baseOffsetScale: Float, + val lrXyz: Float, + val lrFeaturesDc: Float, + val lrOpacity: Float, + val lrScaling: Float, + val lrRotation: Float, + val perturbeSize: Float, + val filterKernelSize: Double, + val scalingBias: Float, + val opacityBias: Float, +) { + companion object { + fun parse(o: JSONObject): GaussianRep { + val rep = o.getJSONObject("rep") + val lr = rep.getJSONObject("lr") + return GaussianRep( + baseOffsetScale = o.getDouble("base_offset_scale").toFloat(), + lrXyz = lr.getDouble("_xyz").toFloat(), + lrFeaturesDc = lr.getDouble("_features_dc").toFloat(), + lrOpacity = lr.getDouble("_opacity").toFloat(), + lrScaling = lr.getDouble("_scaling").toFloat(), + lrRotation = lr.getDouble("_rotation").toFloat(), + perturbeSize = rep.getDouble("perturbe_size").toFloat(), + filterKernelSize = rep.getDouble("3d_filter_kernel_size"), + scalingBias = rep.getDouble("scaling_bias").toFloat(), + opacityBias = rep.getDouble("opacity_bias").toFloat(), + ) + } + } +} + +/** + * The host glue between the graphs, ported line by line from scripts/idol_glue.py, scripts/idol_graphs.py + * (sinusoid, point_pe, lora_gate) and scripts/idol_runner.py (vae_tokens): the same order of float32 operations as + * the torch / numpy code (each transcendental is computed in double and rounded to float32). + */ +object Glue { + const val MAX_LEVEL = 8 + const val PER_ANCHOR = 32 + const val PHASE_THRESHOLD = 0.5 + const val NUM_POINTS = 8192 + const val IN_CH = 16 + const val CAM_CH = 5 + const val COND_TOKENS = 4101 + const val LORA_SCALE = 128.0f / 32.0f + val OCT_BUCKETS = intArrayOf(64, 256, 1024, 4096, 8192) + const val PLY_FLOATS = 17 + val PLY_NAMES = listOf( + "x", "y", "z", "nx", "ny", "nz", "f_dc_0", "f_dc_1", "f_dc_2", "opacity", + "scale_0", "scale_1", "scale_2", "rot_0", "rot_1", "rot_2", "rot_3", + ) + + private val PI_F = PI.toFloat() + + private fun expf(x: Float) = exp(x.toDouble()).toFloat() + private fun logf(x: Float) = ln(x.toDouble()).toFloat() + + /** time_grid: t = r s / (1 + (r - 1) s), s = numpy.linspace(1, 0, steps + 1), float64. */ + fun timeGrid(steps: Int = 5, rescale: Double = 3.0): DoubleArray { + val num = steps + 1 + val step = (0.0 - 1.0) / (num - 1) + val s = DoubleArray(num) { i -> i * step + 1.0 } + s[num - 1] = 0.0 + return DoubleArray(num) { rescale * s[it] / (1 + (rescale - 1) * s[it]) } + } + + /** lora_gate: [1, 64], the slow expert's half (t >= 0.5) or the fast one's at LORA_SCALE. */ + fun loraGate(slow: Boolean): FloatArray = FloatArray(64) { i -> + if ((i < 32) == slow) LORA_SCALE else 0f + } + + /** sinusoid: SinusoidalEmbedder's frequency features of one scalar -> [1, dim]. */ + fun sinusoid(value: Double, dim: Int = 256, maxPeriod: Double = 10000.0, twoPi: Boolean = false): FloatArray { + val half = dim / 2 + val c = (-ln(maxPeriod)).toFloat() + val v = value.toFloat() + val out = FloatArray(dim) + for (k in 0 until half) { + val f = expf((c * k.toFloat()) / half.toFloat()) + var a = v * f + if (twoPi) { + a *= 2f + a *= PI_F + } + out[k] = cos(a.toDouble()).toFloat() + out[half + k] = sin(a.toDouble()).toFloat() + } + return out + } + + /** + * torch.linspace(start, end, steps) in float32 as the reference runs' torch computes it: start + step * i below + * halfway, end - step * (steps - 1 - i) from there, each one fused multiply-add (one rounding). Only called with + * (0, 10, 170): there the product and the sum are exact in double, so the double expression rounds once too. + */ + private fun linspaceF(start: Float, end: Float, steps: Int): FloatArray { + val step = (end - start) / (steps - 1).toFloat() + val halfway = steps / 2 + return FloatArray(steps) { i -> + if (i < halfway) (step.toDouble() * i + start).toFloat() else (end - step.toDouble() * (steps - i - 1)).toFloat() + } + } + + /** + * The frequencies where torch's float32 pow(2, x) is one ulp off the correctly rounded value, read from the Mac + * reference run's point embedding (device/fx/g1_gs_0): with them, all 170 frequency columns of that embedding agree + * to one ulp of sin / cos (without them six columns are off by up to 1.1e-3). + */ + private val TORCH_POW_ULP = mapOf(43 to -1, 72 to -1, 95 to -1, 113 to 1, 138 to -1, 159 to 1) + + /** torch.pow(2.0, torch.linspace(0.0, 10.0, 170)), float32. */ + private val PE_FREQS: FloatArray by lazy { + val lin = linspaceF(0f, 10f, 1024 / 3 / 2) + FloatArray(lin.size) { + val f = 2.0.pow(lin[it].toDouble()).toFloat() + Float.fromBits(f.toRawBits() + (TORCH_POW_ULP[it] ?: 0)) + } + } + + /** point_pe: PcdAbsolutePositionEmbedderv2 of [n] points (x, y, z) -> [n, 1024]. */ + fun pointPe(points: FloatArray, n: Int, channels: Int = 1024): FloatArray { + val freqs = PE_FREQS + val fd = freqs.size + val out = FloatArray(n * channels) + for (p in 0 until n) { + for (c in 0 until 3) { + val v = points[p * 3 + c] + val base = p * channels + c * 2 * fd + for (j in 0 until fd) { + val ang = (v * freqs[j]) * PI_F + out[base + j] = sin(ang.toDouble()).toFloat() + out[base + fd + j] = cos(ang.toDouble()).toFloat() + } + } + } + return out + } + + /** + * vae_tokens: posterior sample with the shared noise (logvar clipped to [-30, 20]), 2 x 2 patchify, bn -> + * [1, 4096, 128]. moments [1, 64, 128, 128] (mean, logvar), noise [1, 32, 128, 128]. + */ + fun vaeTokens(moments: FloatArray, noise: FloatArray, bnMean: FloatArray, bnStd: FloatArray): FloatArray { + val hw = 128 + val plane = hw * hw + val out = FloatArray(4096 * 128) + for (c in 0 until 32) for (y in 0 until hw) for (x in 0 until hw) { + val i = c * plane + y * hw + x + var lv = moments[32 * plane + i] + lv = if (lv < -30f) -30f else if (lv > 20f) 20f else lv + val lat = moments[i] + expf(0.5f * lv) * noise[i] + val k = c * 4 + (y % 2) * 2 + (x % 2) + val t = (y / 2) * 64 + (x / 2) + out[t * 128 + k] = (lat - bnMean[k]) / bnStd[k] + } + return out + } + + /** torch.softmax over the last dimension (8) of [rows, 8] logits: exp(x - max) times 1 / sum. */ + fun softmax8(logits: FloatArray, rows: Int): FloatArray { + val out = FloatArray(rows * 8) + for (r in 0 until rows) { + val b = r * 8 + var mx = logits[b] + for (k in 1 until 8) if (logits[b + k] > mx) mx = logits[b + k] + var sum = 0f + for (k in 0 until 8) { + val e = expf(logits[b + k] - mx) + out[b + k] = e + sum += e + } + val inv = 1f / sum + for (k in 0 until 8) out[b + k] = out[b + k] * inv + } + return out + } + + /** + * systematic_rows: idol's systematic_sample for one batch, row i drawing counts[i] samples from its 8 child + * probabilities with the offset u[i] / counts[i] -> [rows, 8] sample counts. The cumulative sum accumulates in + * double as torch's CPU cumsum does for float32. + */ + fun systematicRows(probs: FloatArray, counts: IntArray, u: FloatArray, rows: Int): IntArray { + val P = 8 + val out = IntArray(rows * P) + val p = FloatArray(P) + val cdf = FloatArray(P) + for (r in 0 until rows) { + val n = counts[r] + if (n == 0) continue + var s = 0f + for (k in 0 until P) { + val v = probs[r * P + k] + p[k] = if (v < 0f) 0f else v + s += p[k] + } + if (s == 0f) { + for (k in 0 until P) p[k] = 1.0f / P + } else { + val d = if (s < 1f) 1f else s + for (k in 0 until P) p[k] = p[k] / d + } + var acc = 0.0 + for (k in 0 until P) { + acc += p[k] + val c = acc.toFloat() + cdf[k] = if (c > 1f) 1f else c + } + val nf = n.toFloat() + val u0 = u[r] / nf + for (j in 0 until n) { + var us = u0 + j.toFloat() / nf + if (us > 1f) us = 1f + var idx = 0 + while (idx < P && cdf[idx] < us) idx++ + if (idx > P - 1) idx = P - 1 + out[r * P + idx]++ + } + } + return out + } + + /** + * octree_descend: idol OctreeProbabilityFixedlenDecoder.sample for one batch. [probsFn] (cells [rows, 3], rows, + * level) -> child probabilities [rows, 8]. pool [8, 8192], leaf [8192, 3] -> anchors [8192, 3] in [0, 1]^3. + * [cellsOut], when given, receives each level's cells (the parity record). + */ + fun octreeDescend( + probsFn: (FloatArray, Int, Int) -> FloatArray, + pool: FloatArray, + leaf: FloatArray, + numPoints: Int = NUM_POINTS, + levels: Int = MAX_LEVEL, + ): FloatArray { + var prev = IntArray(3) + var rows = 1 + var counts = intArrayOf(numPoints) + for (l in 1..levels) { + val res = (1 shl (l - 1)).toFloat() + val cells = FloatArray(rows * 3) { (prev[it].toFloat() + 0.5f) / res } + val probs = probsFn(cells, rows, l) + val u = FloatArray(rows) { pool[(l - 1) * numPoints + it] } + val sampled = systematicRows(probs, counts, u, rows) + var m = 0 + for (v in sampled) if (v > 0) m++ + val next = IntArray(m * 3) + val nextCounts = IntArray(m) + var o = 0 + for (r in 0 until rows) for (c in 0 until 8) { + val s = sampled[r * 8 + c] + if (s <= 0) continue + // CHILD_OFFSET[c] = (i, j, k) with c = k * 4 + j * 2 + i + next[o * 3] = prev[r * 3] * 2 + (c and 1) + next[o * 3 + 1] = prev[r * 3 + 1] * 2 + ((c shr 1) and 1) + next[o * 3 + 2] = prev[r * 3 + 2] * 2 + ((c shr 2) and 1) + nextCounts[o] = s + o++ + } + prev = next + counts = nextCounts + rows = m + } + val out = FloatArray(numPoints * 3) + val scale = (1 shl levels).toFloat() + var p = 0 + for (r in 0 until rows) for (k in 0 until counts[r]) { + check(p < numPoints) { "octree sampling produced more than $numPoints points" } + for (d in 0 until 3) out[p * 3 + d] = (prev[r * 3 + d].toFloat() + leaf[p * 3 + d]) / scale + p++ + } + check(p == numPoints) { "octree sampling produced $p points, not $numPoints" } + return out + } + + private fun softplus(x: Float): Float = if (x > 20f) x else ln1p(expf(x).toDouble()).toFloat() + + private fun sigmoid(x: Float): Float = 1f / (1f + expf(-x)) + + /** + * to_gaussians: GaussianFixedlenDecoder.to_representation + Gaussian.__init__ for one batch. anchors [L, 3], + * features [L, 480] (per anchor: xyz 32x3, f_dc 32x3, scaling 32x3, rotation 32x4, opacity 32, offset scale 32), + * perturbation [32, 3]. + */ + fun toGaussians(anchors: FloatArray, features: FloatArray, rep: GaussianRep, perturbation: FloatArray): Gaussians { + val L = anchors.size / 3 + val n = L * PER_ANCHOR + val width = PER_ANCHOR * (3 + 3 + 3 + 4 + 1 + 1) + require(features.size == L * width) { "features hold ${features.size} values, not $L x $width" } + val xyz = FloatArray(n * 3) + val fdc = FloatArray(n * 3) + val scaling = FloatArray(n * 3) + val rotation = FloatArray(n * 4) + val opacity = FloatArray(n) + // inverse_softplus(scaling_bias), inverse_sigmoid(opacity_bias): float32 tensors in the reference + val sb = rep.scalingBias + val scalingOffset = sb + logf(-expm1(-sb.toDouble()).toFloat()) + val ob = rep.opacityBias + val opacityOffset = logf(ob / (1f - ob)) + val kernel2 = (rep.filterKernelSize * rep.filterKernelSize).toFloat() + val xyzAt = 0 + val fdcAt = PER_ANCHOR * 3 + val scaleAt = PER_ANCHOR * 6 + val rotAt = PER_ANCHOR * 9 + val opaAt = PER_ANCHOR * 13 + val oscAt = PER_ANCHOR * 14 + for (a in 0 until L) { + val f = a * width + for (j in 0 until PER_ANCHOR) { + val g = a * PER_ANCHOR + j + val offsetScale = softplus(features[f + oscAt + j] + rep.baseOffsetScale) + for (c in 0 until 3) { + var off = features[f + xyzAt + j * 3 + c] * rep.lrXyz + perturbation[j * 3 + c] + off = tanh(off.toDouble()).toFloat() * 0.5f * rep.perturbeSize * offsetScale + val raw = off + anchors[a * 3 + c] + xyz[g * 3 + c] = raw * 1.0f - 0.5f + fdc[g * 3 + c] = features[f + fdcAt + j * 3 + c] * rep.lrFeaturesDc + val s = softplus(features[f + scaleAt + j * 3 + c] * rep.lrScaling + scalingOffset) + scaling[g * 3 + c] = sqrt((s * s + kernel2).toDouble()).toFloat() + } + val r0 = features[f + rotAt + j * 4] * rep.lrRotation + 1.0f + val r1 = features[f + rotAt + j * 4 + 1] * rep.lrRotation + 0.0f + val r2 = features[f + rotAt + j * 4 + 2] * rep.lrRotation + 0.0f + val r3 = features[f + rotAt + j * 4 + 3] * rep.lrRotation + 0.0f + var ss = r0 * r0 + ss += r1 * r1 + ss += r2 * r2 + ss += r3 * r3 + val norm = maxOf(sqrt(ss.toDouble()).toFloat(), 1e-12f) + rotation[g * 4] = r0 / norm + rotation[g * 4 + 1] = r1 / norm + rotation[g * 4 + 2] = r2 / norm + rotation[g * 4 + 3] = r3 / norm + opacity[g] = sigmoid(features[f + opaAt + j] * rep.lrOpacity + opacityOffset) + } + } + return Gaussians(n, xyz, fdc, scaling, rotation, opacity) + } + + /** ply_rows: idol Gaussian.save_ply's 17 attributes per Gaussian (z-up turned to y-up), float32. */ + fun plyRows(g: Gaussians): FloatArray { + val n = g.count + val rows = FloatArray(n * PLY_FLOATS) + val eps = 1e-6f + val hi = (1.0 - 1e-6).toFloat() + // quaternion_multiply([sqrt(.5), sqrt(.5), 0, 0], rotation) in float64 (numpy promotes), rounded at the end + val aw = sqrt(0.5) + val ax = sqrt(0.5) + val ay = 0.0 + val az = 0.0 + for (i in 0 until n) { + val o = i * PLY_FLOATS + rows[o] = g.xyz[i * 3] + rows[o + 1] = -g.xyz[i * 3 + 2] + rows[o + 2] = g.xyz[i * 3 + 1] + rows[o + 6] = g.featuresDc[i * 3] + rows[o + 7] = g.featuresDc[i * 3 + 1] + rows[o + 8] = g.featuresDc[i * 3 + 2] + var op = g.opacity[i] + op = if (op < eps) eps else if (op > hi) hi else op + rows[o + 9] = logf(op / (1f - op)) + for (c in 0 until 3) rows[o + 10 + c] = logf(maxOf(g.scaling[i * 3 + c], eps)) + val bw = g.rotation[i * 4].toDouble() + val bx = g.rotation[i * 4 + 1].toDouble() + val by = g.rotation[i * 4 + 2].toDouble() + val bz = g.rotation[i * 4 + 3].toDouble() + rows[o + 13] = (aw * bw - ax * bx - ay * by - az * bz).toFloat() + rows[o + 14] = (aw * bx + ax * bw + ay * bz - az * by).toFloat() + rows[o + 15] = (aw * by - ax * bz + ay * bw + az * bx).toFloat() + rows[o + 16] = (aw * bz + ax * by - ay * bx + az * bw).toFloat() + } + return rows + } + + /** write_ply: binary little-endian PLY with the 17 float properties (the bytes plyfile writes). */ + fun writePly(rows: FloatArray, file: File) { + val n = rows.size / PLY_FLOATS + val header = buildString { + append("ply\nformat binary_little_endian 1.0\nelement vertex ").append(n).append('\n') + for (name in PLY_NAMES) append("property float ").append(name).append('\n') + append("end_header\n") + } + val tmp = File(file.parentFile, file.name + ".tmp") + BufferedOutputStream(FileOutputStream(tmp), 1 shl 20).use { out -> + out.write(header.toByteArray(Charsets.US_ASCII)) + val chunk = 4096 * PLY_FLOATS + val buf = ByteBuffer.allocate(chunk * 4).order(ByteOrder.LITTLE_ENDIAN) + var i = 0 + while (i < rows.size) { + val m = minOf(chunk, rows.size - i) + buf.clear() + buf.asFloatBuffer().put(rows, i, m) + out.write(buf.array(), 0, m * 4) + i += m + } + } + if (!tmp.renameTo(file)) { + tmp.copyTo(file, overwrite = true) + tmp.delete() + } + } +} diff --git a/android/app/src/main/java/com/idol/Graphs.kt b/android/app/src/main/java/com/idol/Graphs.kt new file mode 100644 index 0000000000000000000000000000000000000000..196ccc23470bbaf5a240ab80f51071c0581ecbaa --- /dev/null +++ b/android/app/src/main/java/com/idol/Graphs.kt @@ -0,0 +1,339 @@ +package com.idol + +import android.os.SystemClock +import android.util.Log +import com.google.ai.edge.litert.Accelerator +import com.google.ai.edge.litert.CompiledModel +import com.google.ai.edge.litert.Environment +import com.google.ai.edge.litert.TensorBuffer +import org.json.JSONObject +import java.io.File + +/** + * One compiled graph with its I/O buffers, made once per signature and reused. Inputs written with a tag are not + * written again while the buffer still holds that tag (the RoPE tables, the position embedding, the LoRA gate, the + * conditioning of the current phase...). + */ +class Graph(val spec: GraphSpec, private val model: CompiledModel) : AutoCloseable { + val name: String get() = spec.name + private val inputs = HashMap>() + private val outputs = HashMap>() + private val tags = HashMap>() + var runs = 0 + private set + + init { + try { + for (sig in spec.signatures.values) { + inputs[sig.key] = sig.inputs.map { model.createInputBuffer(it.name, sig.key) } + outputs[sig.key] = sig.outputs.map { model.createOutputBuffer(it.name, sig.key) } + tags[sig.key] = arrayOfNulls(sig.inputs.size) + } + } catch (e: Throwable) { + close() + throw e + } + } + + private fun sig(key: String?): SignatureSpec = spec.signature(key) + + /** Writes [data] to input [index] of signature [key], unless the buffer already holds [tag]. */ + fun write(index: Int, data: FloatArray, tag: Any? = null, key: String? = null) { + val s = sig(key) + val t = tags.getValue(s.key) + if (tag != null && t[index] == tag) return + val want = s.inputs[index].size + require(data.size == want) { "$name ${s.key} input $index: ${data.size} values, expected $want" } + inputs.getValue(s.key)[index].writeFloat(data) + t[index] = tag + } + + /** + * Runs signature [key]. [overrides] replace input buffers by index (another graph's output buffer: the direct + * handoff); the graph's own buffer at that index keeps its content and loses its tag. + */ + fun run(key: String? = null, overrides: Map = emptyMap()) { + val s = sig(key) + val ins = inputs.getValue(s.key) + val outs = outputs.getValue(s.key) + val inMap = HashMap() + for ((i, t) in s.inputs.withIndex()) inMap[t.name] = overrides[i] ?: ins[i] + val outMap = HashMap() + for ((i, t) in s.outputs.withIndex()) outMap[t.name] = outs[i] + model.run(inMap, outMap, s.key) + runs++ + } + + /** Output [index] of signature [key] as float32 (waits for the GPU). */ + fun read(index: Int, key: String? = null): FloatArray { + val s = sig(key) + val want = s.outputs[index].size + val v = outputs.getValue(s.key)[index].readFloat() + check(v.size >= want) { "$name ${s.key} output $index: ${v.size} values, expected $want" } + return if (v.size == want) v else v.copyOf(want) + } + + fun outputBuffer(index: Int, key: String? = null): TensorBuffer = outputs.getValue(sig(key).key)[index] + + /** Forget the tag of input [index] (its buffer was not written by us). */ + fun untag(index: Int, key: String? = null) { + tags.getValue(sig(key).key)[index] = null + } + + override fun close() { + for (l in inputs.values) for (b in l) runCatching { b.close() } + for (l in outputs.values) for (b in l) runCatching { b.close() } + inputs.clear() + outputs.clear() + runCatching { model.close() } + } +} + +/** The memory guard found no graph left to close and the next one does not fit above the reserve. */ +class NotEnoughMemoryException(message: String) : Exception(message) + +/** + * The compiled graphs the app keeps, all on the model thread. All 46 graphs do not fit in a phone's memory at once + * (one denoiser block alone holds about 0.7 GB of GPU memory), so the pool keeps as many as fit above a reserve of + * `MemAvailable` and, when the next graph does not fit, closes the resident graph whose next use in the run's + * schedule is farthest away (Belady's rule over the fixed graph order of a run, read cyclically for the next run). + * Every compile, close and memory shortfall is recorded. + */ +class GraphPool( + private val env: Environment, + private val dir: File, + private val manifest: Manifest, + var reserveKb: Long, + var memoryCheck: Boolean, + var cpuThreads: Int, + var gpuPriority: CompiledModel.GpuOptions.Priority?, +) : AutoCloseable { + class CompileRecord( + val graph: String, + val seconds: Double, + val memAvailableBeforeKb: Long, + val memAvailableAfterKb: Long, + val phase: String, + ) { + fun json(): JSONObject = JSONObject().put("graph", graph).put("seconds", seconds) + .put("mem_available_before_kb", memAvailableBeforeKb).put("mem_available_after_kb", memAvailableAfterKb) + .put("phase", phase) + } + + private val resident = LinkedHashMap() + private val lastUse = HashMap() + private val pinned = HashSet() + private val footprintKb = HashMap() + private var schedule: List = emptyList() + private var cursor = -1 + private var clock = 0L + + /** What the compile records are filed under ("preload", "run", "parity", ...). */ + var phase = "run" + val compiles = ArrayList() + val events = ArrayList() + + val residentNames: List get() = resident.keys.toList() + + fun isResident(name: String) = name in resident + + /** The order the next run will ask for graphs in (resets the position in it). */ + fun setSchedule(order: List) { + schedule = order + cursor = -1 + } + + fun pin(name: String) { + pinned.add(name) + } + + fun unpin(name: String) { + pinned.remove(name) + } + + /** The graph [name], compiled now if it is not resident (closing others first when memory is short). */ + fun get(name: String): Graph { + advance(name) + lastUse[name] = clock++ + resident[name]?.let { return it } + makeRoom(name, evict = true) + return compile(name) + } + + /** Compiles [name] if it fits without closing another graph; false when it does not fit. */ + fun preload(name: String): Boolean { + if (name in resident) return true + val need = estimateKb(name) + reserveKb + val avail = DeviceInfo.memAvailableKb() + if (avail in 0 until need) { + event("preload_stop", name, avail, need) + return false + } + lastUse[name] = clock++ + compile(name) + return true + } + + fun release(name: String) { + resident.remove(name)?.let { + it.close() + event("close", name, DeviceInfo.memAvailableKb(), null) + } + } + + fun releaseIf(predicate: (String) -> Boolean) { + for (n in resident.keys.filter(predicate)) release(n) + } + + fun releaseAll() = releaseIf { true } + + override fun close() = releaseAll() + + private fun advance(name: String) { + if (schedule.isEmpty()) return + for (i in cursor + 1 until schedule.size) if (schedule[i] == name) { + cursor = i + return + } + for (i in 0..minOf(cursor, schedule.size - 1)) if (schedule[i] == name) { + cursor = i + return + } + } + + /** Distance in the schedule to the next use of [name] after the current position (cyclic); MAX when never. */ + private fun nextUse(name: String): Int { + if (schedule.isEmpty()) return Int.MAX_VALUE + for (i in cursor + 1 until schedule.size) if (schedule[i] == name) return i - cursor + for (i in 0..minOf(cursor, schedule.size - 1)) if (schedule[i] == name) return schedule.size - cursor + i + return Int.MAX_VALUE + } + + private fun victim(exclude: String): String? { + val candidates = resident.keys.filter { it != exclude && it !in pinned } + if (candidates.isEmpty()) return null + if (schedule.isEmpty()) return candidates.minByOrNull { lastUse[it] ?: 0L } + return candidates.maxWithOrNull(compareBy({ nextUse(it) }, { -(lastUse[it] ?: 0L) })) + } + + /** + * Closes graphs until `MemAvailable` covers the estimate for [name] plus the reserve. With [memoryCheck], fails + * when nothing is left to close and even the reserve is not there (or the measured size of [name] is not). + */ + private fun makeRoom(name: String, evict: Boolean) { + val est = estimateKb(name) + while (true) { + val avail = DeviceInfo.memAvailableKb() + if (avail < 0 || avail >= est + reserveKb) return + val v = if (evict) victim(name) else null + if (v == null) { + event("short", name, avail, est + reserveKb) + val measured = footprintKb[name] + val floor = reserveKb + (measured ?: 0L) + if (memoryCheck && avail < floor) { + val now = awaitAvailable(floor, MAX_WAIT_SHORT_MS) + event("waited", name, now, floor) + if (now in 0 until floor) { + throw NotEnoughMemoryException( + "Not enough memory for $name: ${now / 1024} MB available, ${floor / 1024} MB needed " + + "(reserve ${reserveKb / 1024} MB). Close other apps and try again.", + ) + } + } + return + } + event("evict", v, avail, est + reserveKb, "for" to name) + val freed = footprintKb[v] + release(v) + // A closed graph can return its GPU memory a few seconds late (Pixel 11 Pro, PowerVR: MemAvailable rose + // from 1.6 to 8.9 GB within 8 s of the closes); without this wait the loop closed every resident graph. + if (freed != null && freed > 0) awaitAvailable(avail + freed / 2, MAX_WAIT_EVICT_MS) + } + } + + /** Polls `MemAvailable` until it reaches [targetKb] or [maxMs] pass; returns the last value read. */ + private fun awaitAvailable(targetKb: Long, maxMs: Long): Long { + val t0 = SystemClock.elapsedRealtime() + var now = DeviceInfo.memAvailableKb() + while (now in 0 until targetKb && SystemClock.elapsedRealtime() - t0 < maxMs) { + Thread.sleep(POLL_MS) + now = DeviceInfo.memAvailableKb() + } + return now + } + + private fun compile(name: String): Graph { + val spec = manifest.graph(name) + val file = File(dir, spec.file) + val options = if (spec.onCpu) { + CompiledModel.Options(Accelerator.CPU).apply { + cpuOptions = CompiledModel.CpuOptions(numThreads = cpuThreads) + } + } else { + CompiledModel.Options(Accelerator.GPU).apply { + gpuOptions = CompiledModel.GpuOptions( + // one copy of the weights for a graph with several signatures (the octree decoder's buckets) + constantTensorSharing = if (spec.signatures.size > 1) true else null, + precision = precisionOf(spec.precision), + priority = gpuPriority, + ) + } + } + val before = DeviceInfo.memAvailableKb() + val t0 = SystemClock.elapsedRealtimeNanos() + val model = CompiledModel.create(file.absolutePath, options, env) + val graph = try { + Graph(spec, model) + } catch (e: Throwable) { + runCatching { model.close() } + throw e + } + val seconds = (SystemClock.elapsedRealtimeNanos() - t0) / 1e9 + val after = DeviceInfo.memAvailableKb() + if (before >= 0 && after >= 0) footprintKb[name] = maxOf(footprintKb[name] ?: 0L, before - after) + resident[name] = graph + compiles.add(CompileRecord(name, seconds, before, after, phase)) + event("compile", name, after, null, "seconds" to seconds) + Log.i(TAG, "compiled $name in ${"%.2f".format(seconds)} s (${spec.accelerator}, ${spec.precision}), " + + "MemAvailable ${before / 1024} -> ${after / 1024} MB, resident ${resident.size}") + return graph + } + + /** + * The memory [name] is expected to take: what it took when it was last compiled, else an estimate from the + * file (weights at the run precision) and its largest tensor (a transformer block keeps a few dozen of those). + */ + fun estimateKb(name: String): Long { + footprintKb[name]?.let { if (it > 0) return it } + val spec = manifest.graph(name) + val elementBytes = if (spec.precision == "fp32" || spec.onCpu) 4L else 2L + val weights = if (elementBytes == 4L) spec.bytes else spec.bytes / 2 + val transformer = name.startsWith("dino_") || name.startsWith("den_") || name == "oct" || + name.startsWith("gs_") || name == "bir" + val activations = spec.maxIoElements * elementBytes * (if (transformer) 26 else 8) + return (weights + activations) / 1024 + } + + private fun event(kind: String, graph: String, availKb: Long, needKb: Long?, vararg extra: Pair) { + val e = JSONObject().put("t_ms", SystemClock.elapsedRealtime()).put("event", kind).put("graph", graph) + .put("mem_available_kb", availKb).put("resident", resident.size) + if (needKb != null) e.put("need_kb", needKb) + for ((k, v) in extra) e.put(k, v) + events.add(e) + if (kind != "compile") Log.i(TAG, "memory: $e") + } + + companion object { + const val TAG = "IdolGraphs" + private const val POLL_MS = 100L + private const val MAX_WAIT_EVICT_MS = 1_500L + private const val MAX_WAIT_SHORT_MS = 8_000L + + fun precisionOf(p: String): CompiledModel.GpuOptions.Precision = when (p) { + "fp32" -> CompiledModel.GpuOptions.Precision.FP32 + "fp16" -> CompiledModel.GpuOptions.Precision.FP16 + "fp16_with_fp32_accum" -> CompiledModel.GpuOptions.Precision.FP16_WITH_FP32_ACCUM + else -> CompiledModel.GpuOptions.Precision.DEFAULT + } + } +} diff --git a/android/app/src/main/java/com/idol/IdolEngine.kt b/android/app/src/main/java/com/idol/IdolEngine.kt new file mode 100644 index 0000000000000000000000000000000000000000..2f176e117711df7903f1a67211664476fed7e416 --- /dev/null +++ b/android/app/src/main/java/com/idol/IdolEngine.kt @@ -0,0 +1,606 @@ +package com.idol + +import android.os.SystemClock +import android.util.Log +import com.google.ai.edge.litert.CompiledModel +import com.google.ai.edge.litert.Environment +import com.google.ai.edge.litert.LiteRtException +import org.json.JSONArray +import org.json.JSONObject +import java.io.File +import java.util.concurrent.atomic.AtomicLong +import kotlin.math.sqrt + +/** The stages the Run screen lists, in order. */ +enum class Stage(val label: String) { + MATTING("Removing background"), + DINO("Image features (DINOv3)"), + VAE("Image latent (VAE)"), + DENOISER("Generating the 3D shape"), + OCTREE("Placing points"), + GAUSSIANS("Decoding Gaussians"), + PLY("Saving the .ply"), +} + +/** What the engine tells the screen while it works (called on the model thread). */ +interface Progress { + fun stage(stage: Stage, detail: String?, seconds: Double?, done: Boolean) {} + fun status(text: String) {} + fun cutout(composite: Image8) {} +} + +class CancelledException : Exception("cancelled") + +data class EngineOptions( + /** Compile the graphs before the first run and keep them (true), or per stage, closed after it (false). */ + val preload: Boolean = true, + /** + * MemAvailable the memory guard keeps free. MemAvailable counts the page cache of the mapped graph files, and with + * 1 GB kept a compile burst outran its reclaim: lmkd killed the app on the Galaxy S26 at about 1.9 GB available. + */ + val reserveMb: Long = 2500, + /** Fail a run with a message when a graph does not fit even after closing all others (instead of trying). */ + val memoryCheck: Boolean = true, + /** Pass a graph's output buffer straight to the next graph instead of copying it through the host. */ + val directHandoff: Boolean = false, + /** "bir" (BiRefNet on the CPU, the author's matting) or "dis" (IS-Net DIS on the GPU). */ + val matting: String = "bir", + val cpuThreads: Int = 6, + val gpuPriority: CompiledModel.GpuOptions.Priority? = null, + val disFloor: Int = Matting.DIS_ALPHA_FLOOR, +) + +class RunResult( + val ply: File, + val gaussians: Int, + /** Seconds: matting, preprocess (matting included, as idol_runner times it), dinov3, vae, denoiser, octree, + * gaussian_decoder, ply, generate (preprocessed pixels ready -> .ply written), all. */ + val times: LinkedHashMap, + val pre: Preprocessed, + val compiles: List, + val events: List, + val noiseSource: String, + val handoff: String, +) + +/** + * The phone pipeline, step for step scripts/idol_runner.py: matting -> preprocessing -> DINOv3 (4 graphs) + VAE + * (6 graphs) -> 5 Euler steps (den_ctx once per LoRA phase, den_in, 24 joint graphs) -> octree (8 levels, row + * buckets) -> Gaussian decoder (8 graphs) -> .ply. Every graph runs on its own GPU call (the 2 s rule of F0.md + * finding 2), every call on the one model thread. Create, use and close it on that thread. + */ +class IdolEngine(val modelDir: File, val manifest: Manifest, val options: EngineOptions) : AutoCloseable { + private val env: Environment = Environment.create() + val pool = GraphPool( + env, modelDir, manifest, options.reserveMb * 1024, options.memoryCheck, options.cpuThreads, options.gpuPriority, + ) + + private val posEmb = side("den_pos_emb.f32") + private val ropeCos = side("dino_rope_cos.f32") + private val ropeSin = side("dino_rope_sin.f32") + private val perturbation = side("gs_perturbation.f32") + private val bnMean: FloatArray + private val bnStd: FloatArray + private val rep: GaussianRep + private val gateSlow = Glue.loraGate(true) + private val gateFast = Glue.loraGate(false) + private val timeGrid = Glue.timeGrid() + + /** Set from the screen thread; the model thread stops at the next graph. */ + @Volatile var cancel = false + + /** Direct handoff, until the first failure (then host copies, and the reason is kept). */ + var direct = options.directHandoff + private set + var directFailure: String? = null + private set + + private var progress: Progress = object : Progress {} + + init { + val bn = JSONObject(File(modelDir, "vae_bn.json").readText()) + val mean = bn.getJSONArray("bn_mean") + val variance = bn.getJSONArray("bn_var") + val eps = bn.getDouble("bn_eps").toFloat() + bnMean = FloatArray(mean.length()) { mean.getDouble(it).toFloat() } + bnStd = FloatArray(variance.length()) { sqrt((variance.getDouble(it).toFloat() + eps).toDouble()).toFloat() } + rep = GaussianRep.parse(JSONObject(File(modelDir, "gs_rep.json").readText())) + } + + private fun side(name: String): FloatArray { + val s = requireNotNull(manifest.sideFiles[name]) { "manifest.json has no side file $name" } + return Fixture.readF32(File(modelDir, name), s.elements.toInt()) + } + + override fun close() { + pool.close() + runCatching { env.close() } + } + + // ------------------------------------------------------------------------------------------------ schedule + // Each chain's graphs in manifest order (manifest.json lists them in pipeline order: dino_0..3, vae_0..3, vae_attn, + // vae_4, den_joint_00..23, gs_0..7), so a stage split into more or fewer graphs needs a new manifest, not new code. + private fun chain(prefix: String): List = manifest.graphs.keys.filter { it.startsWith(prefix) }.also { + require(it.isNotEmpty()) { "manifest.json has no $prefix graphs" } + } + private val dinoNames = chain("dino_") + private val vaeNames = chain("vae_") + private val jointNames = chain("den_joint_") + private val gsNames = chain("gs_") + + /** The order a run asks for graphs in. */ + fun runSchedule(): List { + val s = ArrayList() + s.add(options.matting) + s.addAll(dinoNames) + s.addAll(vaeNames) + var phase: Boolean? = null + for (i in 0 until timeGrid.size - 1) { + val slow = timeGrid[i] >= Glue.PHASE_THRESHOLD + if (slow != phase) s.add("den_ctx") + phase = slow + s.add("den_in") + s.addAll(jointNames) + } + repeat(Glue.MAX_LEVEL) { s.add("oct") } + s.addAll(gsNames) + return s + } + + /** The GPU graphs of a run in first-use order (the CPU matting graph is built for each matting and closed). */ + fun preloadOrder(): List = runSchedule().distinct().filter { !manifest.graph(it).onCpu } + + /** + * preload = true: compiles the graphs before the first run, in first-use order, as many as fit above the reserve. + * Returns the load record (seconds per graph and the total). + */ + fun preload(progress: Progress): JSONObject { + cancel = false + val names = preloadOrder() + pool.phase = "preload" + pool.setSchedule(runSchedule()) + val first = pool.compiles.size + val t0 = SystemClock.elapsedRealtimeNanos() + var loaded = 0 + var stoppedAt: String? = null + try { + for (n in names) { + if (cancel) throw CancelledException() + progress.status("Loading models ${loaded + 1}/${names.size}") + if (!pool.preload(n)) { + stoppedAt = n + break + } + loaded++ + } + } finally { + pool.phase = "run" + } + val total = (SystemClock.elapsedRealtimeNanos() - t0) / 1e9 + progress.status( + if (stoppedAt == null) "Models loaded ($loaded/${names.size})" + else "Loaded $loaded/${names.size} models; the rest load during the run (memory)", + ) + val graphs = JSONArray() + for (r in pool.compiles.subList(first, pool.compiles.size)) graphs.put(r.json()) + return JSONObject().put("graphs", graphs).put("loaded", loaded).put("of", names.size) + .put("stopped_at", stoppedAt ?: JSONObject.NULL).put("total_s", total) + .put("compile_s", pool.compiles.subList(first, pool.compiles.size).sumOf { it.seconds }) + .put("resident", JSONArray(pool.residentNames)) + } + + // ------------------------------------------------------------------------------------------------ the run + /** A picture (or a fixture's raw image) -> .ply at [outPly]. */ + fun generate(image: Image8, mode: String, noise: Noise, outPly: File, progress: Progress): RunResult { + cancel = false + this.progress = progress + val firstCompile = pool.compiles.size + val firstEvent = pool.events.size + pool.phase = "run" + pool.setSchedule(runSchedule()) + val times = LinkedHashMap() + val tAll = now() + + progress.stage(Stage.MATTING, null, null, false) + var mattingS = 0.0 + val tPre = now() + val pre = Preprocess.run(image, mode) { rgb -> + val t = now() + val a = alpha(rgb) + mattingS = since(t) + a + } + times["matting"] = mattingS + times["preprocess"] = since(tPre) + pool.release(options.matting) + progress.cutout(pre.composite) + progress.stage(Stage.MATTING, null, times["preprocess"], true) + val pixels = Preprocess.pixels(pre.composite) + + val tGen = now() + progress.stage(Stage.DINO, null, null, false) + var t = now() + val f1 = dinov3(pixels) + times["dinov3"] = since(t) + progress.stage(Stage.DINO, null, times["dinov3"], true) + if (!options.preload) pool.releaseIf { it.startsWith("dino_") } + + progress.stage(Stage.VAE, null, null, false) + t = now() + val moments = vaeMoments(pixels) + times["vae"] = since(t) + val f2 = cond2(Glue.vaeTokens(moments, noise.vae, bnMean, bnStd)) + progress.stage(Stage.VAE, null, times["vae"], true) + if (!options.preload) pool.releaseIf { it.startsWith("vae_") } + + val tDen = now() + val latent = sample(noise, f1, f2, times) + progress.stage(Stage.DENOISER, null, since(tDen), true) + if (!options.preload) pool.releaseIf { it.startsWith("den_") } + + progress.stage(Stage.OCTREE, null, null, false) + t = now() + val latentTag = nextTag() + val anchors = Glue.octreeDescend({ cells, rows, level -> + progress.stage(Stage.OCTREE, "level $level of ${Glue.MAX_LEVEL}", null, false) + Glue.softmax8(octreeLogits(cells, rows, level, latent, latentTag), rows) + }, noise.pool, noise.leaf) + times["octree"] = since(t) + progress.stage(Stage.OCTREE, null, times["octree"], true) + if (!options.preload) pool.release("oct") + + progress.stage(Stage.GAUSSIANS, null, null, false) + t = now() + val features = gaussianFeatures(anchors, latent, latentTag) + times["gaussian_decoder"] = since(t) + if (!options.preload) pool.releaseIf { it.startsWith("gs_") } + val g = Glue.toGaussians(anchors, features, rep, perturbation) + progress.stage(Stage.GAUSSIANS, null, since(t), true) + + progress.stage(Stage.PLY, null, null, false) + t = now() + Glue.writePly(Glue.plyRows(g), outPly) + times["ply"] = since(t) + progress.stage(Stage.PLY, null, times["ply"], true) + times["generate"] = since(tGen) + times["all"] = since(tAll) + Log.i(TAG, "run done: " + times.entries.joinToString { "${it.key} ${"%.2f".format(it.value)}" }) + return RunResult( + outPly, g.count, times, pre, + pool.compiles.subList(firstCompile, pool.compiles.size).toList(), + pool.events.subList(firstEvent, pool.events.size).toList(), + noise.source, handoffName(), + ) + } + + fun handoffName(): String = when { + direct -> "direct" + options.directHandoff -> "host (direct failed: $directFailure)" + else -> "host" + } + + /** feature2 = 5 zero tokens + the 4,096 VAE tokens: [1, 4101, 128]. */ + private fun cond2(tokens: FloatArray): FloatArray { + val f2 = FloatArray(Glue.COND_TOKENS * 128) + System.arraycopy(tokens, 0, f2, 5 * 128, tokens.size) + return f2 + } + + // ------------------------------------------------------------------------------------------------ stages + private fun graph(name: String): Graph { + if (cancel) throw CancelledException() + if (!pool.isResident(name)) progress.status("Loading $name") + val g = pool.get(name) + progress.status("") + return g + } + + /** The matting graph's uint8 alpha for the resized RGB picture; the graph is closed right after. */ + private fun alpha(rgb: Image8): ByteArray { + val m = options.matting + val input = if (m == "bir") Matting.birefnetInput(rgb) else Matting.disInput(rgb) + try { + val g = graph(m) + g.write(0, input) + g.run() + val map = g.read(0) + pool.release(m) + return if (m == "bir") Matting.birefnetAlpha(map, rgb.width, rgb.height) + else Matting.disAlpha(map, rgb.width, rgb.height, options.disFloor) + } finally { + pool.release(m) + } + } + + private class Link(val from: Graph, val fromOut: Int, val to: Graph, val toIn: Int, val direct: Boolean) + + /** + * The next graph [toName] of a chain, fed [from]'s output [fromOut] at input [toIn]: copied through the host now + * (the read waits for [from]'s GPU work), or handed over as the buffer itself at [run] with direct handoff. + */ + private fun link(from: Graph, fromOut: Int, toName: String, toIn: Int = 0): Link { + if (direct) { + pool.pin(from.name) + try { + return Link(from, fromOut, graph(toName), toIn, true) + } catch (e: Throwable) { + pool.unpin(from.name) + throw e + } + } + val data = from.read(fromOut) + val to = graph(toName) + to.write(toIn, data) + return Link(from, fromOut, to, toIn, false) + } + + private fun run(l: Link) { + if (!l.direct) { + l.to.run() + return + } + try { + l.to.run(overrides = mapOf(l.toIn to l.from.outputBuffer(l.fromOut))) + l.to.untag(l.toIn) + } catch (e: LiteRtException) { + Log.w(TAG, "direct handoff ${l.from.name} -> ${l.to.name} failed, host copies from now on", e) + direct = false + directFailure = "${l.from.name} -> ${l.to.name}: ${e.message}" + l.to.write(l.toIn, l.from.read(l.fromOut)) + l.to.run() + } finally { + pool.unpin(l.from.name) + } + } + + private fun dinov3(pixels: FloatArray): FloatArray { + var g = graph(dinoNames[0]) + g.write(0, pixels) + g.write(1, ropeCos, "rope_cos") + g.write(2, ropeSin, "rope_sin") + g.run() + for (name in dinoNames.drop(1)) { + val l = link(g, 0, name) + l.to.write(1, ropeCos, "rope_cos") + l.to.write(2, ropeSin, "rope_sin") + run(l) + g = l.to + } + return g.read(0) + } + + private fun vaeMoments(pixels: FloatArray): FloatArray { + var g = graph(vaeNames[0]) + g.write(0, pixels) + g.run() + for (name in vaeNames.drop(1)) { + val l = link(g, 0, name) + run(l) + g = l.to + } + return g.read(0) + } + + /** The conditioning of one LoRA phase (den_ctx), kept for the steps of that phase. */ + private class Ctx(val h: FloatArray, val tag: Long) + + /** sample: 5 Euler steps from the shared noise -> the final latent [1, 8192, 16]. */ + private fun sample(noise: Noise, f1: FloatArray, f2: FloatArray, times: MutableMap): FloatArray { + val xLat = noise.latent.copyOf() + val xCam = noise.camera.copyOf() + val ctx = HashMap() + val condTag = nextTag() + var den = 0.0 + val steps = timeGrid.size - 1 + for (i in 0 until steps) { + progress.stage(Stage.DENOISER, "step ${i + 1} of $steps", null, false) + val t0 = now() + val (v, vc) = velocity(xLat, xCam, timeGrid[i], f1, f2, condTag, ctx) + den += since(t0) + val dt = (timeGrid[i] - timeGrid[i + 1]).toFloat() + for (k in xLat.indices) xLat[k] = xLat[k] - v[k] * dt + for (k in xCam.indices) xCam[k] = xCam[k] - vc[k] * dt + } + times["denoiser"] = den + return xLat + } + + /** velocity: one denoiser call at time t -> (latent velocity [1, 8192, 16], camera velocity [1, 1, 5]). */ + private fun velocity( + xLat: FloatArray, + xCam: FloatArray, + t: Double, + f1: FloatArray, + f2: FloatArray, + condTag: Long, + ctx: HashMap, + ): Pair { + val slow = t >= Glue.PHASE_THRESHOLD + val gate = if (slow) gateSlow else gateFast + val gateTag = if (slow) "gate_slow" else "gate_fast" + val c = ctx.getOrPut(slow) { + val g = graph("den_ctx") + g.write(0, f1, "f1_$condTag") + g.write(1, f2, "f2_$condTag") + g.write(2, gate, gateTag) + g.run() + Ctx(g.read(0), nextTag()) + } + val din = graph("den_in") + din.write(0, xLat) + din.write(1, xCam) + din.write(2, Glue.sinusoid(1000 * t)) + din.write(3, posEmb, "pos_emb") + din.write(4, gate, gateTag) + din.run() + val hCam = din.read(1) + val tMod = din.read(2) + val tEmb = din.read(3) + val tModTag = nextTag() + var prev = din + for ((k, name) in jointNames.withIndex()) { + val l = link(prev, 0, name) + val j = l.to + when (k) { + 0 -> { + j.write(1, c.h, c.tag) + j.write(2, hCam) + j.write(3, tMod, tModTag) + j.write(4, gate, gateTag) + } + jointNames.size - 1 -> { + j.write(1, tMod, tModTag) + j.write(2, tEmb) + j.write(3, gate, gateTag) + } + else -> { + j.write(1, tMod, tModTag) + j.write(2, gate, gateTag) + } + } + run(l) + prev = j + } + return prev.read(0) to prev.read(1) + } + + /** octree_logits: the decoder's child logits for [n] cells at [level], rows padded to the bucket -> [n, 8]. */ + private fun octreeLogits(cells: FloatArray, n: Int, level: Int, latent: FloatArray, latentTag: Any?): FloatArray { + val b = Glue.OCT_BUCKETS.first { n <= it } + val key = "L$b" + val c = FloatArray(b * 3) + System.arraycopy(cells, 0, c, 0, n * 3) + val pe = Glue.pointPe(c, b) + val lf = Glue.sinusoid((1 shl level).toDouble(), maxPeriod = 1024.0, twoPi = true) + val g = graph("oct") + g.write(0, c, key = key) + g.write(1, pe, key = key) + g.write(2, lf, key = key) + g.write(3, latent, latentTag, key) + g.run(key) + return g.read(0, key).copyOf(n * 8) + } + + /** The Gaussian decoder's 8 graphs: anchors [8192, 3] + latent -> features [8192, 480]. */ + private fun gaussianFeatures(anchors: FloatArray, latent: FloatArray, latentTag: Any?): FloatArray { + val pe = Glue.pointPe(anchors, Glue.NUM_POINTS) + var g = graph(gsNames[0]) + g.write(0, anchors) + g.write(1, pe) + g.write(2, latent, latentTag) + g.run() + for (name in gsNames.drop(1)) { + val l = link(g, 0, name) + l.to.write(1, latent, latentTag) + run(l) + g = l.to + } + return g.read(0) + } + + // ------------------------------------------------------------------------------------------------ parity + /** + * CONTRACT.md line (a) on the phone: each part gets the fixture's reference input for it (no chaining) and its + * output is written as `..f32` (little-endian float32) in [out]; the matting alpha as `.alpha.u8`. + */ + fun parity(fx: Fixture, out: File, id: String, progress: Progress): JSONObject { + cancel = false + this.progress = progress + pool.phase = "parity" + pool.setSchedule(emptyList()) + val arrays = JSONObject() + val seconds = JSONObject() + fun dump(name: String, data: FloatArray, shape: IntArray) { + Fixture.writeF32(File(out, "$id.$name.f32"), data) + arrays.put(name, JSONObject().put("file", "$id.$name.f32").put("shape", JSONArray(shape.toList()))) + } + val result = JSONObject().put("id", id).put("fixture", fx.dir.name).put("matting", options.matting) + + progress.stage(Stage.MATTING, null, null, false) + var t = now() + val photo = Photo.fromFile(fx.image) + val pre = Preprocess.run(photo.image, fx.mode) { alpha(it) } + File(out, "$id.alpha.u8").writeBytes(pre.alphaRaw) + result.put("alpha", JSONObject().put("file", "$id.alpha.u8").put("width", pre.alphaWidth) + .put("height", pre.alphaHeight).put("box", JSONArray(pre.box.toList()))) + seconds.put("matting", since(t)) + progress.stage(Stage.MATTING, null, since(t), true) + + val pixels = fx.floats("pixels") + progress.stage(Stage.DINO, null, null, false) + t = now() + dump("feature1", dinov3(pixels), intArrayOf(1, Glue.COND_TOKENS, 1280)) + seconds.put("dinov3", since(t)) + pool.releaseIf { it.startsWith("dino_") } + progress.stage(Stage.DINO, null, since(t), true) + + progress.stage(Stage.VAE, null, null, false) + t = now() + val moments = vaeMoments(pixels) + dump("vae_moments", moments, intArrayOf(1, 64, 128, 128)) + dump("vae_tokens", Glue.vaeTokens(moments, fx.floats("vae_noise"), bnMean, bnStd), intArrayOf(1, 4096, 128)) + seconds.put("vae", since(t)) + pool.releaseIf { it.startsWith("vae_") } + progress.stage(Stage.VAE, null, since(t), true) + + t = now() + val f1 = fx.floats("feature1") + val f2 = fx.floats("feature2") + // t comes from the engine's float64 grid, as the reference passes float(ts[i]) (1000 t rounds differently from + // the fixture's float32 copy); the fixture's grid only confirms it is the same schedule + val tg = fx.floats("t_grid") + result.put("t_grid_matches_fixture", tg.size == timeGrid.size && tg.indices.all { tg[it] == timeGrid[it].toFloat() }) + val ctx = HashMap() + val condTag = nextTag() + for (i in 0 until timeGrid.size - 1) { + progress.stage(Stage.DENOISER, "step ${i + 1} of ${timeGrid.size - 1}", null, false) + val (v, vc) = velocity(fx.floats("x_latent_$i"), fx.floats("x_camera_$i"), timeGrid[i], f1, f2, condTag, ctx) + dump("v_latent_$i", v, intArrayOf(1, Glue.NUM_POINTS, Glue.IN_CH)) + dump("v_camera_$i", vc, intArrayOf(1, 1, Glue.CAM_CH)) + } + seconds.put("denoiser", since(t)) + pool.releaseIf { it.startsWith("den_") } + progress.stage(Stage.DENOISER, null, since(t), true) + + t = now() + val latent = fx.floats("latent") + val latentTag = nextTag() + for (level in 1..Glue.MAX_LEVEL) { + progress.stage(Stage.OCTREE, "level $level of ${Glue.MAX_LEVEL}", null, false) + val cells = fx.floats("oct_cells_$level") + val n = cells.size / 3 + dump("oct_logits_$level", octreeLogits(cells, n, level, latent, latentTag), intArrayOf(1, n, 8)) + } + seconds.put("octree", since(t)) + pool.release("oct") + progress.stage(Stage.OCTREE, null, since(t), true) + + progress.stage(Stage.GAUSSIANS, null, null, false) + t = now() + val anchors = fx.floats("anchors") + val feats = gaussianFeatures(anchors, latent, latentTag) + seconds.put("gaussian_decoder", since(t)) + pool.releaseIf { it.startsWith("gs_") } + dump("gs_features", feats, intArrayOf(1, Glue.NUM_POINTS, 480)) + val g = Glue.toGaussians(anchors, feats, rep, perturbation) + dump("g_xyz", g.xyz, intArrayOf(g.count, 3)) + dump("g_features_dc", g.featuresDc, intArrayOf(g.count, 1, 3)) + dump("g_scaling", g.scaling, intArrayOf(g.count, 3)) + dump("g_rotation", g.rotation, intArrayOf(g.count, 4)) + dump("g_opacity", g.opacity, intArrayOf(g.count, 1)) + progress.stage(Stage.GAUSSIANS, null, since(t), true) + pool.phase = "run" + return result.put("arrays", arrays).put("seconds", seconds).put("handoff", handoffName()) + } + + companion object { + const val TAG = "IdolEngine" + private val tags = AtomicLong(1) + + fun nextTag(): Long = tags.getAndIncrement() + + private fun now() = SystemClock.elapsedRealtimeNanos() + + private fun since(t0: Long) = (SystemClock.elapsedRealtimeNanos() - t0) / 1e9 + } +} diff --git a/android/app/src/main/java/com/idol/MainActivity.kt b/android/app/src/main/java/com/idol/MainActivity.kt new file mode 100644 index 0000000000000000000000000000000000000000..baffcde43c455e5c38307cfa820722265bc8059a --- /dev/null +++ b/android/app/src/main/java/com/idol/MainActivity.kt @@ -0,0 +1,581 @@ +package com.idol + +import android.Manifest +import android.content.ActivityNotFoundException +import android.content.ComponentCallbacks2 +import android.content.Intent +import android.content.pm.ApplicationInfo +import android.content.pm.PackageManager +import android.net.Uri +import android.os.Build +import android.os.Bundle +import android.os.Handler +import android.os.Looper +import android.os.SystemClock +import android.provider.OpenableColumns +import android.util.Log +import android.util.TypedValue +import android.view.View +import android.view.WindowManager +import android.webkit.WebView +import android.widget.Button +import android.widget.CheckBox +import android.widget.FrameLayout +import android.widget.ImageView +import android.widget.LinearLayout +import android.widget.RadioGroup +import android.widget.TextView +import androidx.activity.ComponentActivity +import androidx.activity.OnBackPressedCallback +import androidx.activity.result.PickVisualMediaRequest +import androidx.activity.result.contract.ActivityResultContracts +import androidx.core.content.FileProvider +import androidx.core.view.ViewCompat +import androidx.core.view.WindowInsetsCompat +import org.json.JSONObject +import java.io.File +import java.text.SimpleDateFormat +import java.util.Date +import java.util.Locale +import java.util.concurrent.Callable +import java.util.concurrent.CountDownLatch +import java.util.concurrent.ExecutorService +import java.util.concurrent.Executors + +/** How a run ended, for the screen and the autorun. */ +sealed class Outcome { + class Done(val result: RunResult) : Outcome() + class NoObject(val reason: String) : Outcome() + class Missing(val problems: List) : Outcome() + class Failed(val kind: String, val message: String) : Outcome() +} + +/** The model state of the process: one model thread and one engine, kept across activity instances. */ +object Idol { + /** Every LiteRT call (compile, run, close) happens on this one thread. */ + val model: ExecutorService = Executors.newSingleThreadExecutor { r -> Thread(r, "IdolModel") } + @Volatile var engine: IdolEngine? = null + @Volatile var problems: List? = null + @Volatile var started = false +} + +/** + * Home (pick or take a photo) -> Run (the stages with their seconds) -> Result (the splat viewer). Every run, from a + * button or from the debug autorun ([AutoRun]), goes through [startRun]; the work is on the model thread, never on + * the UI thread. + */ +class MainActivity : ComponentActivity() { + private lateinit var home: View + private lateinit var runScreen: View + private lateinit var resultScreen: View + private lateinit var modeGroup: RadioGroup + private lateinit var disCheck: CheckBox + private lateinit var homeMessage: TextView + private lateinit var homeStatus: TextView + private lateinit var picture: ImageView + private lateinit var cutout: ImageView + private lateinit var stageList: LinearLayout + private lateinit var totalText: TextView + private lateinit var runStatus: TextView + private lateinit var viewerBox: FrameLayout + private var web: WebView? = null + + private val ui = Handler(Looper.getMainLooper()) + private val decoder = Executors.newSingleThreadExecutor { r -> Thread(r, "IdolDecode") } + private val stageRows = ArrayList() + + private class StageState { + var detail: String? = null + var seconds: Double? = null + var startedAt = 0L + var running = false + var done = false + } + + private val stages = Array(Stage.entries.size) { StageState() } + @Volatile private var running = false + @Volatile var autorunActive = false + private set + private var lastPly: File? = null + private var deviceLine = "" + + /** Engine progress, posted to the screen. */ + val progress = object : Progress { + override fun stage(stage: Stage, detail: String?, seconds: Double?, done: Boolean) { + ui.post { onStage(stage, detail, seconds, done) } + } + + override fun status(text: String) { + ui.post { + runStatus.text = text + homeStatus.text = if (text.isEmpty()) deviceLine else "$text\n$deviceLine" + } + } + + override fun cutout(composite: Image8) { + val bmp = Photo.displayBitmap(composite, 512) + ui.post { + cutout.setImageBitmap(bmp) + cutout.visibility = View.VISIBLE + } + } + } + + private val pickPhoto = registerForActivityResult(ActivityResultContracts.PickVisualMedia()) { uri -> + if (uri != null) startFromPicture { Photo.fromUri(contentResolver, uri, displayName(uri)) } + } + + private val takePhoto = registerForActivityResult(ActivityResultContracts.TakePicture()) { ok -> + // the camera app can take enough memory for Android to stop this process while it is open; the activity is + // then created again for the result, so the photo's path is a fixed one, not a field + val file = cameraFile() + if (ok && file.isFile) startFromPicture { Photo.fromFile(file) } + } + + private val cameraPermission = registerForActivityResult(ActivityResultContracts.RequestPermission()) { granted -> + if (granted) launchCamera() else showHome(getString(R.string.camera_denied)) + } + + override fun onCreate(savedInstanceState: Bundle?) { + super.onCreate(savedInstanceState) + if (isDebuggable()) { + // the autorun drives a phone nobody unlocked: stay in front (cpuset top-app) with the screen on + window.addFlags(WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON) + if (Build.VERSION.SDK_INT >= 27) { + setShowWhenLocked(true) + setTurnScreenOn(true) + } + } + setContentView(R.layout.activity_main) + // targetSdk 35 draws edge to edge: keep the screens clear of the status and navigation bars + val root = findViewById(R.id.root) + ViewCompat.setOnApplyWindowInsetsListener(root) { v, insets -> + val bars = insets.getInsets(WindowInsetsCompat.Type.systemBars()) + v.setPadding(bars.left, bars.top, bars.right, bars.bottom) + insets + } + home = findViewById(R.id.home) + runScreen = findViewById(R.id.run) + resultScreen = findViewById(R.id.result) + modeGroup = findViewById(R.id.mode) + disCheck = findViewById(R.id.dis) + homeMessage = findViewById(R.id.home_message) + homeStatus = findViewById(R.id.home_status) + picture = findViewById(R.id.picture) + cutout = findViewById(R.id.cutout) + stageList = findViewById(R.id.stages) + totalText = findViewById(R.id.total) + runStatus = findViewById(R.id.run_status) + viewerBox = findViewById(R.id.viewer) + + for (s in Stage.entries) { + val tv = TextView(this) + tv.setTextSize(TypedValue.COMPLEX_UNIT_SP, 16f) + tv.setPadding(0, dp(6), 0, dp(6)) + stageList.addView(tv) + stageRows.add(tv) + } + + findViewById