Initial public release

Co-authored-by: claude opus 4.8 (1m context) <claude opus 4.8 (1m context)@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,86 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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
+.gradio/cached_examples/30/Depth[[:space:]](disparity,[[:space:]]magma)/5e96afc67581e463f100/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/30/Depth[[:space:]](disparity,[[:space:]]magma)/9220cd48d295ca353a6b/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/30/Depth[[:space:]](disparity,[[:space:]]magma)/c8838412e5dee9697701/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/30/Depth[[:space:]](disparity,[[:space:]]magma)/f7095c71dfb49f7e3a2a/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/30/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/a2b2ed4415d6c0310a36/cloud_05cfcce5d5e6.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/30/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/b72cecf121cf47496c82/cloud_394dafbecc42.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/30/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/b8c85f926a8d6d09ee32/cloud_dbf7146aa8d6.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/30/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/bb6dcdbd8cf58c73cd2c/cloud_a05855c7716b.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/30/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/c3390d29364b0fab5237/cloud_14b821bf5517.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/30/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/f90705e1df43a796cd63/cloud_41032084d765.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/30/RGB[[:space:]]image/02ab28426d5807f6943a/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/30/RGB[[:space:]]image/179b01cce83461a9a19a/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/30/RGB[[:space:]]image/1965504ccdd23a0da632/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/30/RGB[[:space:]]image/3f5824c31471ef256471/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/30/RGB[[:space:]]image/7ebf7b4b6db056927956/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/30/RGB[[:space:]]image/dba880c3bec4afd1a9ff/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/Depth[[:space:]](disparity,[[:space:]]magma)/b8e54b0d5e39f055f164/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/Depth[[:space:]](disparity,[[:space:]]magma)/e3b3eed2144fe2f212bb/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/7b2796d4a9ff918e3ff0/cloud_16bed495fda3.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/d21fe0e76f4b13cc03c4/cloud_b91d63f1da9e.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/ff979f9418cc0388e99b/cloud_006af054b659.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/RGB[[:space:]]image/1a6ed5112df1d233314c/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/RGB[[:space:]]image/c14e4a0757ca88e19135/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/RGB[[:space:]]image/e858eac28e2a39b56e43/image.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/alley.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/castle.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/yosemite.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/Depth[[:space:]](disparity,[[:space:]]magma)/00514d1d228e1970f9a5/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/Depth[[:space:]](disparity,[[:space:]]magma)/282b8de47381433b20ba/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/Depth[[:space:]](disparity,[[:space:]]magma)/77eb75a4ae32a2779670/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/Depth[[:space:]](disparity,[[:space:]]magma)/b3a0c160552d198bcca2/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/6351a7c3b27ea01e3670/cloud_45fd4adef349.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/a42239d35b0256820e21/cloud_82bf2e710be0.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/a600814ba3ef00eacfeb/cloud_9a8d39949ba7.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/a6b517c351861e68aa8f/cloud_6f14c0ff2da6.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/a93a97926162980d549a/cloud_53c7df1dbc77.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/e208d09a8d563ea27323/cloud_a7b919815dba.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/RGB[[:space:]]image/4763ceeb1dd670d3f927/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/RGB[[:space:]]image/530eb79823e527aed6bf/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/RGB[[:space:]]image/65c25b8debef84e0027b/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/RGB[[:space:]]image/6d51368904340f3a7f23/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/RGB[[:space:]]image/91174cfa8f5a5909a013/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/31/RGB[[:space:]]image/d68b76efe33346efd7e3/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/32/Depth[[:space:]](disparity,[[:space:]]magma)/38d56c7397d95d524055/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/32/Depth[[:space:]](disparity,[[:space:]]magma)/9a36d2e382db4e333742/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/32/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/674110f2aecaaf5effa8/cloud_f9d62ed7d997.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/32/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/85ee5283a423852aa42d/cloud_8a6b0d1109ae.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/32/Interactive[[:space:]]3D[[:space:]]point[[:space:]]cloud/e7f9024213e4656b8beb/cloud_a6225bea9d7e.glb filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/32/RGB[[:space:]]image/3c2c4fa7881bb773646d/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/32/RGB[[:space:]]image/87aaba9c524a915404cd/image.png filter=lfs diff=lfs merge=lfs -text
+.gradio/cached_examples/32/RGB[[:space:]]image/e64a0aff7688d0635057/image.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.pyc
+.venv/
+flux_rgbd_outputs/
+*.npy
+artifacts/

.gradio/cached_examples/31/Interactive 3D point cloud/6351a7c3b27ea01e3670/cloud_45fd4adef349.glb

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

.gradio/cached_examples/31/Interactive 3D point cloud/a42239d35b0256820e21/cloud_82bf2e710be0.glb

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

.gradio/cached_examples/31/Interactive 3D point cloud/a600814ba3ef00eacfeb/cloud_9a8d39949ba7.glb

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

.gradio/cached_examples/31/Interactive 3D point cloud/a6b517c351861e68aa8f/cloud_6f14c0ff2da6.glb

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

.gradio/cached_examples/31/Interactive 3D point cloud/a93a97926162980d549a/cloud_53c7df1dbc77.glb

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

.gradio/cached_examples/31/Interactive 3D point cloud/e208d09a8d563ea27323/cloud_a7b919815dba.glb

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

.gradio/cached_examples/31/indices.csv

@@ -0,0 +1,6 @@
+0
+1
+2
+3
+4
+5

.gradio/cached_examples/31/log.csv

@@ -0,0 +1,7 @@
+RGB image,"Depth (disparity, magma)",Interactive 3D point cloud,Status,timestamp
+"{""path"": "".gradio/cached_examples/31/RGB image/91174cfa8f5a5909a013/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/dfd5ecf1e795fbfd7f794edd94c8f06e7dc0f0008a15814286a7a11b0e3bda0a/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/31/Depth (disparity, magma)/72f252326e89ba602db3/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/7111df1c084d7d80e87ed0ba942ef806ebb82b0ea3358242a9471e87ed025aca/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/31/Interactive 3D point cloud/e208d09a8d563ea27323/cloud_a7b919815dba.glb"", ""url"": ""/gradio_api/file=/tmp/gradio/b7842ce128885af0de98f62956ea71483eb3ca75a5af89df81202f0c19de2502/cloud_a7b919815dba.glb"", ""size"": null, ""orig_name"": ""cloud_a7b919815dba.glb"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","joint · 17.0 s · depth median=1.06  p5=0.55  p95=1.39 · 249,485 points",2026-06-12 07:14:17.591889
+"{""path"": "".gradio/cached_examples/31/RGB image/d68b76efe33346efd7e3/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/db51824edbde88fd4de0749ffb077c98b27d9ee96017d42db904e2c763579606/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/31/Depth (disparity, magma)/21da509194782c87b704/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/e3e01501ae08e1438d11238cf43d427406c012cf0b287f1820975b5cc461f224/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/31/Interactive 3D point cloud/a93a97926162980d549a/cloud_53c7df1dbc77.glb"", ""url"": ""/gradio_api/file=/tmp/gradio/96342e925c3ba767021690b60f3715015db850198dccd2fd3e5245468dbcafee/cloud_53c7df1dbc77.glb"", ""size"": null, ""orig_name"": ""cloud_53c7df1dbc77.glb"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","joint · 14.8 s · depth median=1.04  p5=0.57  p95=1.37 · 253,113 points",2026-06-12 07:14:32.666816
+"{""path"": "".gradio/cached_examples/31/RGB image/6d51368904340f3a7f23/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/9919ff6eef6ba9aa87da068492e8a80dd0e5f62f7f27dcc20cb3de5e975ef2b0/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/31/Depth (disparity, magma)/00514d1d228e1970f9a5/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/0ffd5b599e783b31637214af5bfb99c6fea27cce8df77021893d3f63e9e1f5aa/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/31/Interactive 3D point cloud/a600814ba3ef00eacfeb/cloud_9a8d39949ba7.glb"", ""url"": ""/gradio_api/file=/tmp/gradio/3e1f8dedd7baf5f8977d2f5e7488413ca6499045d5939ee6371c86fb2b039e54/cloud_9a8d39949ba7.glb"", ""size"": null, ""orig_name"": ""cloud_9a8d39949ba7.glb"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","joint · 14.9 s · depth median=0.84  p5=0.37  p95=2.21 · 97,259 points",2026-06-12 07:14:47.729293
+"{""path"": "".gradio/cached_examples/31/RGB image/530eb79823e527aed6bf/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/0bbd8ccaaeccd0e353685be2f1146bbbdd4e2a42d5350fe13bd182a27f505a61/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/31/Depth (disparity, magma)/b3a0c160552d198bcca2/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/91abce7d476b85c56d457cb3526e279dbf126ff0d36d266cbde7b0ba97d0714f/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/31/Interactive 3D point cloud/a42239d35b0256820e21/cloud_82bf2e710be0.glb"", ""url"": ""/gradio_api/file=/tmp/gradio/fff39c18ed5428d50126dee01f175fe9339a93dedb01c226549e19234c3a18d3/cloud_82bf2e710be0.glb"", ""size"": null, ""orig_name"": ""cloud_82bf2e710be0.glb"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","joint · 14.9 s · depth median=0.98  p5=0.65  p95=1.38 · 239,893 points",2026-06-12 07:15:02.835368
+"{""path"": "".gradio/cached_examples/31/RGB image/65c25b8debef84e0027b/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/eed4a59e42c6c0c9e9cd120b9f5900522c0f4006c54aec5e3466300cce5f5c3d/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/31/Depth (disparity, magma)/77eb75a4ae32a2779670/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/c9434cd1a2b855cdf7615290ea2d29c8c5418e0793a3cd859e55150891bc495a/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/31/Interactive 3D point cloud/a6b517c351861e68aa8f/cloud_6f14c0ff2da6.glb"", ""url"": ""/gradio_api/file=/tmp/gradio/3fe22788a5f8d0a6d274d5b367522db4306e5dca95dbfe1bd1cfb1c2f1c4c6f0/cloud_6f14c0ff2da6.glb"", ""size"": null, ""orig_name"": ""cloud_6f14c0ff2da6.glb"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","joint · 14.8 s · depth median=0.90  p5=0.63  p95=1.65 · 253,509 points",2026-06-12 07:15:17.944777
+"{""path"": "".gradio/cached_examples/31/RGB image/4763ceeb1dd670d3f927/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/71b22a185a03aa169cdca016381da0b4179ed403b377b7fe4cb005c97797b8ee/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/31/Depth (disparity, magma)/282b8de47381433b20ba/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/9ef7335cc9a72e4364a920bf4431a69841eff7326a683f601376302cc1ddaa24/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/31/Interactive 3D point cloud/6351a7c3b27ea01e3670/cloud_45fd4adef349.glb"", ""url"": ""/gradio_api/file=/tmp/gradio/92aa7dea6a27cd503ba5c2235856b6e28f283700cbfc51764f794167491464e4/cloud_45fd4adef349.glb"", ""size"": null, ""orig_name"": ""cloud_45fd4adef349.glb"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","joint · 14.8 s · depth median=0.82  p5=0.07  p95=2.22 · 191,720 points",2026-06-12 07:15:32.973544

.gradio/cached_examples/32/Interactive 3D point cloud/674110f2aecaaf5effa8/cloud_f9d62ed7d997.glb

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

.gradio/cached_examples/32/Interactive 3D point cloud/85ee5283a423852aa42d/cloud_8a6b0d1109ae.glb

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

.gradio/cached_examples/32/Interactive 3D point cloud/e7f9024213e4656b8beb/cloud_a6225bea9d7e.glb

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

.gradio/cached_examples/32/indices.csv

@@ -0,0 +1,3 @@
+0
+1
+2

.gradio/cached_examples/32/log.csv

@@ -0,0 +1,4 @@
+RGB image,"Depth (disparity, magma)",Interactive 3D point cloud,Status,timestamp
+"{""path"": "".gradio/cached_examples/32/RGB image/3c2c4fa7881bb773646d/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/1ed5fba88325554dd8fde55f43d90e4b52d96bd770ae9abb55558a721d2ad4dc/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/32/Depth (disparity, magma)/38d56c7397d95d524055/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/675d608d8081b140ee360ac99a11615574a95a1f575dd55ceb8941f405daac0f/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/32/Interactive 3D point cloud/674110f2aecaaf5effa8/cloud_f9d62ed7d997.glb"", ""url"": ""/gradio_api/file=/tmp/gradio/6396a2958022a63007f08ff0882cd34405d96b6cee56b1a69444e1ef08649811/cloud_f9d62ed7d997.glb"", ""size"": null, ""orig_name"": ""cloud_f9d62ed7d997.glb"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","i2d · 5.4 s · depth median=0.65  p5=0.37  p95=3.31 · 233,069 points",2026-06-12 07:15:38.629118
+"{""path"": "".gradio/cached_examples/32/RGB image/e64a0aff7688d0635057/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/3a7962dcf7f6b1d7db2e1a0e7a625daecd56aebbac87a934bd805b05c026ee6c/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/32/Depth (disparity, magma)/9a36d2e382db4e333742/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/647b0d2a6553dedf70711720e7a0dd7c2d0c1778bbbc29c066ebbdb74137ba43/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/32/Interactive 3D point cloud/85ee5283a423852aa42d/cloud_8a6b0d1109ae.glb"", ""url"": ""/gradio_api/file=/tmp/gradio/9b0b10012948a25b48290aefa846f9ddb57792cae0bccf98ead298a0a18baf6e/cloud_8a6b0d1109ae.glb"", ""size"": null, ""orig_name"": ""cloud_8a6b0d1109ae.glb"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","i2d · 5.0 s · depth median=1.18  p5=0.36  p95=2.33 · 110,475 points",2026-06-12 07:15:43.778321
+"{""path"": "".gradio/cached_examples/32/RGB image/87aaba9c524a915404cd/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/782693cd71c502c69e559adc17a82d06dc75ece451c4745b286f74f8be98284d/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/32/Depth (disparity, magma)/0891e8543770d0db2151/image.png"", ""url"": ""/gradio_api/file=/tmp/gradio/112070edef52867273e2552aa71313e0447940a6eda5b9215a5202cc82c501a2/image.png"", ""size"": null, ""orig_name"": ""image.png"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","{""path"": "".gradio/cached_examples/32/Interactive 3D point cloud/e7f9024213e4656b8beb/cloud_a6225bea9d7e.glb"", ""url"": ""/gradio_api/file=/tmp/gradio/bd563503321e9ee362fcadf84021c707d2ddfb1ae60357b595f562a830465900/cloud_a6225bea9d7e.glb"", ""size"": null, ""orig_name"": ""cloud_a6225bea9d7e.glb"", ""mime_type"": null, ""is_stream"": false, ""meta"": {""_type"": ""gradio.FileData""}}","i2d · 5.0 s · depth median=0.72  p5=0.08  p95=2.41 · 233,820 points",2026-06-12 07:15:49.021708

.ruff_cache/.gitignore

@@ -0,0 +1,2 @@
+# Automatically created by ruff.
+*

.ruff_cache/CACHEDIR.TAG

@@ -0,0 +1 @@
+Signature: 8a477f597d28d172789f06886806bc55

LICENSE

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

LICENSE-WEIGHTS

@@ -0,0 +1,14 @@
+Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0)
+
+Copyright (c) 2026 World Labs.
+
+The MODEL WEIGHTS distributed for this project (e.g. via the Hugging Face
+repository bartduis/modality_forcing) are licensed under the Creative Commons
+Attribution-NonCommercial 4.0 International License. You are free to share and
+adapt the weights for non-commercial purposes, provided you give appropriate
+credit. THE WEIGHTS ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND.
+
+Full license text: https://creativecommons.org/licenses/by-nc/4.0/legalcode
+
+The CODE in this repository is licensed separately under Apache-2.0 (see
+LICENSE).

README.md

@@ -0,0 +1,99 @@
+---
+title: Modality Forcing
+emoji: 🏢
+colorFrom: blue
+colorTo: gray
+sdk: gradio
+sdk_version: 6.16.0
+python_version: '3.12'
+app_file: app.py
+hardware: zero-h200
+pinned: false
+license: apache-2.0
+short_description: Text → RGB + depth + 3D point cloud.
+---
+
+<div align="center">
+
+<h1>Modality Forcing for Scalable<br>Spatial Generation</h1>
+
+[![Project Page](https://img.shields.io/badge/Project_Page-Modality_Forcing-blue?logo=googlechrome&logoColor=white)](https://modality-forcing.github.io/)
+[![Code](https://img.shields.io/badge/Code-GitHub-181717?logo=github)](https://github.com/Duisterhof/modality-forcing)
+[![Model](https://img.shields.io/badge/Model-modality__forcing-yellow?logo=huggingface)](https://huggingface.co/bartduis/modality_forcing)
+[![arXiv](https://img.shields.io/badge/arXiv-2606.13676-b31b1b?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2606.13676)
+
+**[Bardienus Pieter Duisterhof](https://bart-ai.com)<sup>1,2</sup> · [Deva Ramanan](https://www.cs.cmu.edu/~deva/)<sup>1</sup> · [Jeffrey Ichnowski](https://ichnow.ski)<sup>1</sup> · [Justin Johnson](https://web.eecs.umich.edu/~justincj/)<sup>2</sup> · [Keunhong Park](https://keunhong.com)<sup>2</sup>**
+
+**<sup>1</sup> Carnegie Mellon University &nbsp;&nbsp; <sup>2</sup> World Labs**
+
+*Preprint, 2026*
+
+<em>Modality Forcing turns a pretrained text-to-image model into a joint image–depth generator with a simple post-training recipe.</em>
+
+</div>
+
+## Overview
+
+This Space hosts the interactive demo of **Modality Forcing**: joint text → RGB + depth diffusion built on FLUX.2. A single DiT supports every permutation of conditional and joint generation by assigning a separate noise level to each modality:
+
+| Mode | Input | Output |
+|------|-------|--------|
+| **Joint** | text prompt | RGB + depth + 3D point cloud |
+| **Image → depth** | text + image | depth + 3D point cloud (any aspect ratio, letterbox resize) |
+
+## Abstract
+
+Text-to-image (T2I) models contain rich spatial priors.
+Synthesizing photorealistic, cluttered scenes requires an understanding of geometry, including perspective and relative scale.
+Prior works adapt T2I models to leverage this prior for depth prediction, but they require dense depth data and involve complex recipes.
+We propose **Modality Forcing**, a simple, scalable post-training recipe for joint image-depth generation using a single DiT trained on sparse depth data.
+Modality Forcing enables conditional and joint generation of image and depth in any permutation by assigning separate noise levels per modality.
+Per-modality decoders let us train on sparse, real-world depth and achieve strong, generalizable depth prediction.
+We further show that Modality Forcing inherits the scalability of T2I pre-training: by training a set of T2I models from scratch (300M to 3B parameters), we find that larger models trained on more image data produce more accurate depth.
+Our strongest model is competitive with state-of-the-art monocular depth estimators and reduces AbsRel by **57%** relative to existing joint image-depth generative models.
+These results provide strong evidence that image generation is a scalable pre-training objective for spatial perception.
+
+## Space configuration
+
+### Weights
+
+All weights are pulled from public repos — no token required:
+[`bartduis/modality_forcing`](https://huggingface.co/bartduis/modality_forcing)
+(DiT + FLUX.2 autoencoder, CC BY-NC 4.0) and
+[`Qwen/Qwen3-8B`](https://huggingface.co/Qwen/Qwen3-8B) (text encoder).
+Override the model repo via the `WEIGHTS_REPO` Space variable.
+
+### Hardware
+
+Pinned to `zero-h200`. The runner builds the DiT on the `meta` device and
+assigns checkpoint weights straight onto the GPU (`runner.from_pretrained`),
+so the BF16 model (~24 GB) loads in ~8 s on the first invocation — skipping
+~45 s of throwaway random initialization — and stays resident for subsequent
+calls within the same Space instance.
+
+### Avoiding the cold re-download
+
+A Space instance that has gone to sleep starts a fresh container, which
+re-downloads the weights (~24 GB DiT + ~16 GB Qwen3-8B + the FLUX.2 VAE)
+before the load above can run. To skip that, attach **persistent storage**
+(Settings → Storage, or mount an HF Bucket) and point the HF cache at it via
+the `HF_HOME` Space variable so the weights survive restarts.
+
+## License
+
+Code: Apache-2.0 (files derived from the FLUX.2 reference implementation —
+`flux_rgbd/_flux2/`, `flux_rgbd/dit.py` — are Apache-2.0, Copyright Black
+Forest Labs, with World Labs modifications). Model weights: CC BY-NC 4.0.
+
+## Citation
+
+If you find Modality Forcing useful, please consider citing:
+
+```bibtex
+@article{duisterhof2026mofo,
+  title   = {Modality Forcing for Scalable Spatial Generation},
+  author  = {Duisterhof, Bardienus Pieter and Ramanan, Deva and Ichnowski, Jeffrey and Johnson, Justin and Park, Keunhong},
+  journal = {arXiv preprint arXiv:2606.13676},
+  year    = {2026}
+}
+```

app.py

@@ -0,0 +1,642 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 World Labs.
+"""HuggingFace ZeroGPU Space for flux_rgbd.
+
+Wraps the same generate-and-render path used by the local demo with
+``@spaces.GPU`` so HF ZeroGPU can attach a GPU per call. The model
+loads at module-level (CPU), then moves to CUDA inside the first
+decorated call.
+"""
+
+import os
+import sys
+import time
+import uuid
+from pathlib import Path
+
+# The vendored flux_rgbd package lives next to this file.
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+
+from flux_rgbd import FluxRGBDRunner
+from flux_rgbd.pointcloud import depth_edge_mask, statistical_outlier_mask
+
+# BF16 checkpoint, downloaded from the Hub on first use. Override with the
+# WEIGHTS_REPO env var (repo id or local path).
+WEIGHTS_REPO = os.environ.get("WEIGHTS_REPO", "bartduis/modality_forcing")
+
+DEFAULT_PROMPT = (
+    "A warm, inviting kitchen with a rustic-modern feel, where soft morning "
+    "light filters through a small window above the sink."
+)
+
+# Lazy-loaded runner. On ZeroGPU the model is loaded inside the first
+# @spaces.GPU call so the import path costs nothing.
+_runner: FluxRGBDRunner | None = None
+
+
+def _ensure_runner() -> FluxRGBDRunner:
+    global _runner
+    if _runner is None:
+        # Use BF16 Qwen3-8B instead of Qwen3-8B-FP8 to avoid the
+        # finegrained-fp8 / deep-gemm kernel dependency, which currently
+        # hits a metadata.json parse bug on HF's kernels-community.
+        text_encoder = os.environ.get("TEXT_ENCODER_REPO", "Qwen/Qwen3-8B")
+        # Generation resolution must match the checkpoint's training resolution
+        # (512 for the default model, 1024 for the 1024 checkpoint). Set
+        # IMG_RESOLUTION=1024 alongside WEIGHTS_REPO when using the 1024 ckpt.
+        res = int(os.environ.get("IMG_RESOLUTION", "512"))
+        print(f"[boot] loading {WEIGHTS_REPO} @ {res}px (text encoder: {text_encoder})…",
+              flush=True)
+        _runner = FluxRGBDRunner.from_pretrained(
+            WEIGHTS_REPO, device="cuda",
+            dtype=torch.bfloat16, head_dtype=torch.float32,
+            text_encoder=text_encoder, img_hw=(res, res),
+        )
+        print("[boot] runner ready.", flush=True)
+    return _runner
+
+
+# --- helpers (kept inline so the Space repo doesn't depend on demo/app_lib) ---
+
+_SH_C0 = 0.28209479177387814
+
+
+def _letterbox(img: np.ndarray, target: int = 512):
+    """Resize so long side = target, then zero-pad to (target, target)."""
+    import cv2
+    h_in, w_in = img.shape[:2]
+    if h_in >= w_in:
+        h_out, w_out = target, max(1, int(round(w_in * target / h_in)))
+    else:
+        w_out, h_out = target, max(1, int(round(h_in * target / w_in)))
+    resized = cv2.resize(img, (w_out, h_out), interpolation=cv2.INTER_AREA)
+    canvas = np.zeros((target, target, img.shape[2] if img.ndim == 3 else 1),
+                      dtype=img.dtype)
+    if img.ndim == 2:
+        canvas = canvas[..., 0]
+    top = (target - h_out) // 2
+    left = (target - w_out) // 2
+    canvas[top:top + h_out, left:left + w_out] = resized
+    return canvas, (top, left, h_out, w_out)
+
+
+def _depth_to_pointcloud(rgb_u8, depth, *, fov_deg=65.0, max_points=1_200_000,
+                         edge_rtol=0.04, sor=False):
+    h, w = depth.shape
+    fx = w / (2.0 * np.tan(np.deg2rad(fov_deg) / 2.0))
+    cx, cy = w * 0.5, h * 0.5
+    # Keep every valid pixel — no percentile clip. The earlier [1, 99] clip
+    # discarded the nearest 1% of points, carving a hole in the closest
+    # surface (e.g. the front edge of a table) and also dropping the far
+    # background. The i2d depth is clean enough that this clipping isn't
+    # needed and it was cutting off the geometry users care most about.
+    valid = (depth > 0) & np.isfinite(depth)
+    # Depth-edge mask: drop occlusion-boundary "veil" pixels (MoGe-style).
+    if edge_rtol and edge_rtol > 0:
+        valid &= ~depth_edge_mask(depth, rtol=float(edge_rtol))
+    v_idx, u_idx = np.where(valid)
+    z = depth[v_idx, u_idx]
+    x = (u_idx + 0.5 - cx) * z / fx
+    y = (v_idx + 0.5 - cy) * z / fx
+    # glTF / Three.js: +Y up, camera looks down -Z. Flip image-y
+    # (which points down) and depth (which points into the scene).
+    pts = np.stack([x, -y, -z], axis=-1).astype(np.float32)
+    cols = rgb_u8[v_idx, u_idx]
+    if sor:
+        # Statistical outlier rejection: drops isolated floaters, but can
+        # over-trim fine structures — opt-in (the edge mask above is the
+        # default cleanup).
+        inliers = statistical_outlier_mask(pts)
+        pts, cols = pts[inliers], cols[inliers]
+    if pts.shape[0] > max_points:
+        idx = np.random.default_rng(0).choice(pts.shape[0], max_points, replace=False)
+        pts, cols = pts[idx], cols[idx]
+    if pts.shape[0]:
+        pts -= np.median(pts, axis=0, keepdims=True)
+    return pts, cols
+
+
+def _save_glb(path, points, colors):
+    """Colored point cloud → binary glTF, the format gr.Model3D handles cleanly."""
+    import trimesh
+    cloud = trimesh.PointCloud(vertices=points, colors=colors)
+    scene = trimesh.Scene()
+    scene.add_geometry(cloud)
+    scene.export(str(path))
+
+
+def _depth_to_magma(depth: np.ndarray) -> np.ndarray:
+    """Depth → magma-colormapped disparity image (uint8 RGB).
+
+    Visualizes 1/depth (so near = bright) robustly normalized to the 5–95th
+    percentile, matching the depth panel in the reference notebook.
+    """
+    from matplotlib import cm
+    valid = (depth > 0) & np.isfinite(depth)
+    disparity = np.zeros_like(depth, dtype=np.float32)
+    if valid.any():
+        disparity[valid] = 1.0 / np.maximum(depth[valid], 1e-8)
+        lo, hi = np.percentile(disparity[valid], [5, 95])
+        disparity = np.clip((disparity - lo) / max(hi - lo, 1e-8), 0, 1)
+        disparity[~valid] = 0.0
+    return (cm.magma(disparity)[..., :3] * 255).astype(np.uint8)
+
+
+# /tmp is the writable mount on HF Spaces. We write the PLY here from the
+# parent process (i.e. NOT inside the @spaces.GPU subprocess) so Gradio's
+# file route can read it. Unique filename per call so Gradio's content-
+# hashed cache always serves fresh bytes.
+_ARTIFACT_DIR = Path("/tmp/flux_rgbd_artifacts")
+
+# 2h is comfortably longer than any viewer session; keeps a busy Space's
+# /tmp bounded since nothing else ever deletes these.
+_ARTIFACT_TTL_S = 2 * 3600.0
+
+
+def _prune_old_artifacts() -> None:
+    now = time.time()
+    for f in _ARTIFACT_DIR.glob("cloud_*.glb"):
+        try:
+            if now - f.stat().st_mtime > _ARTIFACT_TTL_S:
+                f.unlink()
+        except OSError:
+            pass  # concurrent delete / fs hiccup — never fail a generation
+
+
+_ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
+
+
+@spaces.GPU(duration=120)
+def _sample_on_gpu(prompt: str, input_image,
+                   num_steps: int, cfg_scale: float, seed: int,
+                   refine_depth: bool = True, log2_alpha: float = 5.0):
+    """GPU-only step: text encode + diffusion sample + VAE decode.
+
+    Returns plain numpy arrays so the parent process (which is what
+    serves Gradio files) can do the rest. Writing the PLY here would
+    leave it in the subprocess's filesystem view where the parent's
+    Gradio file route can't find it (returns 404).
+    """
+    import time
+    runner = _ensure_runner()
+    mode = "i2d" if input_image is not None else "joint"
+
+    target = runner.img_hw[0]
+    if mode == "i2d":
+        letterboxed, (top, left, vh, vw) = _letterbox(input_image, target)
+        model_input = letterboxed
+    else:
+        letterboxed = None
+        top = left = 0
+        vh = vw = target
+        model_input = None
+
+    t0 = time.time()
+    if mode == "i2d":
+        # Image given: single image→depth pass at CFG 1.0 (no guidance — the
+        # RGB is fixed, so there is nothing for CFG to steer).
+        result = runner.generate(
+            prompt.strip() if prompt else "",
+            mode="i2d",
+            num_steps=int(num_steps), cfg_scale=1.0, seed=int(seed),
+            clean_rgb_image=model_input,
+        )
+    else:
+        # Text→RGBD. Stage 1 joint at the requested CFG (default 4.0), rgb-first
+        # trajectory (log2_alpha=5) for cleaner depth. When `refine_depth` is on,
+        # a stage 2 re-derives depth via i2d on that RGB at CFG 1.0 for sharper,
+        # RGB-consistent geometry; otherwise the single joint pass is used.
+        result = runner.generate(
+            prompt.strip() if prompt else "",
+            mode="joint",
+            num_steps=int(num_steps), cfg_scale=float(cfg_scale), seed=int(seed),
+            log2_alpha=float(log2_alpha),
+            refine_depth_i2d=bool(refine_depth), i2d_cfg_scale=1.0,
+        )
+    elapsed = time.time() - t0
+
+    rgb_for_pc = (letterboxed[top:top + vh, left:left + vw] if mode == "i2d"
+                  else result["rgb"])
+    depth = result["depth"]
+    if mode == "i2d":
+        depth = depth[top:top + vh, left:left + vw]
+    return rgb_for_pc, depth, mode, elapsed
+
+
+def generate(prompt: str, input_image, num_steps: int, cfg_scale: float, seed: int,
+             refine_depth: bool = True, log2_alpha: float = 5.0,
+             edge_rtol: float = 0.04, sor: bool = False):
+    """Public Gradio handler. Runs the GPU step then does PLY writing
+    here in the parent process so the file persists for Gradio."""
+    rgb_for_pc, depth, mode, elapsed = _sample_on_gpu(
+        prompt, input_image, num_steps, cfg_scale, seed, refine_depth, log2_alpha,
+    )
+
+    pts, cols = _depth_to_pointcloud(rgb_for_pc, depth, edge_rtol=edge_rtol,
+                                     sor=bool(sor))
+    _prune_old_artifacts()
+    cloud_path = str(_ARTIFACT_DIR / f"cloud_{uuid.uuid4().hex[:12]}.glb")
+    _save_glb(cloud_path, pts, cols)
+
+    valid = (depth > 0) & np.isfinite(depth)
+    if valid.any():
+        d = depth[valid]
+        depth_summary = (
+            f"depth median={float(np.median(d)):.2f}  "
+            f"p5={float(np.percentile(d, 5)):.2f}  "
+            f"p95={float(np.percentile(d, 95)):.2f}"
+        )
+    else:
+        depth_summary = "depth has no valid pixels"
+    status = f"{mode} · {elapsed:.1f} s · {depth_summary} · {pts.shape[0]:,} points"
+    return rgb_for_pc, _depth_to_magma(depth), cloud_path, status
+
+
+# --- Presentation layer ----------------------------------------------------
+# Only the Gradio UI definition lives below. The generation/model code above
+# is untouched.
+
+WORLD_LABS_URL = "https://www.worldlabs.ai"
+_PROJECT_URL = "https://modality-forcing.github.io/"
+_ARXIV_URL = "https://arxiv.org/abs/2606.13676"
+_CODE_URL = "https://github.com/Duisterhof/modality-forcing"
+
+# Editorial monochrome: a fully neutral palette, Inter for body, JetBrains
+# Mono for the small uppercase "eyebrow" labels. The serif display face for
+# the title (Gilda Display) is pulled in via @import in the CSS below.
+_THEME = gr.themes.Default(
+    # System fonts only — no Google-fetched web fonts for the body/mono, which
+    # were loading unreliably (falling back to Arial and looking cheap). The
+    # serif display title uses Gilda Display, pulled in via @import in the CSS.
+    font=(
+        "system-ui",
+        "-apple-system",
+        "Segoe UI",
+        "Helvetica Neue",
+        "Arial",
+        "sans-serif",
+    ),
+    font_mono=(
+        "ui-monospace",
+        "SFMono-Regular",
+        "Menlo",
+        "Consolas",
+        "monospace",
+    ),
+    primary_hue=gr.themes.colors.neutral,
+    secondary_hue=gr.themes.colors.neutral,
+    neutral_hue=gr.themes.colors.neutral,
+).set(
+    # Hairline, low-contrast borders; no heavy shadows or filled labels.
+    block_border_width="1px",
+    block_border_color="*neutral_200",
+    block_background_fill="white",
+    block_shadow="none",
+    block_label_background_fill="transparent",
+    block_label_border_width="0px",
+    block_label_text_weight="500",
+    input_border_width="1px",
+    input_border_color="*neutral_200",
+    input_shadow="none",
+    # Near-black, fully-rounded primary button (pill); white secondary.
+    button_large_radius="*radius_xxl",
+    button_small_radius="*radius_xxl",
+    button_primary_background_fill="#111111",
+    button_primary_background_fill_hover="#1f1f1f",
+    button_primary_text_color="white",
+    button_primary_border_color="#111111",
+    button_secondary_background_fill="white",
+    button_secondary_border_color="rgba(0,0,0,0.16)",
+)
+
+_CSS = """
+@import url('https://fonts.googleapis.com/css2?family=Gilda+Display&display=swap');
+
+/* Warm "paper" canvas everywhere (page + app + container) so the white
+   component cards lift off the background and the layout reads premium
+   rather than flat white-on-white. Body uses the system UI font. */
+html, body, gradio-app, .gradio-container, .gradio-container .gap {
+    background: #f4f3ef !important;
+}
+.gradio-container {
+    color: #141414 !important;
+    max-width: 1120px !important;
+    margin: 0 auto !important;
+    padding: 28px 24px 12px !important;
+    font-family: system-ui, -apple-system, "Segoe UI", "Helvetica Neue", Arial, sans-serif !important;
+}
+
+/* Components become quiet white cards: hairline edge, soft round corners,
+   and a whisper of shadow for depth. (Also overrides this Gradio build's
+   hardcoded 3px black .block border.) */
+.gradio-container .block {
+    border: 1px solid rgba(20,20,20,0.07) !important;
+    border-radius: 16px !important;
+    background: #ffffff !important;
+    box-shadow: 0 1px 2px rgba(20,20,20,0.04),
+                0 12px 28px -18px rgba(20,20,20,0.18) !important;
+}
+
+/* Text/HTML blocks float on the page — no card border, fill, or shadow. */
+.gradio-container .mf-bare {
+    border: 0 !important; background: transparent !important;
+    box-shadow: none !important; padding: 0 !important;
+}
+
+/* Hairline rule separating the masthead from the workspace. */
+.mf-rule {
+    height: 1px; border: 0; background: rgba(20,20,20,0.08);
+    max-width: 1080px; margin: 0.75rem auto 1.5rem;
+}
+
+/* ---- Publication header (mirrors the project page) ---- */
+/* Everything in the masthead is centered. Forced with !important because
+   Gradio's prose CSS otherwise left-aligns <p>/<div> inside gr.HTML. */
+.mf-pub, .mf-pub *, .mf-intro, .mf-intro * { text-align: center !important; }
+.mf-pub { margin: 0.25rem auto 0.25rem; }
+.mf-pub-title {
+    font-weight: 600;
+    font-size: clamp(1.9rem, 4.2vw, 3rem);
+    line-height: 1.13; letter-spacing: -0.01em;
+    color: #363636; margin: 0 auto 0.7em; max-width: 900px;
+}
+.mf-authors {
+    font-size: clamp(1rem, 1.4vw, 1.25rem); line-height: 1.5;
+    color: #363636; margin: 0 auto;
+}
+.mf-authors a { color: #3273dc !important; text-decoration: none !important; }
+.mf-authors a:hover { text-decoration: underline !important; }
+.mf-authors .ab { white-space: nowrap; margin: 0 0.15em; }
+.mf-affil {
+    font-size: clamp(0.95rem, 1.3vw, 1.2rem); color: #363636;
+    margin: 0.5em auto 0;
+}
+.mf-affil .ab { margin: 0 0.6em; }
+.mf-logos {
+    display: flex; justify-content: center; align-items: center;
+    gap: 40px; flex-wrap: wrap; margin: 1.4em auto 0;
+}
+.mf-logos img { height: 78px; }
+.mf-venue { font-weight: 700; color: #363636; margin: 1em auto 0; }
+.mf-links {
+    display: flex; justify-content: center; gap: 12px;
+    flex-wrap: wrap; margin: 1.1em auto 0.25rem;
+}
+.mf-btn {
+    display: inline-flex; align-items: center; gap: 7px;
+    padding: 7px 18px; border-radius: 9999px;
+    background: #363636; color: #ffffff !important;
+    font-size: 14px; font-weight: 500; text-decoration: none !important;
+    transition: background .15s ease;
+}
+.mf-btn:hover { background: #4a4a4a; }
+.mf-pub-sub {
+    font-size: clamp(1rem, 1.3vw, 1.2rem); line-height: 1.5;
+    color: #4a4a4a; max-width: 720px; margin: 1.2em auto 0; font-weight: 400;
+}
+
+/* ---- Quiet helper line + section eyebrows ---- */
+.mf-intro {
+    text-align: center; max-width: 640px; margin: 0.25rem auto 1rem;
+    font-size: 14px; line-height: 1.6; color: #6b6b6b;
+}
+.mf-intro b { color: #111111; font-weight: 600; }
+.mf-sec {
+    font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
+    font-size: 11px; font-weight: 500; letter-spacing: 0.18em;
+    text-transform: uppercase; color: #9a9a9a; margin: 4px 2px 2px;
+}
+
+/* ---- Primary button: full pill, ink-black ---- */
+.gradio-container button.primary,
+.gradio-container button.lg.primary {
+    background: #111111 !important; color: #ffffff !important;
+    border: 1px solid #111111 !important; border-radius: 9999px !important;
+    font-weight: 500 !important; letter-spacing: -0.005em !important;
+}
+.gradio-container button.primary:hover { background: #1f1f1f !important; }
+
+/* ---- Examples: borderless, quiet ---- */
+.gradio-container .examples table,
+.gradio-container .examples .tr-head { border: 0 !important; }
+.gradio-container .examples td {
+    border-color: rgba(0,0,0,0.06) !important; font-size: 13px !important;
+}
+
+/* ---- Footer ---- */
+.mf-footer {
+    text-align: center; margin-top: 2rem; padding-top: 1.1rem;
+    border-top: 1px solid rgba(0,0,0,0.08);
+}
+.mf-footer .mf-cta { font-size: 14px; color: #3d3d3d; }
+.mf-footer .mf-cta a {
+    color: #111111 !important; text-decoration: none !important;
+    border-bottom: 1px solid rgba(0,0,0,0.25);
+}
+.mf-footer .mf-credit {
+    font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
+    font-size: 11px; letter-spacing: 0.08em; text-transform: uppercase;
+    color: #9a9a9a; margin-top: 6px;
+}
+"""
+
+_HEADER_HTML = (
+    '<div class="mf-pub">'
+    '<h1 class="mf-pub-title">Modality Forcing for Scalable Spatial Generation</h1>'
+    '<div class="mf-authors">'
+    '<span class="ab"><a href="https://bart-ai.com" target="_blank" rel="noreferrer">Bardienus Pieter Duisterhof</a><sup>1,2</sup>,</span>'
+    '<span class="ab"><a href="https://www.cs.cmu.edu/~deva/" target="_blank" rel="noreferrer">Deva Ramanan</a><sup>1</sup>,</span>'
+    '<span class="ab"><a href="https://ichnow.ski" target="_blank" rel="noreferrer">Jeffrey Ichnowski</a><sup>1</sup>,</span>'
+    '<span class="ab"><a href="https://web.eecs.umich.edu/~justincj/" target="_blank" rel="noreferrer">Justin Johnson</a><sup>2</sup>,</span>'
+    '<span class="ab"><a href="https://keunhong.com" target="_blank" rel="noreferrer">Keunhong Park</a><sup>2</sup></span>'
+    '</div>'
+    '<div class="mf-affil">'
+    '<span class="ab"><sup>1</sup>Carnegie Mellon University</span>'
+    '<span class="ab"><sup>2</sup>World Labs</span>'
+    '</div>'
+    '<div class="mf-logos">'
+    '<img alt="Carnegie Mellon University" src="https://modality-forcing.github.io/static/images/cmu_logo.png">'
+    '<img alt="World Labs" src="https://modality-forcing.github.io/static/images/world_labs_logo.jpg" style="border-radius:12px;">'
+    '</div>'
+    '<div class="mf-venue">Preprint, 2026</div>'
+    '<div class="mf-links">'
+    f'<a class="mf-btn" href="{_PROJECT_URL}" target="_blank" rel="noopener">📄 Project Page</a>'
+    f'<a class="mf-btn" href="{_ARXIV_URL}" target="_blank" rel="noopener">📚 arXiv</a>'
+    f'<a class="mf-btn" href="{_CODE_URL}" target="_blank" rel="noopener">⌨ Code</a>'
+    '</div>'
+    '<div class="mf-pub-sub" style="text-align:center !important;">Modality '
+    'Forcing turns a pretrained text-to-image diffusion transformer into a '
+    'joint image-depth generator with a simple post-training recipe.</div>'
+    '</div>'
+)
+
+_INTRO_HTML = (
+    '<div class="mf-intro">Type a scene and press <b>Generate</b>, or upload an '
+    'image to run <b>image→depth</b> mode instead.</div>'
+)
+
+_EXAMPLE_IMAGES = [
+    ["assets/examples/alley.png"],
+    ["assets/examples/yosemite.png"],
+    ["assets/examples/castle.png"],
+]
+
+_EXAMPLE_PROMPTS = [
+    [DEFAULT_PROMPT],
+    ["A sunlit Scandinavian living room with a linen sofa, a low oak coffee "
+     "table, and tall windows opening onto a snowy courtyard."],
+    ["A misty pine forest at dawn, shafts of golden light cutting between the "
+     "trunks and a narrow dirt trail winding into the distance."],
+    ["A cozy bookshop interior with floor-to-ceiling wooden shelves, a rolling "
+     "ladder, warm pendant lighting, and a worn leather reading chair."],
+    ["A still life on a marble countertop: a bowl of ripe lemons, a ceramic "
+     "pitcher, and a sprig of rosemary lit by soft side light."],
+    ["A coastal cliffside at golden hour overlooking a turquoise bay, with wild "
+     "grass in the foreground and distant sailboats on the water."],
+]
+
+with gr.Blocks(title="Modality Forcing — World Labs") as demo:
+    gr.HTML(_HEADER_HTML, elem_classes="mf-bare")
+    gr.HTML(_INTRO_HTML, elem_classes="mf-bare")
+    gr.HTML('<hr class="mf-rule">', elem_classes="mf-bare")
+
+    with gr.Row(equal_height=False):
+        with gr.Column(scale=2, min_width=320) as left_col:
+            gr.HTML('<div class="mf-sec">Input</div>', elem_classes="mf-bare")
+            prompt = gr.Textbox(
+                value=DEFAULT_PROMPT,
+                lines=4,
+                label="Scene prompt",
+                placeholder="Describe a scene to generate…",
+            )
+            input_image = gr.Image(
+                label="Optional input image (switches to image→depth mode)",
+                type="numpy", height=200, sources=("upload", "clipboard"),
+            )
+            btn = gr.Button("Generate", variant="primary", size="lg")
+
+            with gr.Accordion("Advanced settings", open=False):
+                with gr.Row():
+                    num_steps = gr.Slider(
+                        1, 80, value=50, step=1, label="Sampling steps")
+                    cfg_scale = gr.Slider(
+                        1.0, 8.0, value=4.0, step=0.1,
+                        label="Guidance (CFG) — text mode only")
+                log2_alpha = gr.Slider(
+                    -5.0, 5.0, value=5.0, step=0.5,
+                    label="log2(alpha) — depth trajectory (>0 rgb-first, "
+                          "cleaner depth; 0 diagonal; <0 depth-first) — text mode")
+                edge_rtol = gr.Slider(
+                    0.0, 0.25, value=0.04, step=0.005,
+                    label="Point-cloud depth-edge mask (rtol) — drop pixels at "
+                          "depth jumps > this; lower = more aggressive, 0 = off")
+                sor_toggle = gr.Checkbox(
+                    value=False,
+                    label="Statistical outlier removal (point cloud) — drops "
+                          "isolated floaters; can over-trim fine structures")
+                seed = gr.Number(value=0, precision=0, label="Seed")
+                refine_depth = gr.Checkbox(
+                    value=True,
+                    label="Refine depth: joint (CFG) → image→depth (CFG 1) "
+                          "— text mode only",
+                )
+
+            status = gr.Textbox(label="Status", interactive=False)
+
+        with gr.Column(scale=3, min_width=380):
+            gr.HTML('<div class="mf-sec">Output</div>', elem_classes="mf-bare")
+            with gr.Row():
+                rgb_out = gr.Image(
+                    label="RGB image", type="numpy", height=320, format="png")
+                depth_out = gr.Image(
+                    label="Depth (disparity, magma)",
+                    type="numpy", height=320, format="png")
+            with gr.Group():
+                cloud_out = gr.Model3D(
+                    label="Interactive 3D point cloud",
+                    clear_color=(0.12, 0.12, 0.14, 1.0),
+                    zoom_speed=0.5, pan_speed=0.5,
+                    height=520,
+                )
+
+    def _run_text_example(prompt_text):
+        # Cached example runs pin the UI defaults so the cache stays valid.
+        return generate(prompt_text, None, 50, 4.0, 0, True, 5.0, 0.04, False)
+
+    def _run_image_example(image):
+        return generate("", image, 50, 4.0, 0, True, 5.0, 0.04, False)
+
+    # Defined after the output components exist, rendered back into the left
+    # column. Cached: clicking an example serves precomputed results instead
+    # of spending GPU time ("lazy" = no cache rebuild at startup; the
+    # committed cache ships with the Space repo).
+    _CACHE_MODE = os.environ.get("EXAMPLES_CACHE_MODE", "lazy")
+    with left_col:
+        gr.Examples(
+            examples=_EXAMPLE_PROMPTS,
+            inputs=[prompt],
+            outputs=[rgb_out, depth_out, cloud_out, status],
+            fn=_run_text_example,
+            cache_examples=True,
+            cache_mode=_CACHE_MODE,
+            label="Example prompts",
+        )
+        gr.Examples(
+            examples=_EXAMPLE_IMAGES,
+            inputs=[input_image],
+            outputs=[rgb_out, depth_out, cloud_out, status],
+            fn=_run_image_example,
+            cache_examples=True,
+            cache_mode=_CACHE_MODE,
+            label="Example images (image → depth)",
+        )
+
+    gr.HTML(
+        '<div class="mf-footer">'
+        '<div class="mf-cta">Built by <a href="' + WORLD_LABS_URL + '" '
+        'target="_blank" rel="noopener noreferrer">World Labs</a></div>'
+        '<div class="mf-credit">worldlabs.ai · Modality Forcing</div>'
+        '</div>',
+        elem_classes="mf-bare",
+    )
+
+    btn.click(
+        generate,
+        inputs=[prompt, input_image, num_steps, cfg_scale, seed, refine_depth,
+                log2_alpha, edge_rtol, sor_toggle],
+        outputs=[rgb_out, depth_out, cloud_out, status],
+    )
+
+    # Adding an image switches the demo into image→depth mode, where the prompt
+    # is optional — clear it so it defaults to empty. The user can retype one.
+    input_image.change(
+        lambda img: "" if img is not None else gr.update(),
+        inputs=[input_image],
+        outputs=[prompt],
+    )
+
+# The UI is designed light-only (paper-white cards, explicit light CSS). A
+# visitor whose OS is in dark mode otherwise gets gradio's dark text colors on
+# our light backgrounds — an unreadable mix. Redirect to ?__theme=light in
+# <head>, before gradio boots, so every visitor gets the light theme.
+_FORCE_LIGHT_HEAD = """
+<script>
+  (function () {
+    const url = new URL(window.location);
+    if (url.searchParams.get("__theme") !== "light") {
+      url.searchParams.set("__theme", "light");
+      window.location.replace(url);
+    }
+  })();
+</script>
+"""
+
+if __name__ == "__main__":
+    # Gradio 6 takes theme + css on launch() (not the Blocks constructor),
+    # so they must be passed here to actually apply on the Space.
+    demo.queue(max_size=4).launch(theme=_THEME, css=_CSS, head=_FORCE_LIGHT_HEAD)

flux_rgbd/__init__.py

@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 World Labs.
+"""flux_rgbd — joint text → RGB + depth diffusion on top of FLUX.2."""
+
+__version__ = "0.1.0"
+__all__ = ["FluxRGBDRunner"]
+
+
+def __getattr__(name):
+    if name == "FluxRGBDRunner":
+        from flux_rgbd.runner import FluxRGBDRunner
+        return FluxRGBDRunner
+    raise AttributeError(name)

flux_rgbd/_flux2/__init__.py

@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Black Forest Labs.
+"""Vendored FLUX.2 reference implementation.
+
+This subpackage holds the unmodified FLUX.2 architecture (Apache-2.0,
+Black Forest Labs upstream). Do not edit these files; they are
+upstream-tracking. All modifications and extensions for RGB+D inference
+live in the parent ``flux_rgbd`` package.
+
+Upstream: https://github.com/black-forest-labs/flux2
+"""

flux_rgbd/_flux2/autoencoder.py

@@ -0,0 +1,499 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Black Forest Labs.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""FLUX.2 autoencoder.
+
+Adapted from the FLUX.2 codebase:
+    https://github.com/black-forest-labs/flux2
+"""
+
+import math
+from collections.abc import Sequence
+
+import torch
+from einops import rearrange
+from jaxtyping import Float
+from safetensors.torch import load_file as load_safetensors
+from torch import Tensor, nn
+
+__all__ = ["Flux2Encoder", "Flux2Decoder"]
+
+# FLUX.2-dev autoencoder weights, pre-split into encoder + decoder
+# safetensors. (BFL's canonical `ae.safetensors` is a combined file
+# we'd have to filter by key prefix — easier to host the split files
+# inside the model repo. Override via the FLUX_RGBD_LOCAL_AE_{...}
+# env vars if you want to skip the Hub fetch entirely.)
+DEFAULT_AE_REPO = "bartduis/modality_forcing"
+DEFAULT_ENCODER_FILE = "ae_encoder.safetensors"
+DEFAULT_DECODER_FILE = "ae_decoder.safetensors"
+
+DEFAULT_RESOLUTION = 256
+DEFAULT_IN_CHANNELS = 3
+DEFAULT_CH = 128
+DEFAULT_OUT_CHANNELS = 3
+DEFAULT_CH_MULT = (1, 2, 4, 4)
+DEFAULT_NUM_RES_BLOCKS = 2
+DEFAULT_Z_CHANNELS = 32
+
+
+def _hf_download(repo_id: str, filename: str) -> str:
+    """Resolve a HuggingFace Hub file path. Lazy-import hf_hub to keep the
+    base package importable in environments that pin huggingface_hub via
+    extras."""
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(repo_id=repo_id, filename=filename)
+
+
+def _load_weights(repo_id: str, filename: str) -> dict[str, Tensor]:
+    # Allow callers to point at a local file via env var. Useful for tests
+    # and air-gapped environments.
+    import os
+    env_key = f"FLUX_RGBD_LOCAL_{filename.upper().replace('.', '_')}"
+    override = os.environ.get(env_key)
+    if override:
+        if not os.path.isfile(override):
+            raise FileNotFoundError(f"{env_key}={override} is set but not a file")
+        return load_safetensors(override)
+    path = _hf_download(repo_id, filename)
+    return load_safetensors(path)
+
+
+def _swish(x: Tensor) -> Tensor:
+    return x * torch.sigmoid(x)
+
+
+class AttnBlock(nn.Module):
+    """Attention block for FLUX.2 autoencoder."""
+
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = nn.GroupNorm(
+            num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+        )
+
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        b, c, h, w = q.shape
+        q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
+        k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
+        v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
+        h_ = nn.functional.scaled_dot_product_attention(q, k, v)
+
+        return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+    """Resnet block for FLUX.2 autoencoder."""
+
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+
+        self.norm1 = nn.GroupNorm(
+            num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+        )
+        self.conv1 = nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        self.norm2 = nn.GroupNorm(
+            num_groups=32, num_channels=out_channels, eps=1e-6, affine=True
+        )
+        self.conv2 = nn.Conv2d(
+            out_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        if self.in_channels != self.out_channels:
+            self.nin_shortcut = nn.Conv2d(
+                in_channels, out_channels, kernel_size=1, stride=1, padding=0
+            )
+
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = _swish(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = _swish(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class Downsample(nn.Module):
+    """Downsample block for FLUX.2 autoencoder."""
+
+    def __init__(self, in_channels: int):
+        super().__init__()
+        # no asymmetric padding in torch conv, must do it ourselves
+        self.conv = nn.Conv2d(
+            in_channels, in_channels, kernel_size=3, stride=2, padding=0
+        )
+
+    def forward(self, x: Tensor):
+        pad = (0, 1, 0, 1)
+        x = nn.functional.pad(x, pad, mode="constant", value=0)
+        x = self.conv(x)
+        return x
+
+
+class Upsample(nn.Module):
+    """Upsample block for FLUX.2 autoencoder."""
+
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(
+            in_channels, in_channels, kernel_size=3, stride=1, padding=1
+        )
+
+    def forward(self, x: Tensor):
+        x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        x = self.conv(x)
+        return x
+
+
+class Flux2Encoder(nn.Module):
+    """Encoder for FLUX.2 autoencoder."""
+
+    def __init__(
+        self,
+        resolution: int = DEFAULT_RESOLUTION,
+        in_channels: int = DEFAULT_IN_CHANNELS,
+        ch: int = DEFAULT_CH,
+        ch_mult: Sequence[int] = DEFAULT_CH_MULT,
+        num_res_blocks: int = DEFAULT_NUM_RES_BLOCKS,
+        z_channels: int = DEFAULT_Z_CHANNELS,
+        repo_id: str = DEFAULT_AE_REPO,
+        filename: str = DEFAULT_ENCODER_FILE,
+    ):
+        """Initialize the FLUX.2 encoder.
+
+        Args:
+            resolution: The resolution of the input images.
+            in_channels: The number of channels in the input images.
+            ch: The number of channels in the encoder.
+            ch_mult: The number of channels in the encoder at each resolution.
+            num_res_blocks: The number of ResNet blocks in each downsampling block.
+            z_channels: The number of channels in the latent space.
+            repo_id: HuggingFace Hub repository id from which to load weights.
+            filename: Filename of the safetensors weights within `repo_id`.
+        """
+        super().__init__()
+        self._repo_id = repo_id
+        self._filename = filename
+
+        self.quant_conv = torch.nn.Conv2d(2 * z_channels, 2 * z_channels, 1)
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        # downsampling
+        self.conv_in = nn.Conv2d(
+            in_channels, self.ch, kernel_size=3, stride=1, padding=1
+        )
+
+        curr_res = resolution
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.in_ch_mult = in_ch_mult
+        self.down = nn.ModuleList()
+        block_in = self.ch
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        # end
+        self.norm_out = nn.GroupNorm(
+            num_groups=32, num_channels=block_in, eps=1e-6, affine=True
+        )
+        self.conv_out = nn.Conv2d(
+            block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1
+        )
+
+        self.bn_eps = 1e-4
+        self.bn_momentum = 0.1
+        self.ps = [2, 2]
+        self.bn = torch.nn.BatchNorm2d(
+            math.prod(self.ps) * z_channels,
+            eps=self.bn_eps,
+            momentum=self.bn_momentum,
+            affine=False,
+            track_running_stats=True,
+        )
+
+    def load_weights(self):
+        """Fetch and load weights from HuggingFace Hub."""
+        self.load_state_dict(_load_weights(self._repo_id, self._filename))
+
+    def _forward(self, x: Float[Tensor, "n cx hx wx"]) -> Float[Tensor, "n cz hz wz"]:
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        # end
+        h = self.norm_out(h)
+        h = _swish(h)
+        h = self.conv_out(h)
+        h = self.quant_conv(h)
+        return h
+
+    def _normalize(self, z: Float[Tensor, "n cz hz wz"]) -> Float[Tensor, "n cz hz wz"]:
+        self.bn.eval()
+        return self.bn(z)
+
+    def forward(
+        self, x: Float[Tensor, "... hx wx cx"]
+    ) -> Float[Tensor, "... hz wz cz"]:
+        """Encode images to latents.
+
+        Args:
+            x: Images to encode.
+
+        Returns:
+            The encoded latents. The latents will be normalized using the running
+                statistics from training.
+        """
+        *batch_shape, _, _, _ = x.shape
+        # Convert to NCHW.
+        x = rearrange(x, "... h w c -> (...) c h w")
+
+        # Compute latents.
+        moments = self._forward(x)
+        mean = torch.chunk(moments, 2, dim=1)[0]
+        z = rearrange(
+            mean,
+            "... c (i pi) (j pj)  -> ... (c pi pj) i j",
+            pi=self.ps[0],
+            pj=self.ps[1],
+        )
+        z = self._normalize(z)
+
+        # Convert back to NHWC and restore batch shape.
+        z = rearrange(z, "... c h w -> ... h w c")
+        return z.reshape(*batch_shape, *z.shape[-3:])
+
+
+class Flux2Decoder(nn.Module):
+    """Decoder for FLUX.2 autoencoder."""
+
+    def __init__(
+        self,
+        ch: int = DEFAULT_CH,
+        out_ch: int = DEFAULT_OUT_CHANNELS,
+        ch_mult: Sequence[int] = DEFAULT_CH_MULT,
+        num_res_blocks: int = DEFAULT_NUM_RES_BLOCKS,
+        in_channels: int = DEFAULT_IN_CHANNELS,
+        resolution: int = DEFAULT_RESOLUTION,
+        z_channels: int = DEFAULT_Z_CHANNELS,
+        repo_id: str = DEFAULT_AE_REPO,
+        filename: str = DEFAULT_DECODER_FILE,
+    ):
+        """Initialize the FLUX.2 decoder.
+
+        Args:
+            ch: The number of channels in the decoder.
+            out_ch: The number of channels in the output.
+            ch_mult: The number of channels in the decoder at each resolution.
+            num_res_blocks: The number of ResNet blocks in each upsampling block.
+            in_channels: The number of channels in the input images.
+            resolution: The resolution of the input images.
+            z_channels: The number of channels in the latent space.
+            repo_id: HuggingFace Hub repository id from which to load weights.
+            filename: Filename of the safetensors weights within `repo_id`.
+        """
+        super().__init__()
+        self._repo_id = repo_id
+        self._filename = filename
+
+        self.post_quant_conv = torch.nn.Conv2d(z_channels, z_channels, 1)
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.ffactor = 2 ** (self.num_resolutions - 1)
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+
+        # z to block_in
+        self.conv_in = nn.Conv2d(
+            z_channels, block_in, kernel_size=3, stride=1, padding=1
+        )
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)  # prepend to get consistent order
+
+        # end
+        self.norm_out = nn.GroupNorm(
+            num_groups=32, num_channels=block_in, eps=1e-6, affine=True
+        )
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+        self.bn_eps = 1e-4
+        self.bn_momentum = 0.1
+        self.ps = [2, 2]
+        self.bn = torch.nn.BatchNorm2d(
+            math.prod(self.ps) * z_channels,
+            eps=self.bn_eps,
+            momentum=self.bn_momentum,
+            affine=False,
+            track_running_stats=True,
+        )
+
+    def load_weights(self):
+        """Fetch and load weights from HuggingFace Hub."""
+        self.load_state_dict(_load_weights(self._repo_id, self._filename))
+
+    def _forward(self, z: Float[Tensor, "n cz hz wz"]) -> Float[Tensor, "n cx hx wx"]:
+        z = self.post_quant_conv(z)
+
+        # get dtype for proper tracing
+        upscale_dtype = next(self.up.parameters()).dtype
+
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        # cast to proper dtype
+        h = h.to(upscale_dtype)
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = _swish(h)
+        h = self.conv_out(h)
+        return h
+
+    def _inv_normalize(
+        self, z: Float[Tensor, "n cz hz wz"]
+    ) -> Float[Tensor, "n cz hz wz"]:
+        self.bn.eval()
+        s = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + self.bn_eps)
+        m = self.bn.running_mean.view(1, -1, 1, 1)
+        return z * s + m
+
+    def forward(
+        self, z: Float[Tensor, "... cz hz wz"]
+    ) -> Float[Tensor, "... cx hx wx"]:
+        """Decode latents to images.
+
+        Args:
+            z: Image latents normalized to have zero mean and unit variance.
+
+        Returns:
+            Decoded images.
+        """
+        *batch_shape, _, _, _ = z.shape
+
+        # Convert to NCHW.
+        z = rearrange(z, "... h w c -> (...) c h w")
+
+        # Denormalize and evaluate decoder.
+        z = self._inv_normalize(z)
+        z = rearrange(
+            z,
+            "... (c pi pj) i j -> ... c (i pi) (j pj)",
+            pi=self.ps[0],
+            pj=self.ps[1],
+        )
+        dec = self._forward(z)
+
+        # Convert back to NHWC and restore batch shape.
+        dec = rearrange(dec, "... c h w -> ... h w c")
+        dec = dec.reshape(*batch_shape, *dec.shape[-3:])
+        return dec

flux_rgbd/_flux2/constants.py

@@ -0,0 +1,207 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Black Forest Labs.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Adapted from the FLUX.2 codebase:
+#     https://github.com/black-forest-labs/flux2
+"""
+Architecture parameters and other constants useful for FLUX.2 inference.
+
+This module centralizes all constants used across the FLUX.2 model implementation,
+including model architecture configurations for all public FLUX.2 variants, their
+text encoder parameters, and system prompts for prompt upsampling.
+"""
+
+from types import MappingProxyType
+
+__all__ = [
+    "FLUX2_DEV_PARAMS",
+    "FLUX2_KLEIN_4B_PARAMS",
+    "FLUX2_KLEIN_9B_PARAMS",
+    "OUTPUT_LAYERS_MISTRAL",
+    "OUTPUT_LAYERS_QWEN3",
+    "MAX_LENGTH",
+    "UPSAMPLING_MAX_IMAGE_SIZE",
+    "SYSTEM_MESSAGE",
+    "SYSTEM_MESSAGE_UPSAMPLING_T2I",
+    "SYSTEM_MESSAGE_UPSAMPLING_I2I",
+]
+
+# =============================================================================
+# Model Architecture Parameters
+# =============================================================================
+
+FLUX2_DEV_PARAMS = MappingProxyType(
+    {
+        "in_channels": 128,
+        "context_in_dim": 15360,
+        "hidden_size": 6144,
+        "num_heads": 48,
+        "depth": 8,
+        "depth_single_blocks": 48,
+        "axes_dim": (32, 32, 32, 32),
+        "theta": 2000,
+        "mlp_ratio": 3.0,
+        "use_guidance_embed": True,
+    }
+)
+"""
+Architecture config for FLUX.2 [dev] 32B parameter model (guidance distilled).
+"""
+
+FLUX2_KLEIN_4B_PARAMS = MappingProxyType(
+    {
+        "in_channels": 128,
+        "context_in_dim": 7680,
+        "hidden_size": 3072,
+        "num_heads": 24,
+        "depth": 5,
+        "depth_single_blocks": 20,
+        "axes_dim": (32, 32, 32, 32),
+        "theta": 2000,
+        "mlp_ratio": 3.0,
+        "use_guidance_embed": False,
+    }
+)
+"""
+Architecture config for FLUX.2 [klein] 4B parameter model (size distilled).
+"""
+
+FLUX2_KLEIN_9B_PARAMS = MappingProxyType(
+    {
+        "in_channels": 128,
+        "context_in_dim": 12288,
+        "hidden_size": 4096,
+        "num_heads": 32,
+        "depth": 8,
+        "depth_single_blocks": 24,
+        "axes_dim": (32, 32, 32, 32),
+        "theta": 2000,
+        "mlp_ratio": 3.0,
+        "use_guidance_embed": False,
+    }
+)
+"""
+Immutable architecture configuration for FLUX.2 [klein]-9B model.
+
+This is a mid-size 9B parameter model that balances quality and speed. It uses
+Qwen3-8B as text encoder and provides higher quality than the 4B variant while
+maintaining fast 4-step inference. The model is guidance-distilled.
+
+Key architectural features:
+    - 24 single-stream transformer blocks (between 4B and dev)
+    - 4096 hidden dimensions with 32 attention heads
+    - No guidance embeddings (guidance-distilled into weights)
+    - Context dimension 12288 matches concatenated Qwen3-8B outputs
+
+Default generation parameters (distilled variants):
+    - guidance: 1.0 (fixed, baked into weights)
+    - num_steps: 4 (fixed)
+
+Base variant generation parameters:
+    - guidance: 4.0 (adjustable)
+    - num_steps: 50 (adjustable)
+"""
+
+# =============================================================================
+# Text Encoder Configuration
+# =============================================================================
+
+OUTPUT_LAYERS_MISTRAL = (10, 20, 30)
+"""
+Layer indices to extract hidden states from Mistral-Small-3.2-24B text encoder.
+
+The FLUX.2 [dev] model uses Mistral-Small as its text encoder and concatenates
+hidden states from these three intermediate layers to form rich text embeddings.
+"""
+
+OUTPUT_LAYERS_QWEN3 = (9, 18, 27)
+"""
+Layer indices to extract hidden states from Qwen3 text encoder variants.
+
+The FLUX.2 [klein] models (4B and 9B) use Qwen3 text encoders and concatenate
+hidden states from these three intermediate layers to form text embeddings.
+"""
+
+MAX_LENGTH = 512
+"""
+Maximum sequence length for text encoder tokenization.
+
+Text prompts are truncated or padded to this length during encoding. This
+applies to both Mistral-Small and Qwen3 text encoders.
+"""
+
+UPSAMPLING_MAX_IMAGE_SIZE = 768**2
+"""
+Maximum pixel area for images during prompt upsampling.
+
+When images are provided as conditioning for prompt upsampling (image-to-image
+mode), they are resized to have at most this many pixels while preserving
+aspect ratio. This prevents out-of-memory errors when processing large images.
+"""
+
+# =============================================================================
+# System Messages for Prompt Upsampling
+# =============================================================================
+
+SYSTEM_MESSAGE = """You are an AI that reasons about image descriptions. You give \
+structured responses focusing on object relationships, object attribution and \
+actions without speculation."""
+"""
+Default system message for text encoder when generating image descriptions.
+
+This system message is used for general-purpose text encoding tasks where the
+model needs to reason about images in a structured, factual manner.
+"""
+
+SYSTEM_MESSAGE_UPSAMPLING_T2I = """You are an expert prompt engineer for FLUX.2 \
+by Black Forest Labs. Rewrite user prompts to be more descriptive while \
+strictly preserving their core subject and intent.
+
+Guidelines:
+1. Structure: Keep structured inputs structured (enhance within fields). \
+Convert natural language to detailed paragraphs.
+2. Details: Add concrete visual specifics - form, scale, textures, materials, \
+lighting (quality, direction, color), shadows, spatial relationships, and \
+environmental context.
+3. Text in Images: Put ALL text in quotation marks, matching the prompt's \
+language. Always provide explicit quoted text for objects that would contain \
+text in reality (signs, labels, screens, etc.) - without it, the model \
+generates gibberish.
+
+Output only the revised prompt and nothing else."""
+"""
+System message for text-to-image prompt upsampling.
+
+This prompt engineering template guides the text encoder to expand sparse user
+prompts into detailed, visually-rich descriptions that improve image generation
+quality. It emphasizes adding concrete visual details while preserving the
+original intent.
+"""
+
+SYSTEM_MESSAGE_UPSAMPLING_I2I = """You are FLUX.2 by Black Forest Labs, an \
+image-editing expert. You convert editing requests into one concise \
+instruction (50-80 words, ~30 for brief requests).
+
+Rules:
+- Single instruction only, no commentary
+- Use clear, analytical language (avoid "whimsical," "cascading," etc.)
+- Specify what changes AND what stays the same (face, lighting, composition)
+- Reference actual image elements
+- Turn negatives into positives ("don't change X" → "keep X")
+- Make abstractions concrete ("futuristic" → "glowing cyan neon, metallic panels")
+- Keep content PG-13
+
+Output only the final instruction in plain text and nothing else."""
+"""
+System message for image-to-image prompt upsampling.
+
+This prompt engineering template guides the text encoder to convert free-form
+image editing requests into concise, actionable instructions suitable for
+image-to-image generation tasks.
+"""

flux_rgbd/_flux2/model.py

@@ -0,0 +1,674 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Black Forest Labs.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Adapted from the FLUX.2 codebase:
+#     https://github.com/black-forest-labs/flux2
+"""
+FLUX.2 diffusion transformer architecture for image generation and editing.
+
+This module implements the core transformer architecture for FLUX.2 models from
+Black Forest Labs. The architecture uses dual-stream and single-stream transformer
+blocks to process text and image latents for text-to-image and image-to-image tasks.
+"""
+
+import math
+
+import einops
+import torch
+from torch import Tensor, nn
+
+
+class Flux2(nn.Module):
+    """
+    FLUX.2 diffusion transformer for image generation and editing.
+
+    This is a flow-matching diffusion model that uses a stack of dual-stream and
+    single-stream transformer blocks for text context and image latents. The
+    model supports text-to-image and image-to-image generation tasks.
+
+    Default parameter values match the FLUX.2 [klein] 4B architecture, which is
+    optimized for fast inference. For other model variants (klein-9B or dev),
+    use the parameters from :mod:`flux_rgbd._flux2.constants`.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 128,
+        context_in_dim: int = 7680,
+        hidden_size: int = 3072,
+        num_heads: int = 24,
+        depth: int = 5,
+        depth_single_blocks: int = 20,
+        axes_dim: tuple[int, int, int, int] = (32, 32, 32, 32),
+        theta: int = 2000,
+        mlp_ratio: float = 3.0,
+        use_guidance_embed: bool = False,
+    ):
+        """
+        Args:
+            in_channels: Number of input channels for image latents. Matches the
+                output dimension of the autoencoder used to encode images.
+            context_in_dim: Dimension of text context embeddings from the text
+                encoder. This should match the concatenated output dimension of
+                the text encoder being used (7680 for Qwen3-4B, 12288 for Qwen3-8B,
+                15360 for Mistral-Small).
+            hidden_size: Hidden dimension size for transformer blocks. All
+                attention and MLP operations use this dimension internally.
+            num_heads: Number of attention heads in multi-head attention layers.
+                Must evenly divide `hidden_size`.
+            depth: Number of dual-stream transformer blocks. These blocks process
+                text and image streams separately with cross-attention.
+            depth_single_blocks: Number of single-stream transformer blocks. These
+                blocks process the concatenated text+image sequence.
+            axes_dim: Tuple of 4 integers specifying the dimensionality for each
+                axis in rotary embeddings. Must sum to `hidden_size / num_heads`.
+            theta: Base frequency for rotary position embeddings (RoPE). Higher
+                values result in slower position encoding rotation.
+            mlp_ratio: Expansion ratio for MLP hidden dimension relative to
+                `hidden_size`. MLP hidden dim = `hidden_size * mlp_ratio`.
+            use_guidance_embed: Whether to include guidance scale embeddings for
+                classifier-free guidance. Set to False for distilled models where
+                guidance is baked into weights.
+        """
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        if hidden_size % num_heads != 0:
+            raise ValueError(
+                f"Hidden size {hidden_size} must be divisible by "
+                f"num_heads {num_heads}"
+            )
+        pe_dim = hidden_size // num_heads
+        if sum(axes_dim) != pe_dim:
+            raise ValueError(f"Got {axes_dim} but expected positional dim {pe_dim}")
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.pe_embedder = EmbedND(dim=pe_dim, theta=theta, axes_dim=axes_dim)
+        self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=False)
+        self.time_in = MLPEmbedder(
+            in_dim=256, hidden_dim=self.hidden_size, disable_bias=True
+        )
+        self.txt_in = nn.Linear(context_in_dim, self.hidden_size, bias=False)
+
+        self.use_guidance_embed = use_guidance_embed
+        if self.use_guidance_embed:
+            self.guidance_in = MLPEmbedder(
+                in_dim=256, hidden_dim=self.hidden_size, disable_bias=True
+            )
+
+        double_blocks = [
+            DoubleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio)
+            for _ in range(depth)
+        ]
+        self.double_blocks = nn.ModuleList(double_blocks)
+
+        single_blocks = [
+            SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio)
+            for _ in range(depth_single_blocks)
+        ]
+        self.single_blocks = nn.ModuleList(single_blocks)
+
+        self.double_stream_modulation_img = Modulation(
+            self.hidden_size, double=True, disable_bias=True
+        )
+        self.double_stream_modulation_txt = Modulation(
+            self.hidden_size, double=True, disable_bias=True
+        )
+        self.single_stream_modulation = Modulation(
+            self.hidden_size, double=False, disable_bias=True
+        )
+
+        self.final_layer = LastLayer(self.hidden_size, self.out_channels)
+
+    def forward(
+        self,
+        x: Tensor,
+        x_ids: Tensor,
+        timesteps: Tensor,
+        ctx: Tensor,
+        ctx_ids: Tensor,
+        guidance: Tensor | None,
+    ):
+        num_txt_tokens = ctx.shape[1]
+
+        timestep_emb = timestep_embedding(timesteps, 256)
+        vec = self.time_in(timestep_emb)
+        if self.use_guidance_embed:
+            guidance_emb = timestep_embedding(guidance, 256)
+            vec = vec + self.guidance_in(guidance_emb)
+
+        double_block_mod_img = self.double_stream_modulation_img(vec)
+        double_block_mod_txt = self.double_stream_modulation_txt(vec)
+        single_block_mod, _ = self.single_stream_modulation(vec)
+
+        img = self.img_in(x)
+        txt = self.txt_in(ctx)
+
+        pe_x = self.pe_embedder(x_ids)
+        pe_ctx = self.pe_embedder(ctx_ids)
+
+        for block in self.double_blocks:
+            img, txt = block(
+                img,
+                txt,
+                pe_x,
+                pe_ctx,
+                double_block_mod_img,
+                double_block_mod_txt,
+            )
+
+        img = torch.cat((txt, img), dim=1)
+        pe = torch.cat((pe_ctx, pe_x), dim=2)
+
+        for block in self.single_blocks:
+            img = block(
+                img,
+                pe,
+                single_block_mod,
+            )
+
+        img = img[:, num_txt_tokens:, ...]
+
+        img = self.final_layer(img, vec)
+        return img
+
+
+class SelfAttention(nn.Module):
+    """
+    Multi-head self-attention with QK normalization.
+
+    This module computes query, key, and value projections in a single linear layer,
+    applies RMS normalization to queries and keys, then performs attention and projects
+    the output back to the original dimension.
+    """
+
+    def __init__(self, dim: int, num_heads: int = 8):
+        """
+        Args:
+            dim: Hidden dimension size. Must be divisible by `num_heads`.
+            num_heads: Number of parallel attention heads.
+        """
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.qkv = nn.Linear(dim, dim * 3, bias=False)
+
+        self.norm = QKNorm(head_dim)
+        self.proj = nn.Linear(dim, dim, bias=False)
+
+
+class SiLUActivation(nn.Module):
+    """
+    Gated activation using SiLU (Swish) function.
+
+    This module splits the input tensor along the last dimension, applies SiLU to one
+    half, and element-wise multiplies it with the other half. This is commonly known
+    as SwiGLU when used in MLP layers.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.gate_fn = nn.SiLU()
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x: Input tensor of shape `(..., 2 * dim)` where the last dimension will
+                be split into two equal parts for gating.
+
+        Returns:
+            Gated output of shape `(..., dim)`.
+        """
+        x1, x2 = x.chunk(2, dim=-1)
+        return self.gate_fn(x1) * x2
+
+
+class Modulation(nn.Module):
+    """
+    Adaptive layer normalization (AdaLN) modulation layer.
+
+    This module generates scale, shift, and gate parameters for adaptive normalization
+    from timestep/guidance embeddings. For double-stream blocks, it produces two sets
+    of modulation parameters (one for each stream).
+    """
+
+    def __init__(self, dim: int, double: bool, disable_bias: bool = False):
+        """
+        Args:
+            dim: Hidden dimension size matching the transformer blocks.
+            double: If True, generates parameters for dual-stream blocks (6 params:
+                shift, scale, gate for each stream). If False, generates for
+                single-stream blocks (3 params: shift, scale, gate).
+            disable_bias: If True, the linear layer has no bias term.
+        """
+        super().__init__()
+        self.is_double = double
+        self.multiplier = 6 if double else 3
+        self.lin = nn.Linear(dim, self.multiplier * dim, bias=not disable_bias)
+
+    def forward(self, vec: Tensor):
+        """
+        Args:
+            vec: Timestep/guidance embedding of shape `(batch_size, dim)` or
+                `(batch_size, seq_len, dim)`.
+
+        Returns:
+            Tuple of modulation parameters. For single-stream: `(mod, None)` where
+            `mod` is a 3-tuple of (shift, scale, gate). For double-stream:
+            `(mod1, mod2)` where each is a 3-tuple for different streams.
+        """
+        out = self.lin(nn.functional.silu(vec))
+        if out.ndim == 2:
+            out = out[:, None, :]
+        out = out.chunk(self.multiplier, dim=-1)
+        return out[:3], out[3:] if self.is_double else None
+
+
+class LastLayer(nn.Module):
+    """
+    Final output layer with adaptive layer normalization.
+
+    This module applies AdaLN-modulated normalization followed by a linear projection
+    to map transformer hidden states back to the output space (image latent channels).
+    """
+
+    def __init__(self, hidden_size: int, out_channels: int):
+        """
+        Args:
+            hidden_size: Hidden dimension of transformer blocks.
+            out_channels: Output dimension (number of latent channels).
+        """
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=False)
+        self.adaLN_modulation = nn.Sequential(  # pylint: disable=invalid-name
+            nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=False)
+        )
+
+    def forward(self, x: Tensor, vec: Tensor) -> Tensor:
+        """
+        Args:
+            x: Hidden states from transformer as a tensor of shape
+                `(batch_size, seq_len, hidden_size)`.
+            vec: Timestep embedding of shape `(batch_size, hidden_size)`.
+
+        Returns:
+            Output tensor of shape `(batch_size, seq_len, out_channels)`.
+        """
+        mod = self.adaLN_modulation(vec)
+        shift, scale = mod.chunk(2, dim=-1)
+        if shift.ndim == 2:
+            shift = shift[:, None, :]
+            scale = scale[:, None, :]
+        x = (1 + scale) * self.norm_final(x) + shift
+        x = self.linear(x)
+        return x
+
+
+class SingleStreamBlock(nn.Module):
+    """
+    Single-stream transformer block processing concatenated text+image tokens.
+
+    This block applies self-attention and MLP operations to a unified sequence of
+    text and image tokens. Both operations share pre-normalization and use adaptive
+    modulation from timestep embeddings.
+    """
+
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0):
+        """
+        Args:
+            hidden_size: Hidden dimension size for all linear projections.
+            num_heads: Number of attention heads. Must divide `hidden_size` evenly.
+            mlp_ratio: Ratio of MLP hidden dimension to `hidden_size`. The actual
+                MLP hidden dim is `int(hidden_size * mlp_ratio)`.
+        """
+        super().__init__()
+
+        self.hidden_dim = hidden_size
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.scale = head_dim**-0.5
+        self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.mlp_mult_factor = 2
+
+        self.linear1 = nn.Linear(
+            hidden_size,
+            hidden_size * 3 + self.mlp_hidden_dim * self.mlp_mult_factor,
+            bias=False,
+        )
+        self.linear2 = nn.Linear(
+            hidden_size + self.mlp_hidden_dim, hidden_size, bias=False
+        )
+        self.norm = QKNorm(head_dim)
+
+        self.hidden_size = hidden_size
+        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+
+        self.mlp_act = SiLUActivation()
+
+    def forward(self, x: Tensor, pe: Tensor, mod: tuple[Tensor, Tensor]) -> Tensor:
+        mod_shift, mod_scale, mod_gate = mod
+        x_mod = (1 + mod_scale) * self.pre_norm(x) + mod_shift
+
+        qkv, mlp = torch.split(
+            self.linear1(x_mod),
+            [3 * self.hidden_size, self.mlp_hidden_dim * self.mlp_mult_factor],
+            dim=-1,
+        )
+
+        q, k, v = einops.rearrange(
+            qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
+        )
+        q, k = self.norm(q, k, v)
+
+        attn = attention(q, k, v, pe)
+
+        # Compute activation in mlp stream, cat again and run second linear layer.
+        output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
+        return x + mod_gate * output
+
+
+class DoubleStreamBlock(nn.Module):
+    """
+    Dual-stream transformer block processing text and image tokens separately.
+
+    This block maintains separate streams for text and image tokens, each with their
+    own self-attention and MLP sublayers. Cross-stream information exchange happens
+    through joint attention where Q, K, V from both streams are concatenated.
+    """
+
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float):
+        """
+        Args:
+            hidden_size: Hidden dimension size for all linear projections.
+            num_heads: Number of attention heads. Must divide `hidden_size` evenly.
+            mlp_ratio: Ratio of MLP hidden dimension to `hidden_size`.
+        """
+        super().__init__()
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.num_heads = num_heads
+        assert (
+            hidden_size % num_heads == 0
+        ), f"{hidden_size=} must be divisible by {num_heads=}"
+
+        self.hidden_size = hidden_size
+        self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp_mult_factor = 2
+
+        self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads)
+
+        self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_mlp = nn.Sequential(
+            nn.Linear(hidden_size, mlp_hidden_dim * self.mlp_mult_factor, bias=False),
+            SiLUActivation(),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=False),
+        )
+
+        self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads)
+
+        self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.txt_mlp = nn.Sequential(
+            nn.Linear(
+                hidden_size,
+                mlp_hidden_dim * self.mlp_mult_factor,
+                bias=False,
+            ),
+            SiLUActivation(),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=False),
+        )
+
+    def forward(
+        self,
+        img: Tensor,
+        txt: Tensor,
+        pe: Tensor,
+        pe_ctx: Tensor,
+        mod_img: tuple[Tensor, Tensor],
+        mod_txt: tuple[Tensor, Tensor],
+    ) -> tuple[Tensor, Tensor]:
+        img_mod1, img_mod2 = mod_img
+        txt_mod1, txt_mod2 = mod_txt
+
+        img_mod1_shift, img_mod1_scale, img_mod1_gate = img_mod1
+        img_mod2_shift, img_mod2_scale, img_mod2_gate = img_mod2
+        txt_mod1_shift, txt_mod1_scale, txt_mod1_gate = txt_mod1
+        txt_mod2_shift, txt_mod2_scale, txt_mod2_gate = txt_mod2
+
+        # Prepare image for attention.
+        img_modulated = self.img_norm1(img)
+        img_modulated = (1 + img_mod1_scale) * img_modulated + img_mod1_shift
+
+        img_qkv = self.img_attn.qkv(img_modulated)
+        img_q, img_k, img_v = einops.rearrange(
+            img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
+        )
+        img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
+
+        # Prepare txt for attention.
+        txt_modulated = self.txt_norm1(txt)
+        txt_modulated = (1 + txt_mod1_scale) * txt_modulated + txt_mod1_shift
+
+        txt_qkv = self.txt_attn.qkv(txt_modulated)
+        txt_q, txt_k, txt_v = einops.rearrange(
+            txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
+        )
+        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
+
+        q = torch.cat((txt_q, img_q), dim=2)
+        k = torch.cat((txt_k, img_k), dim=2)
+        v = torch.cat((txt_v, img_v), dim=2)
+
+        pe = torch.cat((pe_ctx, pe), dim=2)
+        attn = attention(q, k, v, pe)
+        txt_attn, img_attn = attn[:, : txt_q.shape[2]], attn[:, txt_q.shape[2] :]
+
+        # Calculate the img blocks.
+        img = img + img_mod1_gate * self.img_attn.proj(img_attn)
+        img = img + img_mod2_gate * self.img_mlp(
+            (1 + img_mod2_scale) * (self.img_norm2(img)) + img_mod2_shift
+        )
+
+        # Calculate the txt blocks.
+        txt = txt + txt_mod1_gate * self.txt_attn.proj(txt_attn)
+        txt = txt + txt_mod2_gate * self.txt_mlp(
+            (1 + txt_mod2_scale) * (self.txt_norm2(txt)) + txt_mod2_shift
+        )
+        return img, txt
+
+
+class MLPEmbedder(nn.Module):
+    """
+    Two-layer MLP for embedding timestep and guidance values.
+
+    This simple MLP transforms scalar timestep or guidance embeddings (after
+    sinusoidal encoding) into the transformer's hidden dimension space.
+    """
+
+    def __init__(self, in_dim: int, hidden_dim: int, disable_bias: bool = False):
+        """
+        Args:
+            in_dim: Input dimension (typically 256 for sinusoidal embeddings).
+            hidden_dim: Output hidden dimension matching transformer blocks.
+            disable_bias: If True, linear layers have no bias terms.
+        """
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=not disable_bias)
+        self.silu = nn.SiLU()
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=not disable_bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x: Input embeddings of shape `(batch_size, in_dim)`.
+
+        Returns:
+            Projected embeddings of shape `(batch_size, hidden_dim)`.
+        """
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class EmbedND(nn.Module):
+    """
+    N-dimensional rotary position embeddings (RoPE) for spatial-temporal tokens.
+
+    This module creates rotary embeddings for multi-dimensional position indices
+    (e.g., time, height, width, sequence). Each dimension gets its own embedding
+    component with configurable dimensions.
+    """
+
+    def __init__(self, dim: int, theta: int, axes_dim: list[int]):
+        """
+        Args:
+            dim: Total position embedding dimension. Should equal `sum(axes_dim)`.
+            theta: Base frequency for RoPE. Higher values give slower rotation.
+            axes_dim: Dimension allocation for each position axis.
+        """
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: Tensor) -> Tensor:
+        """
+        Args:
+            ids: Position indices of shape `(..., num_axes)` where `num_axes`
+                matches `len(axes_dim)`.
+
+        Returns:
+            Rotary embeddings of shape `(..., 1, sum(axes_dim), 2, 2)` suitable
+            for applying rotation to query and key tensors.
+        """
+        emb = torch.cat(
+            [
+                rope(ids[..., i], self.axes_dim[i], self.theta)
+                for i in range(len(self.axes_dim))
+            ],
+            dim=-3,
+        )
+
+        return emb.unsqueeze(1)
+
+
+def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
+    """
+    Create sinusoidal timestep embeddings.
+    :param t: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an (N, D) Tensor of positional embeddings.
+    """
+    t = time_factor * t
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period)
+        * torch.arange(start=0, end=half, device=t.device, dtype=torch.float32)
+        / half
+    )
+
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(t)
+    return embedding
+
+
+class RMSNorm(torch.nn.Module):
+    """
+    Root Mean Square Layer Normalization.
+
+    RMSNorm normalizes using only the variance (RMS) without centering by mean,
+    providing a simpler and often equally effective alternative to LayerNorm.
+    """
+
+    def __init__(self, dim: int):
+        """
+        Args:
+            dim: Dimension to normalize over (last dimension of input).
+        """
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: Tensor):
+        """
+        Args:
+            x: Input tensor of shape `(..., dim)`.
+
+        Returns:
+            Normalized tensor of same shape as input.
+        """
+        x_dtype = x.dtype
+        x = x.float()
+        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
+        return (x * rrms).to(dtype=x_dtype) * self.scale
+
+
+class QKNorm(torch.nn.Module):
+    """
+    Separate RMSNorm for query and key tensors in attention.
+
+    Normalizing queries and keys independently before attention computation improves
+    training stability and can lead to better performance.
+    """
+
+    def __init__(self, dim: int):
+        """
+        Args:
+            dim: Head dimension for queries and keys.
+        """
+        super().__init__()
+        self.query_norm = RMSNorm(dim)
+        self.key_norm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        """
+        Args:
+            q: Query tensor of shape `(..., head_dim)`.
+            k: Key tensor of shape `(..., head_dim)`.
+            v: Value tensor (used only for dtype matching).
+
+        Returns:
+            Tuple of normalized `(query, key)` tensors with dtype matching `v`.
+        """
+        q = self.query_norm(q)
+        k = self.key_norm(k)
+        return q.to(v), k.to(v)
+
+
+def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
+    q, k = apply_rope(q, k, pe)
+
+    x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+    x = einops.rearrange(x, "B H L D -> B L (H D)")
+
+    return x
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+    assert dim % 2 == 0
+    scale = torch.arange(0, dim, 2, dtype=pos.dtype, device=pos.device) / dim
+    omega = 1.0 / (theta**scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack(
+        [torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1
+    )
+    out = einops.rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+    return out.float()
+
+
+def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)

flux_rgbd/depth/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 World Labs.
+"""Depth-side preprocessing and schedule helpers for flux_rgbd."""
+
+from flux_rgbd.depth.preprocess import DepthConfig, decode_depth
+from flux_rgbd.depth.schedule import ScheduleConfig, rollout_timesteps
+
+__all__ = ["DepthConfig", "decode_depth", "ScheduleConfig", "rollout_timesteps"]

flux_rgbd/depth/preprocess.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 World Labs.
+"""Depth-side decoding: model tokens → depth pixel map.
+
+Decoding steps:
+  1. Unpatchify the (tok_h, tok_w, patch_h * patch_w) token grid.
+  2. Divide out the post-normalization scalar `depth_value_scale`.
+  3. Apply the inverse of every recorded normalization stage, in reverse.
+
+Only `"contract"` (the mip-NeRF 360 radial squash, arXiv 2111.12077
+Eq. 10) is genuinely invertible — the other stages discard per-sample
+statistics during encoding and pass through unchanged on decode.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+
+import einops
+import torch
+from torch import Tensor
+
+_KNOWN_STAGES = frozenset({"contract", "normal", "mean", "unit_mean", "none"})
+
+
+@dataclasses.dataclass(frozen=True)
+class DepthConfig:
+    """How depth was normalized at training time. Decoder reads these."""
+
+    depth_normalize_mode: tuple[str, ...] | str = ("unit_mean", "contract")
+    patch_size: int = 16
+    depth_value_scale: float = 1.0
+
+    def __post_init__(self) -> None:
+        stages = self.depth_normalize_mode
+        if isinstance(stages, str):
+            object.__setattr__(self, "depth_normalize_mode", (stages,))
+        unknown = set(self.depth_normalize_mode) - _KNOWN_STAGES
+        if unknown:
+            raise ValueError(f"unknown stages {sorted(unknown)}; expected {sorted(_KNOWN_STAGES)}")
+        if self.depth_value_scale <= 0:
+            raise ValueError(f"depth_value_scale must be > 0; got {self.depth_value_scale}")
+        if self.patch_size <= 0:
+            raise ValueError(f"patch_size must be > 0; got {self.patch_size}")
+
+
+def _inv_contract(z: Tensor) -> Tensor:
+    """Inverse of the mip-NeRF 360 radial contract bijection (operates on last dim).
+
+    Contract maps R^n into the ball of radius 2; this undoes it. For |z| <= 1
+    the contract was identity, so this is also identity in that region.
+    """
+    mag_sq = torch.clamp(z.square().sum(dim=-1, keepdim=True), min=1.0)
+    return z / (2.0 * mag_sq.sqrt() - mag_sq)
+
+
+def _undo_stage(depth: Tensor, stage: str) -> Tensor:
+    if stage != "contract":  # "normal"/"mean"/"unit_mean"/"none" are pass-through.
+        return depth
+    if depth.shape[-1] == 1:
+        return _inv_contract(depth)
+    return _inv_contract(depth.unsqueeze(-1)).squeeze(-1)
+
+
+def decode_depth(tokens: Tensor, config: DepthConfig) -> Tensor:
+    """Token grid -> depth pixel map.
+
+    Input  shape: (..., tok_h, tok_w, patch_h * patch_w)
+    Output shape: (..., tok_h * patch_size, tok_w * patch_size, 1)
+    """
+    p = config.patch_size
+    depth = einops.rearrange(tokens, "... h w (p1 p2 c) -> ... (h p1) (w p2) c", p1=p, p2=p)
+    if config.depth_value_scale != 1.0:
+        depth = depth / config.depth_value_scale
+    for stage in reversed(config.depth_normalize_mode):
+        depth = _undo_stage(depth, stage)
+    return depth
+
+
+def _contract(z: Tensor) -> Tensor:
+    """Forward mip-NeRF 360 radial contract (inverse of `_inv_contract`).
+
+    Identity for |z| <= 1; otherwise squashes radius r into (2 - 1/r).
+    """
+    mag = z.norm(dim=-1, keepdim=True)
+    safe = mag.clamp(min=1e-12)
+    scale = torch.where(mag <= 1.0, torch.ones_like(mag), (2.0 - 1.0 / safe) / safe)
+    return z * scale
+
+
+def _apply_stage(depth: Tensor, stage: str) -> Tensor:
+    """Forward of `_undo_stage`. Only `contract` and `unit_mean` are active.
+
+    `unit_mean` divides by the per-map mean so the depth is scale-normalised
+    the way the model was trained; it is pass-through on decode (the absolute
+    scale is not recoverable), so encode/decode round-trips up to that scale.
+    """
+    if stage == "unit_mean":
+        return depth / depth.mean().clamp(min=1e-12)
+    if stage != "contract":  # "normal"/"mean"/"none" are pass-through.
+        return depth
+    if depth.shape[-1] == 1:
+        return _contract(depth)
+    return _contract(depth.unsqueeze(-1)).squeeze(-1)
+
+
+def encode_depth(depth_map: Tensor, config: DepthConfig) -> Tensor:
+    """Depth pixel map -> token grid (inverse of `decode_depth`).
+
+    Input  shape: (..., H, W, 1)
+    Output shape: (..., H/patch_size, W/patch_size, patch_size**2)
+
+    Used for ``mode="d2i"`` to turn a depth map into the depth-stream tokens
+    the model conditions on. Because `unit_mean` is scale-normalising, only the
+    relative depth structure matters — the input need not be metric.
+    """
+    p = config.patch_size
+    depth = depth_map
+    for stage in config.depth_normalize_mode:
+        depth = _apply_stage(depth, stage)
+    if config.depth_value_scale != 1.0:
+        depth = depth * config.depth_value_scale
+    return einops.rearrange(depth, "... (h p1) (w p2) c -> ... h w (p1 p2 c)", p1=p, p2=p)