SpaceFlow demo

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+wheels/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,201 @@
+                                 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 [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,13 +1,27 @@
 ---
-title: Spaceflow
-emoji: 🏢
-colorFrom: blue
-colorTo: gray
+title: SpaceFlow
+emoji: 🧊
+colorFrom: green
+colorTo: yellow
 sdk: gradio
 sdk_version: 6.30.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Locally controllable 3D generation from primitives
+python_version: "3.10"
+startup_duration_timeout: 1h
+license: apache-2.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SpaceFlow: Locally Controllable 3D Generation
+
+Gradio demo of [SpaceFlow](https://github.com/SpaceFlow3D/spaceflow) ([project page](https://spaceflow3d.github.io/), [paper](https://arxiv.org/abs/2610.12399)).
+It runs the authors' `run_local_tau.py` text pipeline: SpaceControl sparse-structure generation with
+local (high/low) geometry control from superquadric primitives, PartField part features, and
+self-similarity guided TRELLIS appearance flow with per-primitive local texture prompts.
+
+Models: `microsoft/TRELLIS-text-xlarge`, `microsoft/TRELLIS-image-large` (VAE / decoders),
+`openai/clip-vit-large-patch14`, `mikaelaangel/partfield-ckpt`.
+
+Example primitive bundles come from the SpaceFlow repository (Apache-2.0).
+SpaceFlow is Apache-2.0; **PartField is restricted to non-commercial research and educational use**
+(see `THIRD_PARTY_NOTICES.md`).

THIRD_PARTY_NOTICES.md

@@ -0,0 +1,27 @@
+# Attribution and third-party notices
+
+The root `LICENSE` contains SpaceFlow's Apache-2.0 license. It does not replace
+licenses or notices attached to upstream source, checkpoints, or dependencies.
+
+| Component | Source and notice |
+| --- | --- |
+| GuideFlow3D | [Upstream project](https://github.com/GradientSpaces/GuideFlow3D); SpaceFlow builds on its generation and similarity-guidance pipeline. |
+| TRELLIS | [Upstream](https://github.com/microsoft/TRELLIS); retained [MIT license](third_party/TRELLIS/LICENSE). |
+| PartField | [Upstream](https://github.com/nv-tlabs/PartField); retained [NVIDIA license](third_party/PartField/LICENSE). |
+| DINOv2 | [Upstream](https://github.com/facebookresearch/dinov2); [Apache-2.0 license](https://github.com/facebookresearch/dinov2/blob/7764ea0f912e53c92e82eb78a2a1631e92725fc8/LICENSE). Image staging records its source commit and checkpoint hash. |
+| rembg / U2Net | [rembg upstream](https://github.com/danielgatis/rembg); background-removal dependency and downloaded U2Net model. Their own upstream notices apply. |
+| Native extensions | Sources/revisions are recorded in `requirements/native-revisions.json`; installed dependencies retain their own notices. |
+| React, Three.js, and editor dependencies | Versions and package integrity records are in `sq_ui/app/package-lock.json`; dependencies retain their own licenses. |
+
+## PartField use limitation
+
+PartField's license limits use to **non-commercial research and educational
+purposes** (with the exception specified for NVIDIA and its affiliates).
+Keep its license and notices with any redistribution of that component, and
+review the [original terms](third_party/PartField/LICENSE) for the complete rules.
+The same license constraints apply when using the PartField checkpoint.
+
+Model revision records and checkpoint hashes identify the downloaded artifacts;
+they do not grant additional rights. Download models from their recorded upstream
+repositories and retain the associated notices. Generated research media and
+experiment outputs are separate artifacts from the thin source release.

app.py

@@ -0,0 +1,457 @@
+import os
+import sys
+import glob
+import importlib.util
+import subprocess
+
+APP_DIR = os.path.dirname(os.path.abspath(__file__))
+
+# TRELLIS runtime configuration (matches the reference setup: spconv + flash-attn).
+os.environ.setdefault("SPCONV_ALGO", "native")
+os.environ.setdefault("SPARSE_BACKEND", "spconv")
+os.environ.setdefault("ATTN_BACKEND", "flash_attn")
+os.environ.setdefault("SPARSE_ATTN_BACKEND", "flash_attn")
+os.environ["SPACEFLOW_KEEP_ON_GPU"] = "1"  # all weights stay resident on the ZeroGPU device
+
+# TRELLIS needs the mip-splatting fork of diff-gaussian-rasterization (kernel_size / subpixel_offset API).
+# A prebuilt sm_120 wheel ships with the Space; install it before torch is imported.
+if importlib.util.find_spec("diff_gaussian_rasterization") is None:
+    wheel = glob.glob(os.path.join(APP_DIR, "wheels", "diff_gaussian_rasterization-*.whl"))[0]
+    subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-deps", wheel])
+
+import spaces  # noqa: E402  (must come before torch)
+
+import copy  # noqa: E402
+import json  # noqa: E402
+import random  # noqa: E402
+import shutil  # noqa: E402
+import tempfile  # noqa: E402
+import time  # noqa: E402
+
+import gradio as gr  # noqa: E402
+import numpy as np  # noqa: E402
+import torch  # noqa: E402
+import trimesh  # noqa: E402
+from huggingface_hub import hf_hub_download  # noqa: E402
+from omegaconf import OmegaConf  # noqa: E402
+
+sys.path.insert(0, APP_DIR)
+os.chdir(APP_DIR)
+os.makedirs("debug", exist_ok=True)
+
+import run_local_tau as rlt  # noqa: E402
+from third_party.TRELLIS.trellis.pipelines import TrellisTextTo3DPipeline  # noqa: E402
+from third_party.PartField.partfield.model_trainer_pvcnn_only_demo import Model as PartFieldModel  # noqa: E402
+from third_party.PartField.partfield.dataloader import Demo_Dataset  # noqa: E402
+from lib.opt import self_similarity  # noqa: E402
+from lib.util import pointcloud  # noqa: E402
+from lib.util.checkpoint_compat import partfield_checkpoint_scope  # noqa: E402
+
+CFG = OmegaConf.load("config/default.yaml")
+LOW_CONTROL_BBOX_MARGIN = 0.07  # sq_ui default (DEFAULT_LOW_CONTROL_BBOX_MARGIN)
+
+# ---------------------------------------------------------------------------
+# Models (loaded once at module scope, placed on the GPU)
+# ---------------------------------------------------------------------------
+# TRELLIS' text conditioner encodes the empty prompt at init; defer that forward
+# pass to the GPU worker instead of running it at import time.
+_orig_encode_text = TrellisTextTo3DPipeline.encode_text
+TrellisTextTo3DPipeline.encode_text = lambda self, text: None
+PIPE = TrellisTextTo3DPipeline.from_pretrained("config/trellis_pipeline")
+TrellisTextTo3DPipeline.encode_text = _orig_encode_text
+PIPE.cuda()
+PIPE.text_cond_model["model"].to("cuda")
+
+PF_CFG = OmegaConf.load("third_party/PartField/config.yaml")
+PARTFIELD = PartFieldModel(PF_CFG, None)
+_pf_ckpt = hf_hub_download(
+    "mikaelaangel/partfield-ckpt", "model_objaverse.ckpt",
+    revision="90b9b1e08b6a12fdcb6ee26b4854a26235e1765f",
+)
+with partfield_checkpoint_scope():
+    _pf_state = torch.load(_pf_ckpt, map_location="cpu", weights_only=True)["state_dict"]
+_missing, _unexpected = PARTFIELD.load_state_dict(_pf_state, strict=False)
+print(f"[PartField] missing keys: {len(_missing)}, unexpected keys: {len(_unexpected)}")
+del _pf_state
+PARTFIELD.eval().to("cuda")
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+def _seed_all(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def _ensure_null_cond():
+    if PIPE.text_cond_model.get("null_cond") is None:
+        PIPE.text_cond_model["null_cond"] = PIPE.encode_text([""])
+
+
+def _read_metadata(npz_path):
+    with np.load(npz_path, allow_pickle=False) as d:
+        if "spaceflow_metadata.json" in d.files:
+            try:
+                return json.loads(bytes(d["spaceflow_metadata.json"]).decode("utf-8"))
+            except Exception:
+                return {}
+    return {}
+
+
+def _load_arrays(npz_path):
+    with np.load(npz_path, allow_pickle=False) as d:
+        arrays = {k: np.asarray(d[k]) for k in d.files if k != "spaceflow_metadata.json"}
+    n = arrays["scales"].shape[0]
+    arrays.setdefault("control_levels", np.ones((n,)))
+    return arrays
+
+
+def _subset(arrays, idx):
+    out = {}
+    for k in ["scales", "shapes", "translations", "rotations", "control_levels", "tapering", "bending"]:
+        if k in arrays:
+            out[k] = arrays[k][idx]
+    return out
+
+
+def _low_control_bbox(arrays, low_idx, margin=LOW_CONTROL_BBOX_MARGIN):
+    """Port of sq_ui buildLowControlBoundingBoxPrimitive (padded AABB of the low-control primitives)."""
+    verts = []
+    for i in low_idx:
+        tap = arrays["tapering"][i] if "tapering" in arrays else None
+        bend = arrays["bending"][i] if "bending" in arrays else None
+        v, _ = rlt.add_superquadric_compact_rot_mat(
+            arrays["scales"][i], arrays["shapes"][i], arrays["translations"][i], arrays["rotations"][i],
+            tap, bend, resolution=32)
+        verts.append(v)
+    v = np.concatenate(verts, 0)
+    mn, mx = v.min(0), v.max(0)
+    pad = np.linalg.norm(mx - mn) * margin
+    mn, mx = mn - pad, mx + pad
+    center = (mn + mx) / 2
+    half = np.maximum((mx - mn) / 2, 1e-4)
+    rot = np.array([[1.0, 0.0, 0.0], [0.0, 6.123234e-17, -1.0], [0.0, 1.0, 6.123234e-17]])
+    return {
+        "scales": np.array([[half[0], half[2], half[1]]]),
+        "shapes": np.array([[0.05, 0.05]]),
+        "translations": center[None],
+        "rotations": rot[None],
+        "control_levels": np.array([0.0]),
+    }
+
+
+def _parse_local_prompts(text, n_sq):
+    """'6: pink painted wood' lines (1-based primitive index) -> per-primitive list."""
+    prompts = [""] * n_sq
+    for line in (text or "").replace(";", "\n").splitlines():
+        line = line.strip()
+        if not line or ":" not in line:
+            continue
+        idx, prompt = line.split(":", 1)
+        idx = idx.strip().lower().replace("sq", "").strip()
+        if not idx.isdigit():
+            continue
+        k = int(idx) - 1
+        if 0 <= k < n_sq:
+            prompts[k] = prompt.strip()
+    return prompts
+
+
+def _blender_normalized_ply(glb_path, ply_path):
+    """Python port of TRELLIS' Blender `--save_mesh` export used by SpaceFlow.
+
+    Blender imports glTF (Y-up) into its Z-up frame, scales the scene so the
+    largest bbox side is 1, re-centres the bbox at the origin and writes a PLY.
+    """
+    mesh = trimesh.load(glb_path, force="mesh", process=False)
+    v = np.asarray(mesh.vertices, dtype=np.float64)
+    v = np.stack([v[:, 0], -v[:, 2], v[:, 1]], axis=1)
+    mn, mx = v.min(0), v.max(0)
+    v = v * (1.0 / (mx - mn).max())
+    mn, mx = v.min(0), v.max(0)
+    v = v - (mn + mx) / 2
+    trimesh.Trimesh(vertices=v, faces=np.asarray(mesh.faces), process=False).export(ply_path)
+
+
+def _partfield_planes(mesh_path, out_dir):
+    """Equivalent of run_local_tau.predict_part without the Lightning Trainer."""
+    _seed_all(int(PF_CFG.seed))
+    item = Demo_Dataset(mesh_path)[0]
+    batch = {"pc": item["pc"][None].to("cuda"), "uid": [item["uid"]]}
+    with torch.autocast("cuda", dtype=torch.float16):
+        part_planes, uid = PARTFIELD.predict_step(batch, 0)
+    np.save(os.path.join(out_dir, f"part_feat_{uid}_batch_part_plane.npy"), part_planes)
+
+
+def preview_primitives(primitives_file):
+    """Render the uploaded primitives (orange = high control, white = low control) and read saved prompts."""
+    if primitives_file is None:
+        return None, gr.update(), gr.update(), gr.update(), ""
+    path = primitives_file if isinstance(primitives_file, str) else primitives_file.name
+    out = tempfile.mkdtemp()
+    glb = os.path.join(out, "input_superquadrics.glb")
+    rlt.export_colored_superquadrics_glb(path, glb)
+    # Show in a Y-up viewer the same way the result is oriented.
+    meta = _read_metadata(path)
+    arrays = _load_arrays(path)
+    n = arrays["scales"].shape[0]
+    names = meta.get("primitiveNames") or [f"SQ {i + 1}" for i in range(n)]
+    levels = ["high" if float(c) >= 0.5 else "low" for c in arrays["control_levels"]]
+    local = meta.get("localTextureTexts") or [""] * n
+    rows = "\n".join(
+        f"| {i + 1} | {names[i] if i < len(names) else ''} | {levels[i]} | {local[i] if i < len(local) else ''} |"
+        for i in range(n))
+    table = "| # | name | control | local texture |\n|---|---|---|---|\n" + rows
+    shape = meta.get("textPrompt")
+    tex = meta.get("globalTextureText")
+    local_txt = "\n".join(f"{i + 1}: {t}" for i, t in enumerate(local) if str(t or "").strip())
+    return (
+        glb,
+        gr.update(value=shape) if shape else gr.update(),
+        gr.update(value=tex) if tex else gr.update(),
+        gr.update(value=local_txt) if meta else gr.update(),
+        table,
+    )
+
+
+def _duration(primitives_file, shape_prompt, texture_prompt, local_prompts="", low_tau=3.0, high_tau=10.0,
+              polyak_tau=0.18, repaint_steps=10, texture_steps=300, seed=1, attn_boost=2.5, *args, **kwargs):
+    return int(min(300, 45 + 0.5 * int(texture_steps)))
+
+
+@spaces.GPU(duration=_duration)
+def generate(
+    primitives_file,
+    shape_prompt: str,
+    texture_prompt: str,
+    local_prompts: str = "",
+    low_tau: float = 3.0,
+    high_tau: float = 10.0,
+    polyak_tau: float = 0.18,
+    repaint_steps: int = 10,
+    texture_steps: int = 300,
+    seed: int = 1,
+    attn_boost: float = 2.5,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Generate a textured 3D asset from superquadric primitives with SpaceFlow.
+
+    Args:
+        primitives_file: SpaceFlow primitives `.npz` (scales, shapes, translations, rotations, control_levels).
+        shape_prompt: Text describing the object (e.g. "toy elephant").
+        texture_prompt: Global appearance text (e.g. "blue painted wood").
+        local_prompts: Optional per-primitive appearance overrides, one per line as "<index>: <text>" (1-based).
+        low_tau: Structure-control strength for low-control primitives (lower = more freedom).
+        high_tau: Structure-control strength for high-control primitives.
+        polyak_tau: Polyak blending factor between the low- and high-control trajectories.
+        repaint_steps: Repaint resampling steps used to blend the two control levels.
+        texture_steps: Number of self-similarity guided texture flow steps.
+        seed: Seed for sparse-structure sampling.
+        attn_boost: Region-boosted self-attention factor for local texture routing (1.0 = off).
+
+    Returns:
+        Path to the textured GLB and a short run log.
+    """
+    if primitives_file is None:
+        raise gr.Error("Upload a primitives .npz file (or pick an example).")
+    shape_prompt = (shape_prompt or "").strip()
+    if not shape_prompt:
+        raise gr.Error("Enter a shape prompt.")
+    texture_prompt = (texture_prompt or "").strip() or shape_prompt
+    t0 = time.perf_counter()
+    path = primitives_file if isinstance(primitives_file, str) else primitives_file.name
+
+    out_dir = tempfile.mkdtemp(prefix="spaceflow_")
+    arrays = _load_arrays(path)
+    n_sq = arrays["scales"].shape[0]
+    all_path = os.path.join(out_dir, "all.npz")
+    np.savez(all_path, **arrays)
+    high_idx = [i for i in range(n_sq) if float(arrays["control_levels"][i]) >= 0.5]
+    low_idx = [i for i in range(n_sq) if float(arrays["control_levels"][i]) < 0.5]
+    local_mode = bool(high_idx) and bool(low_idx)
+    if local_mode and float(high_tau) < float(low_tau):
+        raise gr.Error("High-control tau must be >= low-control tau.")
+
+    cfg = copy.deepcopy(CFG)
+    cfg.sim_guidance.steps = max(2, int(texture_steps))
+    cfg.sim_guidance.mixed_self_attn_boost = max(1.0, float(attn_boost))
+
+    _seed_all(0)
+    _ensure_null_cond()
+    log = []
+
+    # --- spatial-control meshes from the superquadrics -------------------------------------
+    sc_mesh = os.path.join(out_dir, "spatial_control_mesh.ply")
+    aabb, center, scale = rlt.load_superquadrics(all_path, sc_mesh)
+    sampler = {
+        "steps": rlt.STEPS_SHAPE_GEN,
+        "cfg_strength": rlt.CFG_SHAPE_GEN,
+        "t0_idx_value": float(low_tau),
+        "spatial_control_mesh_path": sc_mesh,
+        "high_control_spatial_control_mesh_path": None,
+        "low_control_superquadric_mask_path": None,
+        "t0_idx_value_high_control": None,
+        "polyak_update_tau": float(polyak_tau) if local_mode else 0.0,
+        "local_tau_mode": "guidance",
+        "n_repaint_steps": int(repaint_steps),
+    }
+    if local_mode:
+        hi_npz = os.path.join(out_dir, "high_control.npz")
+        lo_npz = os.path.join(out_dir, "low_control_bbox.npz")
+        np.savez(hi_npz, **_subset(arrays, high_idx))
+        np.savez(lo_npz, **_low_control_bbox(arrays, low_idx))
+        hi_mesh = os.path.join(out_dir, "high_control_spatial_control_mesh.ply")
+        lo_mesh = os.path.join(out_dir, "low_control_superquadric_mask.ply")
+        rlt.load_superquadrics(hi_npz, hi_mesh, aabb=aabb, center=center, scale=scale)
+        rlt.load_superquadrics(lo_npz, lo_mesh, aabb=aabb, center=center, scale=scale)
+        sampler.update({
+            "high_control_spatial_control_mesh_path": hi_mesh,
+            "low_control_superquadric_mask_path": lo_mesh,
+            "t0_idx_value_high_control": float(high_tau),
+            "local_tau_mode": "low_control_mask",
+        })
+        log.append(f"Local control: {len(high_idx)} high / {len(low_idx)} low primitives "
+                   f"(tau {low_tau} / {high_tau}, polyak {polyak_tau}).")
+    else:
+        log.append(f"Global control: all {n_sq} primitives at tau {low_tau}.")
+
+    # --- 1. SpaceControl sparse structure ---------------------------------------------------
+    coords = PIPE.gen_structure_v2(shape_prompt, seed=int(seed), vis_output_dir=None,
+                                   sparse_structure_sampler_params=sampler)
+    voxels = coords.detach().cpu().numpy()[:, 1:]
+    log.append(f"Structure: {voxels.shape[0]} voxels ({time.perf_counter() - t0:.1f}s).")
+    rlt.sparse_voxels_to_glb(voxels, grid_size=64, output_filename=os.path.join(out_dir, "sample.glb"))
+    struct_mesh = trimesh.load(os.path.join(out_dir, "sample.glb"), force="mesh")
+    struct_mesh = pointcloud.convert_mesh_yup_to_zup(struct_mesh)  # --convert_yup_to_zup
+    struct_zup = os.path.join(out_dir, "struct_mesh_zup.glb")
+    struct_mesh.export(struct_zup)
+
+    render_dir = os.path.join(out_dir, "struct_renders")
+    os.makedirs(render_dir, exist_ok=True)
+    struct_ply = os.path.join(render_dir, "mesh.ply")
+    _blender_normalized_ply(struct_zup, struct_ply)
+    voxel_dir = os.path.join(out_dir, "voxels")
+    os.makedirs(voxel_dir, exist_ok=True)
+    pointcloud.voxelize_mesh(struct_ply, save_path=os.path.join(voxel_dir, "struct_voxels.ply"))
+
+    # --- 2. PartField features on the generated structure -----------------------------------
+    pf_dir = os.path.join(out_dir, "partfield")
+    os.makedirs(pf_dir, exist_ok=True)
+    _partfield_planes(struct_ply, pf_dir)
+
+    # --- 3. self-similarity guided appearance flow with local texture routing ---------------
+    local_list = _parse_local_prompts(local_prompts, n_sq)
+    active = [i + 1 for i, p in enumerate(local_list) if p]
+    local_list = local_list if active else None
+    individual = rlt.build_individual_sq_meshes_normalized(all_path) if local_list else None
+    if active:
+        log.append("Local texture overrides on primitives " + ", ".join(map(str, active)) + ".")
+    self_similarity.optimize_self_similarity(
+        cfg, texture_prompt, "text", out_dir,
+        local_prompts=local_list,
+        local_prompt_type="text" if local_list else None,
+        individual_sq_meshes=individual,
+        generation_pipeline=PIPE,
+        decode_texture=True,
+    )
+    out_glb = os.path.join(out_dir, "out_sim.glb")
+    for wp in glob.glob(os.path.join(out_dir, "*warning*.json")):
+        try:
+            w = json.load(open(wp))
+            items = w if isinstance(w, list) else w.get("warnings", [w])
+            for it in items:
+                if isinstance(it, dict) and it.get("message"):
+                    log.append("Warning: " + it["message"])
+        except Exception:
+            pass
+    log.append(f"Done in {time.perf_counter() - t0:.1f}s.")
+    print("[SpaceFlow] " + " ".join(log))
+    return out_glb, "\n".join(f"- {line}" for line in log)
+
+
+# ---------------------------------------------------------------------------
+# UI
+# ---------------------------------------------------------------------------
+EXAMPLES = [
+    ["examples/toy_elephant.npz", "toy elephant", "blue painted wood", "6: pink painted wood\n9: pink painted wood"],
+    ["examples/sailboat.npz", "Sailboat", "white", None],
+    ["examples/chair.npz", "A chair", "light wood", None],
+    ["examples/hot_air_balloon.npz", "a vintage hot air balloon with a wicker basket",
+     "warm vintage canvas, leather trim, and soft sunlit shading", None],
+    ["examples/toy_rocket.npz", "a toy rocket", "white metal", None],
+    ["examples/windmill.npz", "a rustic windmill with four wooden blades",
+     "weathered white plaster walls with dark wooden beams", None],
+]
+# Fill local prompts for the examples from the authors' saved metadata.
+for _ex in EXAMPLES:
+    if _ex[3] is None:
+        _m = _read_metadata(_ex[0])
+        _ex[3] = "\n".join(f"{i + 1}: {t}" for i, t in enumerate(_m.get("localTextureTexts", [])) if str(t).strip())
+
+CSS = """
+#col-container { max-width: 1200px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(title="SpaceFlow") as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            "# SpaceFlow: Locally Controllable 3D Generation\n"
+            "Training-free 3D generation from **superquadric primitives**: each primitive is marked as "
+            "*high* (stick to it) or *low* (only a loose hint) geometry control, and can carry its own "
+            "**local texture prompt**. Built on TRELLIS-text-xlarge + PartField.  \n"
+            "[Project page](https://spaceflow3d.github.io/) · [Paper](https://arxiv.org/abs/2610.12399) · "
+            "[Code](https://github.com/SpaceFlow3D/spaceflow) · Author primitives in the "
+            "[SpaceFlow editor](https://github.com/SpaceFlow3D/spaceflow#interactive-editor) and upload its `all.npz` here."
+        )
+        with gr.Row():
+            with gr.Column(scale=1):
+                primitives = gr.File(label="Primitives (.npz from the SpaceFlow editor)", file_types=[".npz"],
+                                     type="filepath")
+                input_view = gr.Model3D(label="Input primitives (orange = high control, white = low control)",
+                                        height=300, interactive=False)
+                shape_prompt = gr.Textbox(label="Shape prompt", placeholder="toy elephant")
+                texture_prompt = gr.Textbox(label="Global texture prompt", placeholder="blue painted wood")
+                local_prompts = gr.Textbox(
+                    label="Local texture prompts (optional)", lines=3,
+                    placeholder="One per line, '<primitive #>: <texture>', e.g.\n6: pink painted wood")
+                run_btn = gr.Button("Generate", variant="primary")
+                with gr.Accordion("Primitive list", open=False):
+                    prim_table = gr.Markdown()
+                with gr.Accordion("Advanced settings", open=False):
+                    with gr.Row():
+                        low_tau = gr.Slider(0, 12, value=3.0, step=0.5, label="Low-control tau")
+                        high_tau = gr.Slider(0, 12, value=10.0, step=0.5, label="High-control tau")
+                    with gr.Row():
+                        polyak_tau = gr.Slider(0, 1, value=0.18, step=0.01, label="Polyak blend")
+                        repaint_steps = gr.Slider(0, 20, value=10, step=1, label="Repaint steps")
+                    with gr.Row():
+                        texture_steps = gr.Slider(50, 300, value=300, step=10, label="Texture flow steps")
+                        attn_boost = gr.Slider(1.0, 4.0, value=2.5, step=0.1, label="Local attention boost")
+                    seed = gr.Number(value=1, precision=0, label="Structure seed")
+            with gr.Column(scale=1):
+                output = gr.Model3D(label="Generated textured mesh (GLB)", height=520)
+                info = gr.Markdown()
+
+        inputs = [primitives, shape_prompt, texture_prompt, local_prompts, low_tau, high_tau, polyak_tau,
+                  repaint_steps, texture_steps, seed, attn_boost]
+        primitives.change(preview_primitives, inputs=[primitives],
+                          outputs=[input_view, shape_prompt, texture_prompt, local_prompts, prim_table],
+                          api_visibility="private")
+        run_btn.click(generate, inputs=inputs, outputs=[output, info], api_name="generate")
+
+        gr.Examples(
+            examples=EXAMPLES,
+            inputs=[primitives, shape_prompt, texture_prompt, local_prompts],
+            outputs=[output, info],
+            fn=generate,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+if __name__ == "__main__":
+    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)

config/default.yaml

@@ -0,0 +1,41 @@
+# render
+num_views: 0
+
+# trellis
+trellis_img_model_name: "JeffreyXiang/TRELLIS-image-large"
+# Local pipeline config that mixes TRELLIS-image-large (encoder/decoder/SLAT
+# decoders) with TRELLIS-text-xlarge (sparse-structure & SLAT flow models).
+# Necessary for spatial-control: microsoft/TRELLIS-text-large does not ship a
+# sparse_structure_encoder, but encode_spatial_control() requires one.
+trellis_text_model_name: "config/trellis_pipeline"
+flow_model_in_channels: 8
+
+img_model:
+  rescale_t: 3.0
+  cfg_strength: 5.0
+  cfg_interval: [0.5, 1.0]
+
+text_model:
+  rescale_t: 3.0
+  cfg_strength: 7.5
+  cfg_interval: [0.5, 9.5]
+
+# self-similarity guidance
+sim_guidance:
+  num_part_clusters: 30
+  steps: 300
+  learning_rate: 5e-4
+  loss_weight: 1.0
+  local_condition_dilation: 0
+  local_condition_min_cells: 192
+  local_condition_max_dilation: 1
+  # Multiply self-attention weights by this factor before the softmax when both
+  # voxels share a routed condition. 1.0 disables and keeps stock self-attention.
+  mixed_self_attn_boost: 2.5
+  # "legacy" preserves the original PartField/geometric router.  "active_margin"
+  # lets active local SQs claim nearby occupied cells when the generated structure
+  # is slightly shifted relative to the input superquadric.
+  local_routing_mode: legacy
+  local_routing_margin: 2.0
+
+log_every: 100

config/trellis_pipeline/pipeline.json

@@ -0,0 +1,67 @@
+{
+    "name": "TrellisTextTo3DPipeline",
+    "args": {
+        "models": {
+            "sparse_structure_decoder": "microsoft/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16",
+            "sparse_structure_encoder": "microsoft/TRELLIS-image-large/ckpts/ss_enc_conv3d_16l8_fp16",
+            "sparse_structure_flow_model": "microsoft/TRELLIS-text-xlarge/ckpts/ss_flow_txt_dit_XL_16l8_fp16",
+            "slat_decoder_gs": "microsoft/TRELLIS-image-large/ckpts/slat_dec_gs_swin8_B_64l8gs32_fp16",
+            "slat_decoder_mesh": "microsoft/TRELLIS-image-large/ckpts/slat_dec_mesh_swin8_B_64l8m256c_fp16",
+            "slat_flow_model_text": "microsoft/TRELLIS-text-xlarge/ckpts/slat_flow_txt_dit_XL_64l8p2_fp16"
+        },
+        "sparse_structure_sampler": {
+            "name": "FlowEulerGuidanceIntervalSampler",
+            "args": {
+                "sigma_min": 1e-05
+            },
+            "params": {
+                "steps": 25,
+                "cfg_strength": 7.5,
+                "cfg_interval": [
+                    0.5,
+                    0.95
+                ],
+                "rescale_t": 3.0
+            }
+        },
+        "slat_sampler": {
+            "name": "FlowEulerGuidanceIntervalSampler",
+            "args": {
+                "sigma_min": 1e-05
+            },
+            "params": {
+                "steps": 25,
+                "cfg_strength": 7.5,
+                "cfg_interval": [
+                    0.5,
+                    0.95
+                ],
+                "rescale_t": 3.0
+            }
+        },
+        "slat_normalization": {
+            "mean": [
+                -2.1687545776367188,
+                -0.004347046371549368,
+                -0.13352349400520325,
+                -0.08418072760105133,
+                -0.5271206498146057,
+                0.7238689064979553,
+                -1.1414450407028198,
+                1.2039363384246826
+            ],
+            "std": [
+                2.377650737762451,
+                2.386378288269043,
+                2.124418020248413,
+                2.1748552322387695,
+                2.663944721221924,
+                2.371192216873169,
+                2.6217446327209473,
+                2.684523105621338
+            ]
+        },
+        "text_cond_model": "openai/clip-vit-large-patch14",
+        "image_cond_model": "dinov2_vitl14_reg"
+    }
+}

examples/chair.npz

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

examples/hot_air_balloon.npz

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

examples/sailboat.npz

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

examples/toy_elephant.npz

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

examples/toy_rocket.npz

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

examples/windmill.npz

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

lib/opt/self_similarity.py

@@ -0,0 +1,1072 @@
+import copy
+import json
+import os.path as osp
+from PIL import Image
+import numpy as np
+import torch
+import torch.nn.functional as F
+import utils3d
+import logging
+import open3d_pycg as o3d
+import time
+import os
+_KEEP_ON_GPU = os.environ.get('SPACEFLOW_KEEP_ON_GPU', '') == '1'
+
+import third_party.TRELLIS.trellis.modules.sparse as sp
+from third_party.TRELLIS.trellis.pipelines import TrellisImageTo3DPipeline, TrellisTextTo3DPipeline
+from lib.util import generation, partfield
+from lib.util.pipeline_compat import can_reuse_appearance_pipeline
+
+# Global logger
+log = logging.getLogger(__name__)
+
+
+def _select_slat_flow_model(generation_pipeline, app_type):
+    """Return the SLAT flow model key/model for old and split text/image pipeline configs."""
+    if app_type == 'image':
+        candidates = ['slat_flow_model_image', 'slat_flow_model']
+    else:
+        candidates = ['slat_flow_model_text', 'slat_flow_model']
+
+    for key in candidates:
+        if key in generation_pipeline.models:
+            log.info(f"Using SLAT flow model: {key}")
+            return key, generation_pipeline.models[key]
+
+    available = ', '.join(sorted(generation_pipeline.models.keys()))
+    raise KeyError(
+        f"No compatible SLAT flow model found for {app_type} conditioning. "
+        f"Tried {candidates}; available models: {available}"
+    )
+
+
+def _text_conditioner_to(generation_pipeline, device):
+    text_cond_model = getattr(generation_pipeline, 'text_cond_model', None)
+    if not text_cond_model:
+        return
+    if 'model' in text_cond_model:
+        text_cond_model['model'].to(device)
+    if 'null_cond' in text_cond_model:
+        text_cond_model['null_cond'] = text_cond_model['null_cond'].to(device)
+
+
+def _preprocess_condition_image(generation_pipeline, path):
+    image = Image.open(path).convert('RGB')
+    if hasattr(generation_pipeline, 'preprocess_image'):
+        return generation_pipeline.preprocess_image(image)
+    return image
+
+
+def _voxelize_mesh(mesh):
+    """Voxelize an Open3D mesh into a (1,1,64,64,64) binary float32 tensor on CPU."""
+    mesh = copy.deepcopy(mesh)
+    vertices = np.clip(np.asarray(mesh.vertices), -0.5 + 1e-6, 0.5 - 1e-6)
+    mesh.vertices = o3d.utility.Vector3dVector(vertices)
+    voxel_grid = o3d.geometry.VoxelGrid.create_from_triangle_mesh_within_bounds(
+        mesh, voxel_size=1/64, min_bound=(-0.5, -0.5, -0.5), max_bound=(0.5, 0.5, 0.5))
+    voxel_indices = np.array([v.grid_index for v in voxel_grid.get_voxels()])
+    coords_dense = torch.zeros(1, 1, 64, 64, 64, dtype=torch.float32)
+    if len(voxel_indices) > 0:
+        coords_dense[0, 0, voxel_indices[:, 0], voxel_indices[:, 1], voxel_indices[:, 2]] = 1.0
+    return coords_dense
+
+
+def _dense_cell_groups(struct_coords, device='cuda'):
+    """Return occupied 32^3 cells and the 64^3 source-voxel group for each point."""
+    dense_coords = (struct_coords[:, 1:].to(device).long() // 2).clamp(0, 31)
+    linear = (
+        dense_coords[:, 0] * 32 * 32
+        + dense_coords[:, 1] * 32
+        + dense_coords[:, 2]
+    )
+    unique_linear, inverse = torch.unique(linear, sorted=True, return_inverse=True)
+    unique_coords = torch.stack(
+        [
+            unique_linear // (32 * 32),
+            (unique_linear // 32) % 32,
+            unique_linear % 32,
+        ],
+        dim=1,
+    )
+    return unique_coords, inverse
+
+
+def _surface_distance_matrix(generated_coords, sq_coords_list):
+    """Distance from each generated 64^3 voxel to each SQ surface."""
+    n_sq = len(sq_coords_list)
+    distances = torch.full(
+        (n_sq, generated_coords.shape[0]),
+        float('inf'),
+        dtype=generated_coords.dtype,
+        device=generated_coords.device,
+    )
+    for idx, sq_coords in enumerate(sq_coords_list):
+        if sq_coords.shape[0] > 0:
+            distances[idx] = torch.cdist(sq_coords, generated_coords).min(0).values
+    return distances
+
+
+def _aggregate_dense_min_distances(point_distances, inverse, n_cells):
+    """Pool point-to-SQ distances into occupied 32^3 cells without last-write bias."""
+    n_sq = point_distances.shape[0]
+    cell_distances = torch.full(
+        (n_sq, n_cells),
+        float('inf'),
+        dtype=point_distances.dtype,
+        device=point_distances.device,
+    )
+    for sq_idx in range(n_sq):
+        cell_distances[sq_idx].scatter_reduce_(
+            0,
+            inverse,
+            point_distances[sq_idx],
+            reduce='amin',
+            include_self=True,
+        )
+    return cell_distances
+
+
+def _aggregate_dense_votes(point_labels, inverse, n_cells, n_labels):
+    """Pool per-point labels into 32^3 cells by majority vote."""
+    votes = torch.zeros(
+        (n_cells, n_labels),
+        dtype=torch.int32,
+        device=point_labels.device,
+    )
+    votes.index_put_(
+        (inverse, point_labels.long()),
+        torch.ones_like(point_labels, dtype=torch.int32),
+        accumulate=True,
+    )
+    return votes.argmax(dim=1)
+
+
+def compute_active_margin_indices(
+    struct_coords,
+    individual_sq_meshes,
+    active_indices,
+    device='cuda',
+    margin=2.0,
+):
+    """Route active local conditions to nearby occupied cells with a distance margin.
+
+    The legacy router assigns a cell only to the single nearest SQ.  This is brittle
+    when the generated structure shifts a thin part by a few 64^3 voxels.  Here an
+    active SQ may claim a cell when it is within ``margin`` voxels of the best SQ;
+    competing active SQs are resolved by their distance to the cell.  The routing is
+    still restricted to occupied generated cells, so the margin cannot create texture
+    in empty space.
+    """
+    coords_dense_indices = torch.zeros(
+        1, 1, 32, 32, 32, dtype=torch.int32, device=device
+    )
+    n_sq = len(individual_sq_meshes)
+    active_indices = [idx for idx in active_indices if 0 <= idx < n_sq]
+    if n_sq == 0 or not active_indices or struct_coords.shape[0] == 0:
+        return coords_dense_indices
+
+    generated_coords = struct_coords[:, 1:].float().to(device)
+    sq_coords_list = []
+    for mesh in individual_sq_meshes:
+        sq_vox = _voxelize_mesh(mesh)
+        sq_coords = torch.argwhere(sq_vox)[:, 2:].float().to(device)
+        sq_coords_list.append(sq_coords)
+
+    dense_coords, inverse = _dense_cell_groups(struct_coords, device=device)
+    point_distances = _surface_distance_matrix(generated_coords, sq_coords_list)
+    cell_distances = _aggregate_dense_min_distances(
+        point_distances,
+        inverse,
+        dense_coords.shape[0],
+    )
+
+    active = torch.as_tensor(active_indices, dtype=torch.long, device=device)
+    active_distances = cell_distances.index_select(0, active)
+    best_active_distances, best_active_pos = active_distances.min(dim=0)
+    best_all_distances = cell_distances.min(dim=0).values
+    selected = best_active_distances <= (best_all_distances + float(margin))
+    if selected.any():
+        selected_coords = dense_coords[selected]
+        selected_sq = active[best_active_pos[selected]] + 1
+        coords_dense_indices[
+            0,
+            0,
+            selected_coords[:, 0],
+            selected_coords[:, 1],
+            selected_coords[:, 2],
+        ] = selected_sq.to(torch.int32)
+    return coords_dense_indices
+
+
+def compute_surface_coverage_indices(
+    struct_coords,
+    individual_sq_meshes,
+    active_indices,
+    device='cuda',
+    min_cells=64,
+    max_cells=512,
+    candidate_margin=3.0,
+):
+    """Route local conditions from exact SQ surface coverage.
+
+    PartField and nearest-SQ routing can lose thin or small parts after the
+    structure is pooled into the 32^3 SLAT grid.  This router first reserves
+    cells touched by the exact 64^3 SQ surface voxelization, then fills only a
+    compact nearest-cell budget for each active SQ.  Small details therefore
+    receive a guaranteed local condition without a broad, unconstrained
+    dilation into neighboring parts.
+
+    Values are SQ indices (1-indexed; zero means global/unassigned), matching
+    ``compute_coords_dense_indices`` and the later condition remapping.
+    """
+    coords_dense_indices = torch.zeros(
+        1, 1, 32, 32, 32, dtype=torch.int32, device=device
+    )
+    n_sq = len(individual_sq_meshes)
+    active_indices = [idx for idx in active_indices if 0 <= idx < n_sq]
+    min_cells = max(1, int(min_cells))
+    max_cells = max(min_cells, int(max_cells))
+    if n_sq == 0 or not active_indices or struct_coords.shape[0] == 0:
+        return coords_dense_indices, {
+            'routing_source': 'surface_coverage',
+            'exact_surface_cells': {},
+            'reserved_cells': {},
+        }
+
+    generated_coords = struct_coords[:, 1:].float().to(device)
+    dense_coords, inverse = _dense_cell_groups(struct_coords, device=device)
+    n_cells = int(dense_coords.shape[0])
+    if n_cells == 0:
+        return coords_dense_indices, {
+            'routing_source': 'surface_coverage',
+            'exact_surface_cells': {},
+            'reserved_cells': {},
+        }
+
+    sq_coords_list = []
+    for mesh in individual_sq_meshes:
+        sq_vox = _voxelize_mesh(mesh)
+        sq_coords_list.append(torch.argwhere(sq_vox)[:, 2:].float().to(device))
+    point_distances = _surface_distance_matrix(generated_coords, sq_coords_list)
+    cell_distances = _aggregate_dense_min_distances(
+        point_distances,
+        inverse,
+        n_cells,
+    )
+
+    dense_linear = (
+        dense_coords[:, 0] * 32 * 32
+        + dense_coords[:, 1] * 32
+        + dense_coords[:, 2]
+    ).long()
+    assigned = torch.zeros(n_cells, dtype=torch.int32, device=device)
+    exact_counts = {}
+    targets = {}
+    candidate_orders = {}
+
+    # Build a compact candidate list for each active SQ. Exact surface cells
+    # are always preferred; the nearest cells fill the minimum budget only
+    # when the detail is too small or shifted relative to the structure grid.
+    for sq_idx in active_indices:
+        surface = sq_coords_list[sq_idx]
+        if surface.shape[0] > 0:
+            surface_dense = (surface.long() // 2).clamp(0, 31)
+            surface_linear = (
+                surface_dense[:, 0] * 32 * 32
+                + surface_dense[:, 1] * 32
+                + surface_dense[:, 2]
+            ).long()
+            positions = torch.searchsorted(dense_linear, surface_linear)
+            valid = positions < n_cells
+            safe_positions = positions.clamp(max=n_cells - 1)
+            valid = valid & (dense_linear[safe_positions] == surface_linear)
+            exact = torch.unique(safe_positions[valid])
+        else:
+            exact = torch.empty(0, dtype=torch.long, device=device)
+
+        exact_counts[sq_idx] = int(exact.numel())
+        target = min(max_cells, max(min_cells, int(exact.numel())))
+        targets[sq_idx] = target
+        order = torch.argsort(cell_distances[sq_idx], stable=True)
+        if candidate_margin > 0 and order.numel() > 0:
+            best = cell_distances[sq_idx, order[0]]
+            near = order[cell_distances[sq_idx, order] <= best + float(candidate_margin)]
+            if near.numel() >= min_cells:
+                order = near
+        if exact.numel() > 0:
+            order = torch.cat([exact, order])
+            order = torch.unique(order, sorted=False)
+        candidate_orders[sq_idx] = order[: min(max_cells, int(order.numel()))]
+
+    # Reserve compact regions for the smallest details first. This prevents a
+    # large neighboring part from consuming all cells needed by a tiny part.
+    priority = sorted(active_indices, key=lambda idx: (exact_counts[idx], idx))
+    for sq_idx in priority:
+        order = candidate_orders[sq_idx]
+        available = order[assigned[order] == 0]
+        selected = available[: targets[sq_idx]]
+        if selected.numel() < targets[sq_idx]:
+            # Overlapping SQs may have no fully disjoint candidate budget. A
+            # detail still gets its nearest cells, with later assignments
+            # intentionally unable to overwrite this reservation.
+            selected = order[: targets[sq_idx]]
+        assigned[selected] = int(sq_idx + 1)
+
+    coords_dense_indices[
+        0,
+        0,
+        dense_coords[:, 0],
+        dense_coords[:, 1],
+        dense_coords[:, 2],
+    ] = assigned
+    reserved_counts = {
+        sq_idx: int((assigned == sq_idx + 1).sum().item())
+        for sq_idx in active_indices
+    }
+    return coords_dense_indices, {
+        'routing_source': 'surface_coverage',
+        'exact_surface_cells': exact_counts,
+        'reserved_cells': reserved_counts,
+        'target_cells': targets,
+    }
+
+
+def compute_coords_dense_indices(struct_coords, individual_sq_meshes, device='cuda', vox_cluster_labels=None):
+    """Map each 64^3 voxel to its nearest superquadric (1-indexed; 0=unassigned).
+
+    When vox_cluster_labels is provided (PartField cluster label per voxel), assigns
+    voxels via cluster centroids: each cluster centroid is matched to the nearest SQ,
+    and all voxels in that cluster inherit the SQ assignment. This is more accurate
+    than per-voxel geometric nearest neighbor because PartField clusters capture
+    semantic part structure.
+
+    struct_coords: (M, 4) int tensor [batch, x, y, z] in 64x64x64 space.
+    vox_cluster_labels: (M,) int tensor of PartField cluster indices per voxel.
+    Returns a (1, 1, 32, 32, 32) int32 tensor.
+    """
+    coords_dense_indices = torch.zeros(1, 1, 32, 32, 32, dtype=torch.int32, device=device)
+    n_sq = len(individual_sq_meshes)
+    if n_sq == 0:
+        return coords_dense_indices
+
+    generated_coords = struct_coords[:, 1:].float().to(device)  # (M, 3)
+
+    # Precompute SQ surface voxel coords once
+    sq_coords_list = []
+    for mesh in individual_sq_meshes:
+        sq_vox = _voxelize_mesh(mesh)  # (1,1,64,64,64)
+        sq_c = torch.argwhere(sq_vox)[:, 2:].float().to(device)  # (K,3)
+        sq_coords_list.append(sq_c)
+
+    if vox_cluster_labels is not None:
+        # PartField-based assignment: cluster centroid → nearest SQ → voxels
+        labels = vox_cluster_labels.to(device)
+        num_clusters = int(labels.max().item()) + 1
+        cluster_to_sq = torch.zeros(num_clusters, dtype=torch.long, device=device)
+        for c in range(num_clusters):
+            mask = labels == c
+            if not mask.any():
+                continue
+            cluster_voxels = generated_coords[mask]
+            mean_pos = cluster_voxels.mean(0, keepdim=True)
+            centroid = cluster_voxels[torch.cdist(cluster_voxels, mean_pos).argmin()].unsqueeze(0)  # (1, 3) actual voxel
+            min_dists = [
+                torch.cdist(sq_c, centroid).min().item() if sq_c.shape[0] > 0 else float('inf')
+                for sq_c in sq_coords_list
+            ]
+            cluster_to_sq[c] = int(np.argmin(min_dists))
+        point_labels = cluster_to_sq[labels]
+        dense_coords, inverse = _dense_cell_groups(struct_coords, device=device)
+        dense_labels = _aggregate_dense_votes(
+            point_labels,
+            inverse,
+            dense_coords.shape[0],
+            n_sq,
+        )
+    else:
+        # Fallback: per-voxel geometric nearest neighbor
+        min_distances = _surface_distance_matrix(generated_coords, sq_coords_list)
+        dense_coords, inverse = _dense_cell_groups(struct_coords, device=device)
+        cell_distances = _aggregate_dense_min_distances(
+            min_distances,
+            inverse,
+            dense_coords.shape[0],
+        )
+        dense_labels = cell_distances.argmin(dim=0)
+
+    coords_dense_indices[
+        0,
+        0,
+        dense_coords[:, 0],
+        dense_coords[:, 1],
+        dense_coords[:, 2],
+    ] = (dense_labels + 1).to(torch.int32)
+
+    return coords_dense_indices
+
+def _dense_sq_counts(coords_dense_indices, n_sq):
+    flat = coords_dense_indices.reshape(-1)
+    return {
+        sq_idx: int((flat == sq_idx + 1).sum().item())
+        for sq_idx in range(n_sq)
+    }
+
+def _dense_condition_counts(coords_dense_indices, n_conditions):
+    flat = coords_dense_indices.reshape(-1)
+    return {
+        cond_idx: int((flat == cond_idx).sum().item())
+        for cond_idx in range(1, n_conditions)
+    }
+
+def _active_route_coverage(sq_counts, active_indices):
+    counts = [sq_counts.get(idx, 0) for idx in active_indices]
+    return sum(count > 0 for count in counts), sum(counts)
+
+def _format_nonzero_counts(counts):
+    nonzero = {idx: count for idx, count in counts.items() if count > 0}
+    return nonzero if nonzero else {}
+
+def _write_routing_warning(output_dir, warning):
+    warning_path = osp.join(output_dir, "routing_warnings.json")
+    try:
+        if osp.isfile(warning_path):
+            with open(warning_path, "r", encoding="utf-8") as fh:
+                payload = json.load(fh)
+        else:
+            payload = {"warnings": []}
+        warnings = payload.get("warnings")
+        if not isinstance(warnings, list):
+            warnings = []
+            payload["warnings"] = warnings
+        warnings.append(warning)
+        with open(warning_path, "w", encoding="utf-8") as fh:
+            json.dump(payload, fh, indent=2)
+    except Exception as exc:  # noqa: BLE001
+        log.warning("Could not write routing warning metadata: %s", exc)
+
+def _choose_local_routing(
+    struct_coords,
+    individual_sq_meshes,
+    struct_labels,
+    active_indices,
+    device='cuda',
+    routing_mode='legacy',
+    routing_margin=2.0,
+    surface_min_cells=64,
+    surface_max_cells=512,
+    surface_candidate_margin=3.0,
+):
+    if routing_mode in {'surface_coverage', 'surface_topk'} and active_indices:
+        surface_routed, surface_stats = compute_surface_coverage_indices(
+            struct_coords,
+            individual_sq_meshes,
+            active_indices,
+            device=device,
+            min_cells=surface_min_cells,
+            max_cells=surface_max_cells,
+            candidate_margin=surface_candidate_margin,
+        )
+        surface_counts = _dense_sq_counts(surface_routed, len(individual_sq_meshes))
+        log.info(
+            "Surface-coverage local routing (min=%d, max=%d, margin=%.2f) SQ counts: %s; exact cells: %s",
+            int(surface_min_cells),
+            int(surface_max_cells),
+            float(surface_candidate_margin),
+            _format_nonzero_counts(surface_counts),
+            surface_stats.get('exact_surface_cells', {}),
+        )
+        return surface_routed, 'surface_coverage', surface_counts
+
+    if routing_mode == 'active_margin' and active_indices:
+        margin_routed = compute_active_margin_indices(
+            struct_coords,
+            individual_sq_meshes,
+            active_indices,
+            device=device,
+            margin=routing_margin,
+        )
+        margin_counts = _dense_sq_counts(margin_routed, len(individual_sq_meshes))
+        log.info(
+            "Active-margin local routing (margin=%.2f) SQ counts: %s",
+            routing_margin,
+            _format_nonzero_counts(margin_counts),
+        )
+        return margin_routed, 'active_margin', margin_counts
+
+    clustered = compute_coords_dense_indices(
+        struct_coords, individual_sq_meshes, device, vox_cluster_labels=struct_labels)
+    n_sq = len(individual_sq_meshes)
+    clustered_counts = _dense_sq_counts(clustered, n_sq)
+    log.info(f"PartField local routing SQ counts: {_format_nonzero_counts(clustered_counts)}")
+
+    if not active_indices:
+        return clustered, "partfield", clustered_counts
+
+    clustered_coverage = _active_route_coverage(clustered_counts, active_indices)
+    if clustered_coverage[0] == len(active_indices):
+        return clustered, "partfield", clustered_counts
+
+    geometric = compute_coords_dense_indices(
+        struct_coords, individual_sq_meshes, device, vox_cluster_labels=None)
+    geometric_counts = _dense_sq_counts(geometric, n_sq)
+    geometric_coverage = _active_route_coverage(geometric_counts, active_indices)
+    log.info(f"Geometric local routing SQ counts: {_format_nonzero_counts(geometric_counts)}")
+
+    if geometric_coverage > clustered_coverage:
+        missing = [idx for idx in active_indices if clustered_counts.get(idx, 0) == 0]
+        log.warning(
+            "PartField local routing missed active SQ(s) %s; using geometric routing "
+            "for local texture conditions instead.",
+            missing,
+        )
+        return geometric, "geometric", geometric_counts
+
+    missing = [idx for idx in active_indices if clustered_counts.get(idx, 0) == 0]
+    if missing:
+        log.warning(
+            "PartField local routing missed active SQ(s) %s, and geometric routing "
+            "did not improve coverage; keeping PartField routing.",
+            missing,
+        )
+    return clustered, "partfield", clustered_counts
+
+def _dilate_condition_indices(coords_dense_indices, n_conditions, radius=1):
+    if radius <= 0 or n_conditions <= 1:
+        return coords_dense_indices
+
+    dilated_indices = coords_dense_indices.clone()
+    kernel_size = 2 * radius + 1
+    for cond_idx in range(1, n_conditions):
+        mask = (coords_dense_indices == cond_idx).float()
+        dilated = F.max_pool3d(mask, kernel_size=kernel_size, stride=1, padding=radius) > 0
+        dilated_indices[(dilated_indices == 0) & dilated] = cond_idx
+    return dilated_indices
+
+def _nonnegative_int(value, default=0):
+    if value is None:
+        return default
+    try:
+        return max(0, int(value))
+    except (TypeError, ValueError):
+        return default
+
+def _select_compact_candidate_cells(mask, candidates, needed):
+    if needed <= 0:
+        return candidates[:0]
+    candidate_coords = candidates.nonzero(as_tuple=False)
+    if candidate_coords.shape[0] <= needed:
+        return candidate_coords
+
+    kernel = torch.ones((1, 1, 3, 3, 3), dtype=mask.dtype, device=mask.device)
+    support = F.conv3d(mask, kernel, stride=1, padding=1)[candidates]
+    source_coords = mask.nonzero(as_tuple=False)[:, 2:].float()
+    centroid = source_coords.mean(dim=0, keepdim=True)
+    spatial = candidate_coords[:, 2:].float()
+    dist2 = ((spatial - centroid) ** 2).sum(dim=1)
+    linear = (
+        candidate_coords[:, 2] * mask.shape[-2] * mask.shape[-1]
+        + candidate_coords[:, 3] * mask.shape[-1]
+        + candidate_coords[:, 4]
+    ).float()
+    priority = support.float() * 1_000_000.0 - dist2 - linear / 1_000_000.0
+    selected = torch.topk(priority, k=needed, largest=True, sorted=False).indices
+    return candidate_coords[selected]
+
+def _expand_small_condition_indices(coords_dense_indices, n_conditions, min_cells=0, max_dilation=0):
+    """Grow under-covered local conditions into global cells.
+
+    The local routing grid is 32^3, so small SQs such as wheels can collapse to
+    very few routed cells. This keeps local prompts active by expanding only
+    conditions that are below min_cells, never overwriting another local prompt.
+    """
+    min_cells = _nonnegative_int(min_cells, 0)
+    max_dilation = _nonnegative_int(max_dilation, 0)
+    if min_cells <= 0 or max_dilation <= 0 or n_conditions <= 1:
+        return coords_dense_indices, {}
+
+    expanded_indices = coords_dense_indices.clone()
+    stats = {}
+    for cond_idx in range(1, n_conditions):
+        initial_count = int((expanded_indices == cond_idx).sum().item())
+        stats[cond_idx] = {
+            "initial": initial_count,
+            "final": initial_count,
+            "dilation": 0,
+            "added": 0,
+        }
+        if initial_count <= 0 or initial_count >= min_cells:
+            continue
+
+        count = initial_count
+        for radius in range(1, max_dilation + 1):
+            mask = (expanded_indices == cond_idx).float()
+            grown = F.max_pool3d(mask, kernel_size=3, stride=1, padding=1) > 0
+            candidates = grown & (expanded_indices == 0)
+            if not candidates.any():
+                break
+            needed = min_cells - count
+            selected = _select_compact_candidate_cells(mask, candidates, needed)
+            expanded_indices[
+                selected[:, 0],
+                selected[:, 1],
+                selected[:, 2],
+                selected[:, 3],
+                selected[:, 4],
+            ] = cond_idx
+            count = int((expanded_indices == cond_idx).sum().item())
+            stats[cond_idx]["final"] = count
+            stats[cond_idx]["dilation"] = radius
+            stats[cond_idx]["added"] = count - initial_count
+            if count >= min_cells:
+                break
+
+    return expanded_indices, stats
+
+def attn_cosine_sim(x, eps=1e-08):
+    x = x[0]  # TEMP: getting rid of redundant dimension, TBF
+    norm1 = x.norm(dim=2, keepdim=True)
+    factor = torch.clamp(norm1 @ norm1.permute(0, 2, 1), min=eps)
+    sim_matrix = (x @ x.permute(0, 2, 1)) / factor
+    return sim_matrix
+
+def _auto_contrastive_chunk_size(num_voxels):
+    if num_voxels <= 16_000:
+        return 1024
+    if num_voxels <= 28_000:
+        return 512
+    return 256
+
+
+def chunked_contrastive_loss(feats, labels, eps=1e-08, chunk_size=None):
+    """Supervised contrastive loss without materializing the full N×N similarity matrix.
+    Peak memory is O(chunk_size × N) instead of O(N²).
+    """
+    # feats: (1, 1, N, C) -> (N, C)
+    x = feats[0, 0]
+    x_norm = x / x.norm(dim=1, keepdim=True).clamp(min=eps)  # (N, C), L2-normalized
+    labels = labels.view(-1)  # (N,)
+    N = x_norm.shape[0]
+    chunk_size = chunk_size or _auto_contrastive_chunk_size(N)
+
+    numerator = torch.zeros(N, device=x.device)
+    denominator = torch.zeros(N, device=x.device)
+    valid = torch.zeros(N, dtype=torch.bool, device=x.device)
+
+    for i in range(0, N, chunk_size):
+        i_end = min(i + chunk_size, N)
+        sim_row = x_norm[i:i_end] @ x_norm.T  # (cs, N)
+        exp_sim = torch.exp(sim_row)
+
+        labels_i = labels[i:i_end].unsqueeze(1)  # (cs, 1)
+        same_label = labels_i == labels.unsqueeze(0)  # (cs, N)
+
+        # Exclude self-similarity on diagonal
+        diag_idx = torch.arange(i_end - i, device=x.device)
+        self_mask = torch.zeros_like(same_label)
+        self_mask[diag_idx, i + diag_idx] = True
+        positive_mask = same_label & ~self_mask
+
+        denominator[i:i_end] = exp_sim.masked_fill(self_mask, 0.0).sum(dim=1)
+        valid[i:i_end] = positive_mask.any(dim=1)
+        numerator[i:i_end] = exp_sim.masked_fill(~positive_mask, 0.0).sum(dim=1)
+
+    loss = -torch.log(numerator / (denominator + 1e-8))
+    return loss[valid].mean()
+
+def optimize_self_similarity(cfg, app, app_type, output_dir,
+                             local_prompts=None, local_prompt_type=None,
+                             individual_sq_meshes=None,
+                             generation_pipeline=None,
+                             decode_texture=True):
+    overall_start = time.perf_counter()
+    log.info("Starting self-similarity optimization...")
+
+    pipeline_start = time.perf_counter()
+    if not can_reuse_appearance_pipeline(generation_pipeline, app_type):
+        if generation_pipeline is not None:
+            log.info("Loading the image appearance pipeline; the structure/text pipeline has no image preprocessing interface.")
+        if app_type == 'image':
+            generation_pipeline = TrellisImageTo3DPipeline.from_pretrained(cfg.trellis_img_model_name)
+        else:
+            generation_pipeline = TrellisTextTo3DPipeline.from_pretrained(cfg.trellis_text_model_name)
+    else:
+        log.info("Reusing preloaded TRELLIS pipeline for self-similarity guidance.")
+    generation_pipeline.cuda()
+    _text_conditioner_to(generation_pipeline, 'cuda')
+    log.info("Prepared TRELLIS pipeline for self-similarity in %.2fs", time.perf_counter() - pipeline_start)
+
+    if app_type == 'image':
+        app = _preprocess_condition_image(generation_pipeline, osp.join(output_dir, 'app_image.png'))
+    
+    # Load Structure Data
+    structure_start = time.perf_counter()
+    struct_coords = utils3d.io.read_ply(osp.join(output_dir, 'voxels', 'struct_voxels.ply'))[0]
+    struct_coords = torch.from_numpy(struct_coords).float().cuda()
+    struct_coords = ((struct_coords + 0.5) * 64).long()
+    
+    zeros = torch.zeros((struct_coords.size(0), 1), dtype=struct_coords.dtype, device=struct_coords.device)
+    struct_coords = torch.cat([zeros, struct_coords], dim=1)
+    
+    # Load partfield planes (extracted from Blender's struct_renders/mesh.ply — same mesh as voxels)
+    path = osp.join(output_dir, "partfield", "part_feat_mesh_batch_part_plane.npy")
+    struct_part_planes = torch.from_numpy(np.load(path, allow_pickle=True)).cuda()
+
+    struct_labels = partfield.cluster_geoms(struct_coords, struct_part_planes, num_clusters=cfg.sim_guidance.num_part_clusters)
+
+    # Optimization Starts...
+    struct_labels = torch.from_numpy(struct_labels.flatten()).cuda()
+    struct_feats_params = torch.nn.Parameter(torch.randn((struct_coords.shape[0], cfg.flow_model_in_channels)), requires_grad=True)
+    log.info(
+        "Loaded structure voxels, PartField planes, and labels in %.2fs (%d voxels)",
+        time.perf_counter() - structure_start,
+        struct_coords.shape[0],
+    )
+    
+    param_list = [struct_feats_params]
+    optimizer = torch.optim.AdamW(param_list, lr=cfg.sim_guidance.learning_rate)
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda x: 1)
+    
+    best_loss = float('inf')
+    feats = None
+    
+    cond_start = time.perf_counter()
+    if app_type == "image":
+        cond = generation_pipeline.get_cond([app]) if hasattr(generation_pipeline, 'get_cond') else generation_pipeline.get_cond_image([app])
+    else:
+        cond = generation_pipeline.get_cond_text([app])
+    log.info("Encoded global %s guidance condition in %.2fs", app_type, time.perf_counter() - cond_start)
+
+    # Build per-SQ local conditioning if local prompts were provided.
+    cond_list = None
+    coords_dense_indices = None
+    if local_prompts and individual_sq_meshes:
+        local_start = time.perf_counter()
+        global_emb = cond['cond']  # (1, seq_len, dim)
+        n_sq = len(local_prompts)
+        active_indices = [
+            idx
+            for idx, prompt in enumerate(local_prompts)
+            if str(prompt or "").strip()
+        ]
+        log.info(
+            f"Building local {local_prompt_type} conditioning for "
+            f"{len(active_indices)} / {n_sq} superquadrics; active SQ indices: {active_indices}"
+        )
+        if local_prompt_type == 'text':
+            local_embs = []
+            for sq_idx, prompt in enumerate(local_prompts):
+                prompt = str(prompt or "").strip()
+                if prompt:
+                    log.info(
+                        "Local text condition for SQ %d: %r",
+                        sq_idx,
+                        prompt,
+                    )
+                local_embs.append(
+                    generation_pipeline.encode_text([prompt])
+                    if prompt
+                    else global_emb
+                )
+        else:  # 'image'
+            local_embs = []
+            for sq_idx, path in enumerate(local_prompts):
+                path = str(path or '').strip()
+                if path:
+                    log.info(
+                        "Local image condition for SQ %d: %s",
+                        sq_idx,
+                        path,
+                    )
+                    local_image = _preprocess_condition_image(generation_pipeline, path)
+                    local_embs.append(generation_pipeline.encode_image([local_image]))
+                else:
+                    local_embs.append(global_emb)
+        # Only include SQs that have a real (non-global-fallback) local condition.
+        conditioned_mask = [emb is not global_emb for emb in local_embs]
+        conditioned_embs = [emb for emb, is_cond in zip(local_embs, conditioned_mask) if is_cond]
+        cond_list = [global_emb] + conditioned_embs
+
+        # Build remap: SQ 1-index → new condition index (0=global fallback, 1..n_conditioned).
+        remap = torch.zeros(n_sq + 1, dtype=torch.long, device='cuda')
+        new_idx = 1
+        condition_to_sq = {}
+        for sq_i, is_cond in enumerate(conditioned_mask):
+            remap[sq_i + 1] = new_idx if is_cond else 0
+            if is_cond:
+                condition_to_sq[new_idx] = sq_i
+                new_idx += 1
+        sq_cond_map = remap[1:].tolist()  # condition index per SQ (0=global)
+        log.info(f"SQ condition mapping (0=global): {sq_cond_map}")
+        log.info(f"Built cond_list with {len(cond_list)} entries "
+                 f"(1 global + {len(conditioned_embs)} real local out of {n_sq} SQs)")
+
+        local_routing_mode = str(
+            getattr(cfg.sim_guidance, 'local_routing_mode', 'legacy')
+        ).strip().lower()
+        local_routing_margin = float(
+            getattr(cfg.sim_guidance, 'local_routing_margin', 2.0)
+        )
+        surface_min_cells = _nonnegative_int(
+            getattr(cfg.sim_guidance, 'surface_min_cells', 64), 64
+        )
+        surface_max_cells = _nonnegative_int(
+            getattr(cfg.sim_guidance, 'surface_max_cells', 512), 512
+        )
+        surface_candidate_margin = float(
+            getattr(cfg.sim_guidance, 'surface_candidate_margin', 3.0)
+        )
+        coords_dense_indices, routing_source, sq_route_counts = _choose_local_routing(
+            struct_coords,
+            individual_sq_meshes,
+            struct_labels,
+            active_indices,
+            device='cuda',
+            routing_mode=local_routing_mode,
+            routing_margin=local_routing_margin,
+            surface_min_cells=surface_min_cells,
+            surface_max_cells=surface_max_cells,
+            surface_candidate_margin=surface_candidate_margin,
+        )
+        coords_dense_indices = remap[coords_dense_indices.long()]
+        condition_counts_before_dilation = _dense_condition_counts(coords_dense_indices, len(cond_list))
+        local_condition_dilation = _nonnegative_int(
+            getattr(cfg.sim_guidance, 'local_condition_dilation', 0), 0)
+        coords_dense_indices = _dilate_condition_indices(
+            coords_dense_indices, len(cond_list), radius=local_condition_dilation)
+        condition_counts_after_dilation = _dense_condition_counts(coords_dense_indices, len(cond_list))
+        local_condition_min_cells = _nonnegative_int(
+            getattr(cfg.sim_guidance, 'local_condition_min_cells', 0), 0)
+        local_condition_max_dilation = _nonnegative_int(
+            getattr(cfg.sim_guidance, 'local_condition_max_dilation', 0), 0)
+        coords_dense_indices, adaptive_expansion_stats = _expand_small_condition_indices(
+            coords_dense_indices,
+            len(cond_list),
+            min_cells=local_condition_min_cells,
+            max_dilation=local_condition_max_dilation,
+        )
+        condition_counts = _dense_condition_counts(coords_dense_indices, len(cond_list))
+        missing_condition_sqs = [
+            condition_to_sq.get(cond_idx, 'unknown')
+            for cond_idx in range(1, len(cond_list))
+            if condition_counts.get(cond_idx, 0) == 0
+        ]
+        if missing_condition_sqs:
+            missing_sq_numbers = [
+                sq_idx + 1 if isinstance(sq_idx, int) else sq_idx
+                for sq_idx in missing_condition_sqs
+            ]
+            sq_label = ", ".join(str(sq_num) for sq_num in missing_sq_numbers)
+            warning_message = (
+                f"Local {local_prompt_type} prompt(s) for SQs {sq_label} "
+                "received zero routed cells; those local overrides had no effect."
+            )
+            log.warning(
+                "Local %s condition(s) for SQ(s) %s have zero routed cells; "
+                "those local overrides will have no effect. Check condition_routing.mp4 "
+                "or enlarge/adjust the selected superquadric.",
+                local_prompt_type,
+                missing_condition_sqs,
+            )
+            _write_routing_warning(
+                output_dir,
+                {
+                    "kind": "local_routing_zero_cells",
+                    "severity": "warning",
+                    "message": warning_message,
+                    "local_prompt_type": local_prompt_type,
+                    "sq_indices": [
+                        sq_idx for sq_idx in missing_condition_sqs
+                        if isinstance(sq_idx, int)
+                    ],
+                    "sq_numbers": missing_sq_numbers,
+                    "routing_source": routing_source,
+                    "active_sq_counts": {
+                        str(idx + 1): int(sq_route_counts.get(idx, 0))
+                        for idx in active_indices
+                    },
+                    "condition_counts": {
+                        str(cond_idx): int(count)
+                        for cond_idx, count in condition_counts.items()
+                    },
+                },
+            )
+        active_sq_counts = {idx: sq_route_counts.get(idx, 0) for idx in active_indices}
+        log.info(
+            f"Local routing source: {routing_source}; active SQ counts: {active_sq_counts}"
+        )
+        if adaptive_expansion_stats:
+            active_expansions = {
+                f"cond_{cond_idx}_sq_{condition_to_sq.get(cond_idx, 'unknown')}": values
+                for cond_idx, values in adaptive_expansion_stats.items()
+                if values["dilation"] > 0 or values["initial"] < local_condition_min_cells
+            }
+            if active_expansions:
+                log.info(
+                    "Adaptive local condition expansion: min_cells=%d, max_dilation=%d, stats=%s",
+                    local_condition_min_cells,
+                    local_condition_max_dilation,
+                    active_expansions,
+                )
+        log.info(
+            f"coords_dense_indices: shape={coords_dense_indices.shape}, "
+            f"condition_counts_before_dilation={condition_counts_before_dilation}, "
+            f"condition_counts_after_dilation={condition_counts_after_dilation}, "
+            f"dilation_radius={local_condition_dilation}, "
+            f"adaptive_min_cells={local_condition_min_cells}, "
+            f"adaptive_max_dilation={local_condition_max_dilation}, "
+            f"condition_counts={condition_counts}, non-zero={int((coords_dense_indices > 0).sum())}"
+        )
+
+        log.info("Skipping condition routing visualization; routing indices are still used for local conditioning.")
+        torch.cuda.empty_cache()
+        log.info("Prepared local condition routing in %.2fs", time.perf_counter() - local_start)
+
+        # condition routing visualization for routing
+        log.info("Visualizing condition routing on structure mesh...")
+        try:
+            if _KEEP_ON_GPU:
+                raise RuntimeError("routing visualization disabled in the Space")
+            import trimesh
+            from lib.util.visualization import visualize_and_save, map_voxel_labels_to_vertices
+            _mesh_path = osp.join(output_dir, 'struct_renders', 'mesh.ply')
+            if not osp.isfile(_mesh_path):
+                _mesh_path = osp.join(output_dir, 'spatial_control_mesh.ply')
+            if osp.isfile(_mesh_path):
+                _sv_norm = ((struct_coords[:, 1:].float() + 0.5) / 64 - 0.5).cpu().numpy()
+                _sq_labels = coords_dense_indices[0, 0,
+                    struct_coords[:, 1] // 2,
+                    struct_coords[:, 2] // 2,
+                    struct_coords[:, 3] // 2].cpu().numpy()
+                _mesh_vis = trimesh.load(_mesh_path, force='mesh')
+                _vtx_labels = map_voxel_labels_to_vertices(_mesh_vis.vertices, _sv_norm, _sq_labels)
+                visualize_and_save(_mesh_vis, _vtx_labels, output_dir, output_name='condition_routing.mp4')
+                del _sv_norm, _sq_labels, _mesh_vis, _vtx_labels
+                torch.cuda.empty_cache()
+            else:
+                log.info("No structure mesh found for routing visualization; continuing.")
+        except Exception as _vis_err:
+            log.warning("Could not render condition routing video: %s; continuing texture optimization.", _vis_err)
+
+    else:
+        log.info("No local SQ texture overrides provided; all voxels use the global condition.")
+
+    if not _KEEP_ON_GPU:
+        _text_conditioner_to(generation_pipeline, 'cpu')
+
+    flow_model_key, flow_model = _select_slat_flow_model(generation_pipeline, app_type)
+
+    # get_cond() is done — only the selected SLAT flow model is needed for the loop. Offload everything else.
+    for k, m in generation_pipeline.models.items():
+        if k != flow_model_key and not _KEEP_ON_GPU:
+            m.cpu()
+    torch.cuda.empty_cache()
+    
+    if app_type == 'image':
+        sampler_params = {
+            "cfg_strength": cfg.img_model.cfg_strength,
+            "cfg_interval": cfg.img_model.cfg_interval,
+        }
+        rescale_t = cfg.img_model.rescale_t
+    
+    else:
+        sampler_params = {
+            "cfg_strength": cfg.text_model.cfg_strength,
+            "cfg_interval": cfg.text_model.cfg_interval,
+        }
+        rescale_t = cfg.text_model.rescale_t
+
+    t_seq = np.linspace(1, 0, cfg.sim_guidance.steps + 1)
+    t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq)
+    t_pairs = list((t_seq[i], t_seq[i + 1]) for i in range(cfg.sim_guidance.steps))
+    
+    std = torch.tensor(generation_pipeline.slat_normalization['std'])[None].cuda()
+    mean = torch.tensor(generation_pipeline.slat_normalization['mean'])[None].cuda()
+
+    loop_start = time.perf_counter()
+    log.info(f"Beginning self-similarity guidance + flow sampling loop for {len(t_pairs)} steps...")
+
+    self_attn_region = None
+    mixed_self_attn_boost = 1.0
+    if coords_dense_indices is not None:
+        mixed_self_attn_boost = float(
+            getattr(cfg.sim_guidance, 'mixed_self_attn_boost', 1.0) or 1.0)
+        if mixed_self_attn_boost > 1.0:
+            self_attn_region = coords_dense_indices
+            log.info(
+                "Region-boosted self-attention ON: in-region weights x%.2f before softmax "
+                "over %d condition regions, all %d steps.",
+                mixed_self_attn_boost, int(coords_dense_indices.max().item()) + 1,
+                len(t_pairs),
+            )
+        else:
+            log.info("Region-boosted self-attention off (mixed_self_attn_boost=%.2f); "
+                     "self-attention is stock.", mixed_self_attn_boost)
+
+    for iteration, (t, t_prev) in enumerate(t_pairs):
+        optimizer.zero_grad()
+        
+        # Diffusion
+        struct_feats_params_clone = struct_feats_params.clone().cuda()
+        noise = sp.SparseTensor(
+            feats = struct_feats_params_clone,
+            coords = struct_coords.int(),
+        ).cuda()
+        
+        local_kwargs = {}
+        if cond_list is not None:
+            local_kwargs['cond_list'] = cond_list
+            local_kwargs['coords_dense_indices'] = coords_dense_indices
+        if self_attn_region is not None:
+            local_kwargs['self_attn_region'] = self_attn_region
+            local_kwargs['self_attn_region_boost'] = mixed_self_attn_boost
+
+        with torch.no_grad():
+            out = generation_pipeline.slat_sampler.sample_once(flow_model, noise, t, t_prev, **cond, **sampler_params, **local_kwargs)
+            
+        sample = out.pred_x_prev
+        struct_feats_params.data = sample.feats
+
+        # guideflow loss put to 0
+        if cfg.sim_guidance.loss_weight == 0:
+            feats = struct_feats_params.detach() * std + mean
+    
+        # Optimization - Structure Loss
+        if cfg.sim_guidance.loss_weight > 0 and iteration < len(t_pairs) - 1:
+            struct_loss = chunked_contrastive_loss(struct_feats_params[None, None, ...], struct_labels)
+
+            total_loss = cfg.sim_guidance.loss_weight * struct_loss
+            total_loss.backward()
+            optimizer.step()
+            scheduler.step()
+            
+            if (iteration == 0) or (iteration + 1) % cfg.log_every == 0:
+                message = f"Step: {iteration}, Structure Loss: {struct_loss.item():.4f}, Total Loss: {total_loss.item():.4f}"
+                log.info(message)
+                
+            if total_loss < best_loss:
+                best_loss = total_loss.item()
+                feats = struct_feats_params.detach() * std + mean
+    log.info(
+        "Completed self-similarity guidance + flow sampling loop in %.2fs",
+        time.perf_counter() - loop_start,
+    )
+    
+    # Move SLAT decoders back to GPU for decoding.
+    decoder_load_start = time.perf_counter()
+    decoder_keys = ['slat_decoder_mesh']
+    if decode_texture:
+        decoder_keys.append('slat_decoder_gs')
+    for k in decoder_keys:
+        if k in generation_pipeline.models:
+            generation_pipeline.models[k].cuda()
+        elif decode_texture:
+            raise KeyError(f"Texture decode requested but required decoder is missing: {k}")
+    log.info("Moved SLAT decoder(s) to CUDA in %.2fs: %s", time.perf_counter() - decoder_load_start, decoder_keys)
+
+    # Decode SLAT
+    log.info("Decoding output SLAT...")
+    decode_start = time.perf_counter()
+    out_meshpath = osp.join(output_dir,  'out_sim.glb')
+    generation.decode_slat(generation_pipeline, feats, struct_coords, out_meshpath, None, texture=decode_texture)
+    log.info("Decoded output SLAT in %.2fs: %s", time.perf_counter() - decode_start, out_meshpath)
+    log.info("Completed self-similarity optimization in %.2fs", time.perf_counter() - overall_start)

lib/util/checkpoint_compat.py

@@ -0,0 +1,18 @@
+"""Scoped compatibility for the recorded upstream PartField checkpoint."""
+
+from contextlib import contextmanager
+
+
+@contextmanager
+def partfield_checkpoint_scope():
+    """Allow the checkpoint's YACS configuration with PyTorch weights-only loading.
+
+    PartField stores a CfgNode alongside its tensors. Lightning 2.2 delegates
+    loading to torch.load, whose default changed in PyTorch 2.6. Keep that
+    default and allow only this configuration class while PartField loads.
+    """
+    import torch
+    from yacs.config import CfgNode
+
+    with torch.serialization.safe_globals([CfgNode]):
+        yield

lib/util/common.py

@@ -0,0 +1,101 @@
+import json
+import os
+import os.path as osp
+import pickle
+import re
+from pathlib import Path
+from typing import Any
+import torch 
+import random
+import numpy as np
+
+def make_dir(dir_path: str) -> None:
+    """Creates a directory if it does not exist."""
+    if not Path(dir_path).exists():
+        Path(dir_path).mkdir(parents=True, exist_ok=True)
+
+def ensure_dir(path: str) -> None:
+    """
+    Ensures that a directory exists; creates it if it does not.
+    """
+    if not osp.exists(path):
+        os.makedirs(path)
+
+def assert_dir(path: str) -> None:
+    """Asserts that a directory exists."""
+    assert osp.exists(path)
+
+def load_pkl_data(filename: str) -> Any:
+    """Loads data from a pickle file."""
+    with open(filename, 'rb') as handle:
+        data_dict = pickle.load(handle)
+    return data_dict
+
+def write_pkl_data(data_dict: Any, filename: str) -> None:
+    """Writes data to a pickle file."""
+    with open(filename, 'wb') as handle:
+        pickle.dump(data_dict, handle, protocol=pickle.HIGHEST_PROTOCOL)
+
+def load_json(filename: str) -> Any:
+    """Loads data from a JSON file."""
+    file = open(filename)
+    data = json.load(file)
+    file.close()
+    return data
+
+def write_json(data_dict: Any, filename: str) -> None:
+    """Writes data to a JSON file with indentation."""
+    json_obj = json.dumps(data_dict, indent=4)
+ 
+    with open(filename, "w") as outfile:
+        outfile.write(json_obj)
+
+def get_print_format(value: Any) -> str:
+    """Determines the appropriate format string for a given value."""
+    if isinstance(value, int):
+        return 'd'
+    if isinstance(value, str):
+        return 's'
+    if value == 0:
+        return '.3f'
+    if value < 1e-6:
+        return '.3e'
+    if value < 1e-3:
+        return '.6f'
+    return '.6f'
+
+
+def get_format_strings(kv_pairs: list) -> list:
+    """Generates format strings for a list of key-value pairs."""
+    log_strings = []
+    for key, value in kv_pairs:
+        fmt = get_print_format(value)
+        format_string = '{}: {:' + fmt + '}'
+        log_strings.append(format_string.format(key, value))
+    return log_strings
+
+def get_first_index_batch(x: Any) -> Any:
+    """Retrieves the first index from a batch, handling different data types."""
+    if isinstance(x, list):
+        x = x[0]
+    elif isinstance(x, torch.Tensor):
+        x = x.squeeze(0)
+    elif isinstance(x, dict):
+        x = {key: get_first_index_batch(value) for key, value in x.items()}
+    return x
+
+def split_sentence(sentence: str) -> list:
+    """Splits a sentence into individual sentences based on periods."""
+    sentence = re.split(r'[.]', sentence)
+    sentence = [s.strip() for s in sentence]
+    sentence = [s for s in sentence if len(s) > 0]
+    return sentence
+
+def set_random_seed(seed: int) -> None:
+    """Sets the random seed for reproducibility."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False

lib/util/generation.py

@@ -0,0 +1,80 @@
+import os
+import torch
+import numpy as np
+import imageio
+import trimesh
+
+import third_party.TRELLIS.trellis.modules.sparse as sp
+from third_party.TRELLIS.trellis.utils import render_utils, postprocessing_utils
+
+def _white_geometry_mesh(vertices, faces):
+    mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False)
+    mesh.visual = trimesh.visual.TextureVisuals(
+        material=trimesh.visual.material.PBRMaterial(
+            name='geometry_white',
+            baseColorFactor=[1.0, 1.0, 1.0, 1.0],
+            metallicFactor=0.0,
+            roughnessFactor=0.85,
+        )
+    )
+    mesh.vertex_normals
+    mesh.face_normals
+    return mesh
+
+
+def _postprocess_mesh_geometry(mesh, simplify=0.95):
+    vertices = mesh.vertices.cpu().numpy()
+    faces = mesh.faces.cpu().numpy()
+    vertices, faces = postprocessing_utils.postprocess_mesh(
+        vertices,
+        faces,
+        simplify=simplify > 0,
+        simplify_ratio=simplify,
+        fill_holes=True,
+        fill_holes_max_hole_size=0.04,
+        fill_holes_max_hole_nbe=int(250 * np.sqrt(1 - simplify)),
+        fill_holes_resolution=1024,
+        fill_holes_num_views=1000,
+        debug=False,
+        verbose=False,
+    )
+    return vertices, faces
+
+
+def decode_slat(generation_pipeline, feats, coords, out_meshpath, out_gspath, texture=True):
+    # Decode Output SLAT
+    slat = sp.SparseTensor(
+            feats = feats.float(),
+            coords = coords.int(),
+        ).cuda()            
+    formats = ['mesh', 'gaussian'] if texture else ['mesh']
+    with torch.no_grad():
+        outputs = generation_pipeline.decode_slat(slat, formats)
+
+    out_geometry_path = f"{os.path.splitext(out_meshpath)[0]}_geometry.glb"
+
+    if texture:
+        mesh_textured = postprocessing_utils.to_glb(
+                        outputs['gaussian'][0],
+                        outputs['mesh'][0],
+                        # Optional parameters
+                        simplify=0.95,          # Ratio of triangles to remove in the simplification process
+                        texture_size=1024,      # Size of the texture used for the GLB
+                        verbose=False,          # Print logs
+                    )
+        mesh_textured.export(out_meshpath)
+        mesh_geometry = _white_geometry_mesh(
+            mesh_textured.vertices.copy(),
+            mesh_textured.faces.copy(),
+        )
+        mesh_geometry.export(out_geometry_path)
+
+        if out_gspath:
+            video = render_utils.render_video(outputs['gaussian'][0])['color']
+            imageio.mimsave(out_gspath, video, fps=30)
+        return
+
+    vertices, faces = _postprocess_mesh_geometry(outputs['mesh'][0], simplify=0.95)
+    mesh_geometry = _white_geometry_mesh(vertices, faces)
+    mesh_geometry.export(out_meshpath)
+    mesh_geometry.export(out_geometry_path)

lib/util/model_revisions.py

@@ -0,0 +1,16 @@
+"""Recorded source revision for TRELLIS's DINOv2 image conditioner."""
+
+import json
+from pathlib import Path
+import re
+
+
+MODEL_REVISIONS = Path(__file__).resolve().parents[2] / "requirements/model-revisions.json"
+
+
+def dinov2_hub_repository(pins_path: Path = MODEL_REVISIONS) -> str:
+    """Use an immutable upstream commit instead of PyTorch Hub's moving main."""
+    record = json.loads(pins_path.read_text())["dinov2"]
+    if record["repository"] != "facebookresearch/dinov2" or not re.fullmatch(r"[0-9a-f]{40}", record["revision"]):
+        raise ValueError("DINOv2 requires the recorded upstream repository and a full commit SHA.")
+    return f"{record['repository']}:{record['revision']}"

lib/util/partfield.py

@@ -0,0 +1,108 @@
+import numpy as np
+from sklearn.cluster import KMeans
+from sklearn.neighbors import KDTree
+
+import torch
+import torch.nn.functional as F
+
+def get_voxel_partfeats(voxel_coords, part_planes):
+    voxel_coords = ((voxel_coords[:, 1:] + 0.5) / 64 - 0.5).cpu().numpy()
+    bbmin = voxel_coords.min(0)
+    bbmax = voxel_coords.max(0)
+    center = (bbmin + bbmax) * 0.5
+    scale = 2.0 * 0.9 / (bbmax - bbmin).max()
+    voxel_coords = (voxel_coords - center) * scale
+
+    tensor_vertices = torch.from_numpy(voxel_coords).unsqueeze(0).reshape(1, -1, 3).cuda().to(torch.float16)
+    part_feats = sample_triplane_feat(part_planes, tensor_vertices) # N, M, C
+    part_feats = part_feats.cpu().numpy().reshape(-1, 448)
+    
+    return part_feats
+
+def sample_triplane_feat(feature_triplane, normalized_pos):
+    '''
+        normalized_pos [-1, 1]
+    '''
+    tri_plane = torch.unbind(feature_triplane, dim=1)
+
+    x_feat = F.grid_sample(
+        tri_plane[0],
+        torch.cat(
+            [normalized_pos[:, :, 0:1], normalized_pos[:, :, 1:2]],
+            dim=-1).unsqueeze(dim=1), padding_mode='border',
+        align_corners=True)
+    y_feat = F.grid_sample(
+        tri_plane[1],
+        torch.cat(
+            [normalized_pos[:, :, 1:2], normalized_pos[:, :, 2:3]],
+            dim=-1).unsqueeze(dim=1), padding_mode='border',
+        align_corners=True)
+
+    z_feat = F.grid_sample(
+        tri_plane[2],
+        torch.cat(
+            [normalized_pos[:, :, 0:1], normalized_pos[:, :, 2:3]],
+            dim=-1).unsqueeze(dim=1), padding_mode='border',
+        align_corners=True)
+    final_feat = (x_feat + y_feat + z_feat)
+    final_feat = final_feat.squeeze(dim=2).permute(0, 2, 1)  # 32dimension
+    return final_feat
+
+def cosegment_part(app_coords, app_part_planes, struct_coords, struct_part_planes, num_clusters=30):
+    struct_partfield_feats = get_voxel_partfeats(struct_coords, struct_part_planes)
+    app_partfield_feats = get_voxel_partfeats(app_coords, app_part_planes)
+    
+    point_feat1 = app_partfield_feats
+    point_feat2 = struct_partfield_feats
+    
+    point_feat1 = point_feat1 / np.linalg.norm(point_feat1, axis=-1, keepdims=True)
+    point_feat2 = point_feat2 / np.linalg.norm(point_feat2, axis=-1, keepdims=True)
+    
+    clustering1 = KMeans(n_clusters=num_clusters, random_state=0, n_init="auto").fit(point_feat1)
+    # Get feature means per cluster
+    feature_means1 = []
+    for j in range(num_clusters):
+        all_cluster_feat = point_feat1[clustering1.labels_==j]
+        mean_feat = np.mean(all_cluster_feat, axis=0)
+        feature_means1.append(mean_feat)
+    
+    labels1 = clustering1.labels_
+
+    feature_means1 = np.array(feature_means1)
+    tree = KDTree(feature_means1)
+    
+    init_mode = np.array(feature_means1)
+    
+    point_feat2 = point_feat2 / np.linalg.norm(point_feat2, axis=-1, keepdims=True)
+    clustering2 = KMeans(n_clusters=num_clusters, random_state=0, init=init_mode).fit(point_feat2)
+
+    ### Get feature means per cluster
+    feature_means2 = []
+    for j in range(num_clusters):
+        all_cluster_feat = point_feat2[clustering2.labels_==j]
+        mean_feat = np.mean(all_cluster_feat, axis=0)
+        feature_means2.append(mean_feat)
+
+    feature_means2 = np.array(feature_means2)
+    _, nn_idx = tree.query(feature_means2, k=1)
+    relabelled_2 = nn_idx[clustering2.labels_]
+    
+    return labels1, relabelled_2, point_feat1, point_feat2
+
+def cluster_geoms(struct_coords, struct_part_planes, num_clusters=10):
+    struct_partfield_feats = get_voxel_partfeats(struct_coords, struct_part_planes)
+    
+    point_feat = struct_partfield_feats
+    point_feat = point_feat / np.linalg.norm(point_feat, axis=-1, keepdims=True)
+    
+    
+    clustering = KMeans(n_clusters=num_clusters, random_state=0, n_init="auto").fit(point_feat)
+    # Get feature means per cluster
+    feature_means = []
+    for j in range(num_clusters):
+        all_cluster_feat = point_feat[clustering.labels_==j]
+        mean_feat = np.mean(all_cluster_feat, axis=0)
+        feature_means.append(mean_feat)
+    
+    labels = clustering.labels_
+    return labels

lib/util/pipeline_compat.py

@@ -0,0 +1,10 @@
+"""Conditioning compatibility checks without importing GPU dependencies."""
+
+
+def can_reuse_appearance_pipeline(pipeline, appearance_type):
+    """Image appearance needs the image pipeline's preprocessing interface."""
+    if pipeline is None:
+        return False
+    if appearance_type == "image":
+        return callable(getattr(pipeline, "preprocess_image", None))
+    return True

lib/util/pointcloud.py

@@ -0,0 +1,24 @@
+import numpy as np
+import trimesh
+import utils3d
+import open3d_pycg as o3d
+
+def convert_mesh_yup_to_zup(mesh):
+    # Row vectors v' = v @ R. Use the inverse of the old rotation so +Y maps to +Z vs -Z the other way.
+    mesh.vertices = mesh.vertices @ np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]])
+    return mesh
+
+def voxelize_mesh(mesh_file, save_path):
+    assert mesh_file.endswith('.ply') and save_path.endswith('.ply'), 'Voxelization only supports .ply files'
+
+    mesh = o3d.io.read_triangle_mesh(mesh_file)
+
+    # clamp vertices to the range [-0.5, 0.5]
+    vertices = np.clip(np.asarray(mesh.vertices), -0.5 + 1e-6, 0.5 - 1e-6)
+    mesh.vertices = o3d.utility.Vector3dVector(vertices)
+    voxel_grid = o3d.geometry.VoxelGrid.create_from_triangle_mesh_within_bounds(mesh, voxel_size=1/64, min_bound=(-0.5, -0.5, -0.5), max_bound=(0.5, 0.5, 0.5))
+    vertices = np.array([voxel.grid_index for voxel in voxel_grid.get_voxels()])
+    assert np.all(vertices >= 0) and np.all(vertices < 64), "Some vertices are out of bounds"
+    vertices = (vertices + 0.5) / 64 - 0.5
+
+    utils3d.io.write_ply(save_path, vertices)

lib/util/render.py

@@ -0,0 +1,133 @@
+import os
+import json
+from pathlib import Path
+import shutil
+import subprocess
+from urllib.request import urlretrieve
+import numpy as np
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+BLENDER_LINK = 'https://download.blender.org/release/Blender3.0/blender-3.0.1-linux-x64.tar.xz'
+BLENDER_INSTALLATION_PATH = os.environ.get('SPACEFLOW_BLENDER_ROOT', str(REPO_ROOT))
+BLENDER_PATH = os.environ.get(
+    'SPACEFLOW_BLENDER_PATH',
+    f'{BLENDER_INSTALLATION_PATH}/blender-3.0.1-linux-x64/blender',
+)
+BLENDER_RENDER_SCRIPT = os.environ.get(
+    'SPACEFLOW_BLENDER_RENDER_SCRIPT',
+    str(REPO_ROOT / 'third_party' / 'TRELLIS' / 'dataset_toolkits' / 'blender_script' / 'render.py'),
+)
+
+def _install_blender():
+    if shutil.which(BLENDER_PATH):
+        return
+    if os.environ.get('SPACEFLOW_BLENDER_PATH'):
+        raise FileNotFoundError(
+            f'Blender is not executable at SPACEFLOW_BLENDER_PATH={BLENDER_PATH}. '
+            'Set it to an installed Blender 3.x executable.'
+        )
+    installation = Path(BLENDER_INSTALLATION_PATH)
+    installation.mkdir(parents=True, exist_ok=True)
+    archive = installation / 'blender-3.0.1-linux-x64.tar.xz'
+    try:
+        urlretrieve(BLENDER_LINK, archive)
+        subprocess.run(['tar', '-xf', str(archive), '-C', str(installation)], check=True)
+    except (OSError, subprocess.CalledProcessError) as exc:
+        raise RuntimeError(
+            'Blender installation failed. Install Blender 3.x separately and set '
+            'SPACEFLOW_BLENDER_PATH to its executable.'
+        ) from exc
+    if not shutil.which(BLENDER_PATH):
+        raise FileNotFoundError(f'Blender installation did not produce an executable: {BLENDER_PATH}')
+
+
+def _run_blender(args):
+    result = subprocess.run(args, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
+    if result.returncode:
+        raise RuntimeError(
+            f'Blender failed with exit code {result.returncode}. '
+            f'Executable: {BLENDER_PATH}. Error output:\n{result.stderr[-3000:]}'
+        )
+
+def render_all_views(file_path, output_folder, num_views=150):
+    _install_blender()
+    # Build camera {yaw, pitch, radius, fov}
+    yaws = []
+    pitchs = []
+    offset = (np.random.rand(), np.random.rand())
+    for i in range(num_views):
+        y, p = sphere_hammersley_sequence(i, num_views, offset)
+        yaws.append(y)
+        pitchs.append(p)
+    radius = [2] * num_views
+    fov = [40 / 180 * np.pi] * num_views
+    views = [{'yaw': y, 'pitch': p, 'radius': r, 'fov': f} for y, p, r, f in zip(yaws, pitchs, radius, fov)]
+    
+    args = [
+        BLENDER_PATH, '-b', '-P', BLENDER_RENDER_SCRIPT,
+        '--',
+        '--views', json.dumps(views),
+        '--object', os.path.expanduser(file_path),
+        '--resolution', '512',
+        '--output_folder', output_folder,
+        '--engine', 'CYCLES',
+        '--save_mesh',
+    ]
+    if file_path.endswith('.blend'):
+        args.insert(1, file_path)
+    
+    _run_blender(args)
+    
+    if os.path.exists(os.path.join(output_folder, 'transforms.json')):
+        return True
+
+
+def export_normalized_mesh(file_path, output_folder):
+    _install_blender()
+    args = [
+        BLENDER_PATH, '-b', '-P', BLENDER_RENDER_SCRIPT,
+        '--',
+        '--views', '[]',
+        '--object', os.path.expanduser(file_path),
+        '--resolution', '512',
+        '--output_folder', output_folder,
+        '--engine', 'CYCLES',
+        '--save_mesh',
+    ]
+    if file_path.endswith('.blend'):
+        args.insert(1, file_path)
+
+    _run_blender(args)
+
+    if os.path.exists(os.path.join(output_folder, 'mesh.ply')):
+        return True
+
+# ===============LOW DISCREPANCY SEQUENCES================
+
+PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53]
+
+def radical_inverse(base, n):
+    val = 0
+    inv_base = 1.0 / base
+    inv_base_n = inv_base
+    while n > 0:
+        digit = n % base
+        val += digit * inv_base_n
+        n //= base
+        inv_base_n *= inv_base
+    return val
+
+def halton_sequence(dim, n):
+    return [radical_inverse(PRIMES[dim], n) for dim in range(dim)]
+
+def hammersley_sequence(dim, n, num_samples):
+    return [n / num_samples] + halton_sequence(dim - 1, n)
+
+def sphere_hammersley_sequence(n, num_samples, offset=(0, 0)):
+    u, v = hammersley_sequence(2, n, num_samples)
+    u += offset[0] / num_samples
+    v += offset[1]
+    u = 2 * u if u < 0.25 else 2 / 3 * u + 1 / 3
+    theta = np.arccos(1 - 2 * u) - np.pi / 2
+    phi = v * 2 * np.pi
+    return [phi, theta]

lib/util/visualization.py

@@ -0,0 +1,93 @@
+import logging as log
+import os.path as osp
+
+import imageio
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.colors as mcolors
+import numpy as np
+import torch
+from sklearn.neighbors import KDTree
+from tqdm import tqdm
+
+from third_party.TRELLIS.trellis.representations.mesh.cube2mesh import MeshExtractResult
+from third_party.TRELLIS.trellis.renderers.mesh_renderer import MeshRenderer
+from third_party.TRELLIS.trellis.utils.render_utils import yaw_pitch_r_fov_to_extrinsics_intrinsics
+
+
+def map_voxel_labels_to_vertices(mesh_vertices, voxel_coords, voxel_labels):
+    """Assign each mesh vertex the label of its nearest voxel via KD-tree lookup.
+
+    Args:
+        mesh_vertices: np.ndarray (V, 3) — mesh vertex positions
+        voxel_coords:  np.ndarray (M, 3) — voxel center positions (same coordinate space)
+        voxel_labels:  np.ndarray (M,)   — integer label per voxel
+
+    Returns:
+        np.ndarray (V,) — integer label per vertex
+    """
+    tree = KDTree(voxel_coords)
+    _, idx = tree.query(mesh_vertices, k=1)
+    return voxel_labels[idx.flatten()]
+
+
+def visualize_and_save(structure, partfield_clusters, output_folder,
+                       output_name='cluster_visualization.mp4',
+                       num_frames=300, resolution=512):
+    """Render a smooth rotating orbit video of the structure colored by cluster labels.
+
+    Uses the TRELLIS MeshRenderer (nvdiffrast, ssaa=4) and the same yaw/pitch
+    camera trajectory as the pipeline's Gaussian scene videos so all sides are visible.
+
+    Args:
+        structure:          trimesh.Trimesh — normalized structure mesh
+                            (vertices in [-0.5, 0.5])
+        partfield_clusters: np.ndarray (V,) — per-vertex integer cluster labels,
+                            where V matches structure.vertices.shape[0]
+        output_folder:      str  — directory in which to save the video
+        output_name:        str  — filename for the saved video
+                            (default 'cluster_visualization.mp4')
+        num_frames:         int  — number of video frames (default 300)
+        resolution:         int  — output resolution in pixels (default 512)
+    """
+    partfield_clusters = np.asarray(partfield_clusters, dtype=int)
+    num_clusters = int(partfield_clusters.max()) + 1
+
+    # Evenly-spaced HSV hues → maximally distinct colors
+    hues = np.linspace(0, 1, num_clusters, endpoint=False)
+    cluster_colors = np.array([mcolors.hsv_to_rgb([h, 0.90, 0.95]) for h in hues],
+                               dtype=np.float32)  # (num_clusters, 3)
+    vertex_colors = cluster_colors[partfield_clusters]  # (V, 3)
+
+    # Build MeshExtractResult for the TRELLIS MeshRenderer
+    verts_t = torch.tensor(structure.vertices, dtype=torch.float32, device='cuda')
+    faces_t = torch.tensor(structure.faces, dtype=torch.long, device='cuda')
+    colors_t = torch.tensor(vertex_colors, dtype=torch.float32, device='cuda')
+    mesh_result = MeshExtractResult(vertices=verts_t, faces=faces_t, vertex_attrs=colors_t)
+
+    # Camera trajectory: identical to render_utils.render_video
+    yaws = torch.linspace(0, 2 * torch.pi, num_frames).tolist()
+    pitch = (0.25 + 0.5 * torch.sin(torch.linspace(0, 2 * torch.pi, num_frames))).tolist()
+    extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitch, 2, 40)
+
+    renderer = MeshRenderer(rendering_options={
+        "resolution": resolution,
+        "near": 1,
+        "far": 100,
+        "ssaa": 4,
+    })
+
+    frames = []
+    log.info(f"Rendering {num_frames} frames → {output_name} ...")
+    with torch.no_grad():
+        for extr, intr in tqdm(zip(extrinsics, intrinsics), total=num_frames,
+                               desc=f'Rendering {output_name}'):
+            res = renderer.render(mesh_result, extr, intr,
+                                  return_types=["color", "mask"])
+            # composite cluster colors over black background
+            frame = (res['color'] * res['mask']).permute(1, 2, 0).clamp(0.0, 1.0)
+            frames.append((frame.cpu().numpy() * 255).astype(np.uint8))
+
+    video_path = osp.join(output_folder, output_name)
+    imageio.mimsave(video_path, frames, fps=30)
+    log.info(f"Saved {video_path}")

packages.txt

@@ -0,0 +1,4 @@
+libgl1
+libglib2.0-0
+libgomp1
+libxrender1

requirements.txt

@@ -0,0 +1,38 @@
+setuptools==80.9.0
+torch==2.8.0
+torchvision==0.23.0
+numpy==1.26.4
+scipy==1.15.3
+pillow==12.0.0
+imageio==2.37.3
+imageio-ffmpeg==0.6.0
+tqdm
+easydict==1.13
+opencv-python-headless==4.11.0.86
+rembg==2.0.69
+onnxruntime==1.23.2
+trimesh==4.11.4
+open3d_pycg-cpu==0.19.0
+xatlas==0.0.11
+pyvista==0.47.1
+pymeshfix==0.18.0
+igraph==1.0.0
+transformers==5.3.0
+safetensors
+omegaconf==2.3.0
+lightning==2.2.0
+yacs==0.1.8
+scikit-image==0.25.2
+scikit-learn==1.7.2
+loguru==0.7.3
+boto3
+psutil
+pytz
+einops==0.8.2
+matplotlib==3.10.8
+plyfile==1.1.3
+spconv-cu120==2.3.6
+git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057d642626897ec8
+https://data.pyg.org/whl/torch-2.8.0%2Bcu128/torch_scatter-2.1.2%2Bpt28cu128-cp310-cp310-linux_x86_64.whl
+https://huggingface.co/datasets/multimodalart/zerogpu-blackwell-wheels/resolve/main/wheels/pt28-cu128-cp310/flash_attn-2.8.3-cp310-cp310-linux_x86_64.whl
+https://huggingface.co/datasets/multimodalart/zerogpu-blackwell-wheels/resolve/main/wheels/pt28-cu128-cp310/nvdiffrast-0.4.0-cp310-cp310-linux_x86_64.whl

requirements/model-revisions.json

@@ -0,0 +1,23 @@
+{
+  "huggingface": {
+    "openai/clip-vit-large-patch14": "32bd64288804d66eefd0ccbe215aa642df71cc41",
+    "microsoft/TRELLIS-text-xlarge": "e0b00432b8e3a8ecee0df806ab1df9f7281f2be4",
+    "microsoft/TRELLIS-image-large": "25e0d31ffbebe4b5a97464dd851910efc3002d96",
+    "mikaelaangel/partfield-ckpt": "90b9b1e08b6a12fdcb6ee26b4854a26235e1765f"
+  },
+  "partfield_sha256": "463efc8a3afd3913142aa025e0125c00f16ef452b8de6a132ebe32bbe7877ee4",
+  "dinov2": {
+    "repository": "facebookresearch/dinov2",
+    "revision": "7764ea0f912e53c92e82eb78a2a1631e92725fc8",
+    "model": "dinov2_vitl14_reg",
+    "checkpoint": "dinov2_vitl14_reg4_pretrain.pth",
+    "sha256": "36e4deffbaef061a2576705b0c36f93621e2ae20bf6274694821b0b492551b51",
+    "source": "Upstream source pinned for fresh release verification; official Meta checkpoint downloaded and hashed on pf-pc69."
+  },
+  "u2net": {
+    "url": "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx",
+    "md5": "60024c5c889badc19c04ad937298a77b",
+    "source": "Official download URL and checksum in the pinned rembg 2.0.69 U2netSession."
+  },
+  "source": "TRELLIS/CLIP revisions recovered from original cache refs; PartField pinned from its official linked model repository."
+}

run_local_tau.py

@@ -0,0 +1,773 @@
+import copy
+import json
+import os
+import os.path as osp
+import gc
+import shutil
+import trimesh
+from PIL import Image
+import logging as log
+from omegaconf import OmegaConf
+import argparse
+import random
+import time
+import numpy as np
+from skimage import measure
+
+import torch
+from lightning.pytorch import seed_everything, Trainer
+import open3d_pycg as o3d
+
+import utils3d
+
+import sys
+sys.path.append('.')
+
+from third_party.PartField.partfield.model_trainer_pvcnn_only_demo import Model
+from lib.opt import self_similarity
+from lib.util import common, render, pointcloud
+from lib.util.checkpoint_compat import partfield_checkpoint_scope
+from third_party.TRELLIS.trellis.pipelines import TrellisTextTo3DPipeline
+from utils import merge_meshes
+
+log.getLogger().setLevel(log.INFO)
+log.basicConfig(level=log.INFO,
+                format='%(asctime)s - %(levelname)s - %(message)s',
+                datefmt='%Y-%m-%d %H:%M:%S')
+
+STEPS_SHAPE_GEN = 12
+CFG_SHAPE_GEN = 7.5
+INPUT_SQ_COLORED_GLB = "input_superquadrics_colored.glb"
+INPUT_SQ_HIGH_COLOR = np.array([0xf5, 0x9e, 0x0b, 0xff], dtype=np.uint8)
+INPUT_SQ_LOW_COLOR = np.array([0xf8, 0xfa, 0xfc, 0xff], dtype=np.uint8)
+
+def reset_run_state(seed=0):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def trellis_pipeline_path_from_args(args, cfg):
+    return (
+        args.trellis_pipeline_path
+        or os.environ.get("SPACEFLOW_TRELLIS_PIPELINE_PATH")
+        or cfg.trellis_text_model_name
+    )
+
+
+def load_trellis_pipeline(args, cfg):
+    trellis_pipeline_path = trellis_pipeline_path_from_args(args, cfg)
+    log.info(f"Loading TRELLIS pipeline from: {trellis_pipeline_path}")
+    start = time.perf_counter()
+    pipeline = TrellisTextTo3DPipeline.from_pretrained(trellis_pipeline_path)
+    log.info(f"Loaded TRELLIS pipeline weights and CLIP in {time.perf_counter() - start:.2f}s")
+    return pipeline
+
+
+def move_trellis_text_conditioner(pipeline, device):
+    text_cond_model = getattr(pipeline, 'text_cond_model', None)
+    if not text_cond_model:
+        return
+    if 'model' in text_cond_model:
+        text_cond_model['model'].to(device)
+    if 'null_cond' in text_cond_model:
+        text_cond_model['null_cond'] = text_cond_model['null_cond'].to(device)
+
+
+def offload_trellis_pipeline(pipeline):
+    if pipeline is None:
+        return
+    pipeline.cpu()
+    move_trellis_text_conditioner(pipeline, 'cpu')
+    torch.cuda.empty_cache()
+
+
+def init_args(argv=None):
+    parser = argparse.ArgumentParser(description='SpaceFlow - controllable 3D generation')
+
+    # Guidance mode selection
+    parser.add_argument('--guidance_mode', type=str, required=True, choices=['similarity'],
+                        help='Guidance mode. Only "similarity" is supported.')
+    parser.add_argument('--output_dir', type=str, required=True,
+                        help='Output directory for results')
+    parser.add_argument('--convert_yup_to_zup', action='store_true',
+                        help='Convert Y-up coordinate system to Z-up')
+
+    parser.add_argument('--appearance_image', type=str,
+                        help='Path to global image prompt for similarity guidance')
+    parser.add_argument('--appearance_text', type=str, default='',
+                        help='Global text prompt for similarity guidance')
+
+    # Spatial control parameters
+    parser.add_argument('--shape_superquadric_path', type=str, required=True,
+                        help='Path to shape superquadrics file')
+    parser.add_argument('--shape_superquadric_high_control_path', type=str, default=None,
+                        help='Path to high control shape superquadrics file (used local_tau_mode is guidance or masking)')
+    parser.add_argument('--low_control_superquadric_mask_path', type=str, default=None,
+                        help='Path to low control superquadric mask (Only used when local_tau_mode is low_control_mask)')
+
+
+    parser.add_argument('--spatial_control_mesh_path', type=str, default=None,
+                        help='Path to save the spatial control mesh (defaults to <output_dir>/spatial_control_mesh.ply)')
+    
+    parser.add_argument('--shape_tau', type=float, default=6.0, required=True,
+                        help='Value of tau for superquadric control')
+    parser.add_argument('--shape_tau_high_control', type=float, default=None, required=False,
+                        help='Value of tau for superquadric control')
+    parser.add_argument('--polyak_update_tau', type=float, default=0.18, required=False,
+                        help='Tau value for Polyak averaging of the high-control model (if using high control). If set to 0 it just perfoms spacecontrol with the lowest provided tau value.')
+    parser.add_argument('--local_tau_mode', type=str, choices=['guidance', 'masking', 'low_control_mask'], default='guidance',
+                        help='Whether to use local tau guidance, masking or low control mask mode. ')
+    parser.add_argument('--full_pipeline', action='store_true',
+                        help='Continue past structure generation into PartField and similarity optimization. Default keeps the legacy structure-only behavior.')
+    parser.add_argument('--n_repaint_steps', type=int, default=10,
+                        help='Number of repaint resampling steps to perform during structure generation to improve blending (default: 10). Set to 0 to disable.')                        
+    parser.add_argument('--structure_seed', type=int, default=1,
+                        help='Random seed for TRELLIS sparse-structure generation (default: 1).')
+    parser.add_argument('--texture_optim_steps', type=int, default=None,
+                        help='Override config/default.yaml sim_guidance.steps for texture similarity optimization. Minimum: 2.')
+    parser.add_argument('--mixed_self_attn_boost', type=float, default=None,
+                        help='Override config/default.yaml sim_guidance.mixed_self_attn_boost. '
+                             'Multiply self-attention weights by this factor before the softmax '
+                             'between voxels sharing a routed condition. 1.0 = stock attention.')
+    parser.add_argument('--trellis_pipeline_path', type=str, default=None,
+                        help='TRELLIS pipeline config/model path. Defaults to SPACEFLOW_TRELLIS_PIPELINE_PATH or config/default.yaml trellis_text_model_name.')
+    parser.add_argument('--geometry_only_decode', action='store_true',
+                        help='Decode out_sim.glb as white geometry only. Use with text-only lightweight pipeline configs that do not include a Gaussian decoder.')
+
+
+    parser.add_argument('--text_prompt', type=str, required=True,
+                        help='Text prompt for 3D shape generation')
+    parser.add_argument('--local_text_prompts', type=str, default=None,
+                        help='JSON-encoded list of per-SQ local text prompts (text-similarity mode)')
+    parser.add_argument('--local_image_paths', type=str, default=None,
+                        help='JSON-encoded list of per-SQ local image file paths (image-similarity mode)')
+
+    args = parser.parse_args(argv)
+
+    if args.appearance_text and args.appearance_image:
+        parser.error("Provide either --appearance_image or --appearance_text for similarity guidance, not both.")
+
+    if not args.appearance_text and not args.appearance_image:
+        parser.error("Provide either --appearance_image or --appearance_text for similarity guidance.")
+
+    return args
+
+def add_superquadric_compact_rot_mat(
+    scalings: np.array=np.array([1.0, 1.0, 1.0]),
+    exponents: np.array=np.array([2.0, 2.0, 2.0]),
+    translation: np.array=np.array([0.0, 0.0, 0.0]),
+    rotation: np.array=np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0],[0.0, 0.0,1.0]]),
+    tapering=None,
+    bending=None,
+    resolution: int=10,
+    visible: bool=True):
+    """Build a superquadric mesh in scene coordinates."""
+
+    def apply_taper(x, y, z, c, kx, ky):
+        c = float(c) if abs(float(c)) > 1e-8 else 1e-8
+        z_norm = z / c
+        x *= float(kx) * z_norm + 1.0
+        y *= float(ky) * z_norm + 1.0
+
+    def apply_bending_axis(x, y, z, kb, alpha, axis):
+        kb = float(kb)
+        if abs(kb) < 1e-3:
+            return
+        alpha = float(alpha)
+        if axis == "z":
+            u, v, w = x.copy(), y.copy(), z.copy()
+        elif axis == "x":
+            u, v, w = y.copy(), z.copy(), x.copy()
+        elif axis == "y":
+            u, v, w = z.copy(), x.copy(), y.copy()
+        else:
+            raise ValueError(axis)
+
+        sin_alpha = np.sin(alpha)
+        cos_alpha = np.cos(alpha)
+        beta = np.arctan2(v, u)
+        r = np.sqrt(u * u + v * v) * np.cos(alpha - beta)
+        inv_kb = 1.0 / kb
+        gamma = w * kb
+        rho = inv_kb - r
+        rb = inv_kb - rho * np.cos(gamma)
+        expr = rb - r
+        u = u + expr * cos_alpha
+        v = v + expr * sin_alpha
+        w = rho * np.sin(gamma)
+
+        if axis == "z":
+            x[:], y[:], z[:] = u, v, w
+        elif axis == "x":
+            x[:], y[:], z[:] = w, u, v
+        else:
+            x[:], y[:], z[:] = v, w, u
+
+    def create_superquadric_mesh(A, B, C, e1, e2, N):
+        def f(o, m):
+            return np.sign(np.sin(o)) * np.abs(np.sin(o))**m
+        def g(o, m):
+            return np.sign(np.cos(o)) * np.abs(np.cos(o))**m
+        u = np.linspace(-np.pi, np.pi, N, endpoint=True)
+        v = np.linspace(-np.pi/2.0, np.pi/2.0, N, endpoint=True)
+        u = np.tile(u, N)
+        v = (np.repeat(v, N))
+        if np.linalg.det(rotation) < 0:
+            u = u[::-1]
+        triangles = []
+
+        x = A * g(v, e1) * g(u, e2)
+        y = B * g(v, e1) * f(u, e2)
+        z = C * f(v, e1)
+        # Set poles to zero to account for numerical instabilities in f and g due to ** operator
+        x[:N] = 0.0
+        x[-N:] = 0.0
+        if tapering is not None:
+            apply_taper(x, y, z, C, tapering[0], tapering[1])
+        if bending is not None:
+            # Packed as [k_z, alpha_z, k_x, alpha_x, k_y, alpha_y].
+            apply_bending_axis(x, y, z, bending[4], bending[5], "y")
+            apply_bending_axis(x, y, z, bending[2], bending[3], "x")
+            apply_bending_axis(x, y, z, bending[0], bending[1], "z")
+        vertices =  np.concatenate([np.expand_dims(x, 1),
+                                    np.expand_dims(y, 1),
+                                    np.expand_dims(z, 1)], axis=1)
+        vertices = (rotation @ vertices.T).T + translation
+
+        triangles = []
+        for i in range(N-1):
+            for j in range(N-1):
+                triangles.append([i*N+j, i*N+j+1, (i+1)*N+j])
+                triangles.append([(i+1)*N+j, i*N+j+1, (i+1)*N+(j+1)])
+        # Connect first and last vertex in each row
+        for i in range(N - 1):
+            triangles.append([i * N + (N - 1), i * N, (i + 1) * N + (N - 1)])
+            triangles.append([(i + 1) * N + (N - 1), i * N, (i + 1) * N])
+
+        triangles.append([(N-1)*N+(N-1), (N-1)*N, (N-1)])
+        triangles.append([(N-1), (N-1)*N, 0])
+
+        return vertices, triangles
+
+
+    vertices, triangles = create_superquadric_mesh(scalings[0], scalings[1], scalings[2],
+                                                exponents[0], exponents[1],
+                                                resolution)
+    return vertices, triangles
+
+def load_superquadric_from_file(file_path: str) -> list:
+    par_dict = np.load(file_path)
+    print(par_dict)
+    scale = par_dict['scales']        # 3 (3x1 vector)
+    rotate = par_dict['rotations']    # 3 (3x3 rotation matrix)
+    shapes = par_dict['shapes']       # 2 (2x1 vector)
+    trans = par_dict['translations']  # 3 (3x1 vector)
+    num_el = scale.shape[0]           # number of superquadrics
+    tapering = par_dict['tapering'] if 'tapering' in par_dict else np.zeros((num_el, 2))
+    bending = par_dict['bending'] if 'bending' in par_dict else np.zeros((num_el, 6))
+    control_levels = par_dict['control_levels'] if 'control_levels' in par_dict else np.ones((num_el,))
+
+    superquadrics = {}
+    for k in range(num_el):
+        superquadric_dict = {}
+        superquadric_dict['scale'] = scale[k, :]
+        superquadric_dict['shape'] = shapes[k]
+        superquadric_dict['rotation'] = rotate[k, :]
+        superquadric_dict['translation'] = trans[k, :]
+        superquadric_dict['tapering'] = tapering[k, :]
+        superquadric_dict['bending'] = bending[k, :]
+        superquadric_dict['control_level'] = 'low' if float(control_levels[k]) < 0.5 else 'high'
+        superquadric_dict['color'] = [90, 200, 255]
+        superquadrics[k] = superquadric_dict
+    return superquadrics
+
+def export_colored_superquadrics_glb(npz_path, output_path, aabb=None, center=None, scale=None):
+    start = time.perf_counter()
+    superquadrics = load_superquadric_from_file(npz_path)
+    meshes = []
+    raw_meshes = []
+    for sq_id in superquadrics:
+        sq = superquadrics[sq_id]
+        vertices, triangles = add_superquadric_compact_rot_mat(
+            sq['scale'],
+            sq['shape'],
+            sq['translation'],
+            sq['rotation'],
+            sq['tapering'],
+            sq['bending'],
+            resolution=100,
+        )
+        raw_mesh = trimesh.Trimesh(vertices=vertices, faces=np.asarray(triangles), process=False)
+        raw_meshes.append(raw_mesh)
+
+    if not raw_meshes:
+        raise ValueError(f"No superquadrics found in {npz_path}")
+
+    if aabb is None or center is None or scale is None:
+        all_vertices = np.concatenate([mesh.vertices for mesh in raw_meshes], axis=0)
+        aabb = np.stack([all_vertices.min(0), all_vertices.max(0)])
+        center = (aabb[0] + aabb[1]) / 2
+        scale = 1.0 / ((aabb[1] - aabb[0]).max())
+
+    for raw_mesh, sq_id in zip(raw_meshes, superquadrics):
+        sq = superquadrics[sq_id]
+        mesh = raw_mesh.copy()
+        mesh.vertices = (mesh.vertices - center) * scale
+        color = INPUT_SQ_LOW_COLOR if sq.get('control_level') == 'low' else INPUT_SQ_HIGH_COLOR
+        mesh.visual = trimesh.visual.ColorVisuals(
+            mesh,
+            vertex_colors=np.tile(color, (mesh.vertices.shape[0], 1)),
+        )
+        meshes.append(mesh)
+
+    scene = trimesh.Scene()
+    for sq_id, mesh in zip(superquadrics, meshes):
+        scene.add_geometry(mesh, node_name=f"superquadric_{sq_id}", geom_name=f"superquadric_{sq_id}")
+    scene.export(output_path)
+    log.info(
+        "Colored input superquadrics GLB exported in %.2fs: %s (%d SQs)",
+        time.perf_counter() - start,
+        output_path,
+        len(meshes),
+    )
+
+def load_superquadrics(path, spatial_control_mesh_path, aabb=None, center=None, scale=None):
+    start = time.perf_counter()
+    # Generate spatial control mesh from superquadric primitives and write to spatial_control_mesh_path
+    superquadrics = load_superquadric_from_file(path)
+
+    meshes = []
+    for superquadric_id in superquadrics.keys():
+        vertices, triangles = add_superquadric_compact_rot_mat(
+        superquadrics[superquadric_id]['scale'],
+        superquadrics[superquadric_id]['shape'],
+        superquadrics[superquadric_id]['translation'],
+        superquadrics[superquadric_id]['rotation'],
+        superquadrics[superquadric_id]['tapering'],
+        superquadrics[superquadric_id]['bending'],
+        resolution=100)
+        mesh = o3d.geometry.TriangleMesh()
+        mesh.vertices = o3d.utility.Vector3dVector(vertices)
+        mesh.triangles = o3d.utility.Vector3iVector(triangles)
+        meshes.append(mesh)
+    merged_mesh = merge_meshes(meshes)
+    aabb = np.stack([np.asarray(merged_mesh.vertices).min(0), np.asarray(merged_mesh.vertices).max(0)]) if aabb is None else aabb
+    center = (aabb[0] + aabb[1]) / 2 if center is None else center
+    scale = 1/((aabb[1] - aabb[0]).max())  if scale is None else scale
+
+    merged_mesh.translate(-center)
+    merged_mesh.scale(scale, (0,0,0))
+    o3d.io.write_triangle_mesh(spatial_control_mesh_path, merged_mesh)
+    log.info(
+        "Spatial control mesh generated from superquadrics in %.2fs: %s (%d SQs)",
+        time.perf_counter() - start,
+        spatial_control_mesh_path,
+        len(superquadrics),
+    )
+
+    if aabb is not None and center is not None and scale is not None:
+        return aabb, center, scale
+
+def build_individual_sq_meshes_normalized(npz_path):
+    """Build individual superquadric meshes in the same normalised space as the merged mesh.
+
+    The merged mesh is centred and scaled to unit size (same as load_superquadrics).
+    Returns a list of open3d TriangleMesh objects in the [-0.5, 0.5] coordinate space.
+    """
+    superquadrics = load_superquadric_from_file(npz_path)
+    meshes = []
+    for sq_id in superquadrics:
+        vertices, triangles = add_superquadric_compact_rot_mat(
+            superquadrics[sq_id]['scale'],
+            superquadrics[sq_id]['shape'],
+            superquadrics[sq_id]['translation'],
+            superquadrics[sq_id]['rotation'],
+            superquadrics[sq_id]['tapering'],
+            superquadrics[sq_id]['bending'],
+            resolution=100)
+        mesh = o3d.geometry.TriangleMesh()
+        mesh.vertices = o3d.utility.Vector3dVector(vertices)
+        mesh.triangles = o3d.utility.Vector3iVector(triangles)
+        meshes.append(mesh)
+
+    # Compute normalization from merged mesh (mirrors load_superquadrics)
+    merged = merge_meshes(meshes)
+    all_verts = np.asarray(merged.vertices)
+    aabb = np.stack([all_verts.min(0), all_verts.max(0)])
+    center = (aabb[0] + aabb[1]) / 2
+    scale = 1.0 / ((aabb[1] - aabb[0]).max())
+
+    normalized = []
+    for mesh in meshes:
+        m = copy.deepcopy(mesh)
+        m.translate(-center)
+        m.scale(scale, (0, 0, 0))
+        normalized.append(m)
+    return normalized
+
+
+def count_superquadrics(npz_path):
+    with np.load(npz_path) as par_dict:
+        return int(par_dict['scales'].shape[0])
+
+
+def parse_local_condition_list(raw_json, num_superquadrics, label, require_existing_files=False):
+    if not raw_json:
+        return None
+    values = json.loads(raw_json)
+    if not isinstance(values, list):
+        raise ValueError(f"{label} must be a JSON list")
+    if len(values) > num_superquadrics:
+        raise ValueError(f"{label} has {len(values)} entries but only {num_superquadrics} superquadrics")
+    normalized = [str(value or "").strip() for value in values]
+    normalized.extend([""] * (num_superquadrics - len(normalized)))
+    if require_existing_files:
+        missing = [path for path in normalized if path and not osp.exists(path)]
+        if missing:
+            raise FileNotFoundError(f"Local image path(s) not found: {missing}")
+    return normalized if any(value for value in normalized) else None
+
+
+def sparse_voxels_to_glb(sparse_points, grid_size=64, output_filename="output.glb"):
+    """
+    Converts sparse voxel coordinates to a GLB mesh.
+
+    :param sparse_points: List or array of (x, y, z) coordinates (integers).
+    :param grid_size: The size of the voxel bounding box (e.g., 64).
+    :param output_filename: The name of the output GLB file.
+    """
+
+    start = time.perf_counter()
+    log.info("Creating %dx%dx%d sparse voxel grid...", grid_size, grid_size, grid_size)
+    # Init grid
+    voxel_grid = np.zeros((grid_size, grid_size, grid_size), dtype=bool)
+
+    # Populate grid with sparse points
+    sparse_points = np.round(sparse_points).astype(int)
+    for x, y, z in sparse_points:
+        # Check bounds just in case
+        if 0 <= x < grid_size and 0 <= y < grid_size and 0 <= z < grid_size:
+            voxel_grid[x, y, z] = True
+
+    # Padding the grid (Needed for marching cubes)
+    padded_grid = np.pad(voxel_grid, pad_width=1, mode='constant', constant_values=False)
+
+    log.info("Running Marching Cubes algorithm...")
+    # Marching Cubes (Level set at 0.5 to extract the surface between occupied and empty voxels)
+    verts, faces, normals, values = measure.marching_cubes(padded_grid, level=0.5)
+
+    # Shift vertices back by 1 to account for the padding we added
+    verts = verts - 1.0
+
+    log.info("Generating mesh and exporting to GLB...")
+    mesh = trimesh.Trimesh(vertices=verts, faces=faces, vertex_normals=normals)
+
+    # # Smoothing of the mesh
+    # trimesh.smoothing.filter_taubin(mesh, iterations=50)
+
+    # Export to GLB format
+    mesh.export(output_filename)
+    log.info(
+        "Successfully exported sparse voxel mesh in %.2fs: %s",
+        time.perf_counter() - start,
+        output_filename,
+    )
+
+def predict_part(obj_path, output_dir):
+    start = time.perf_counter()
+    log.info("Extracting PartField feature planes...")
+    partfield_config = 'third_party/PartField/config.yaml'
+    partfield_cfg = OmegaConf.load(partfield_config)
+    partfield_cfg.dataset.val_num_workers = 0
+    partfield_ckpt = partfield_cfg.continue_ckpt
+    if partfield_ckpt and not osp.isabs(partfield_ckpt):
+        partfield_ckpt = osp.join(osp.dirname(osp.abspath(partfield_config)), partfield_ckpt)
+
+    seed_everything(partfield_cfg.seed)
+
+    torch.manual_seed(0)
+    random.seed(0)
+    np.random.seed(0)
+
+    # Lightning defaults to ./lightning_logs under cwd (repo root); team members often
+    # cannot mkdir there. Keep all PL artifacts under this run's output_dir.
+    pl_root = osp.join(output_dir, 'pl_partfield')
+    common.ensure_dir(pl_root)
+
+    trainer = Trainer(devices=1,
+                      accelerator="gpu",
+                      precision="16-mixed",
+                      max_epochs=partfield_cfg.training_epochs,
+                      log_every_n_steps=1,
+                      limit_train_batches=3500,
+                      limit_val_batches=None,
+                      default_root_dir=pl_root,
+                      logger=False,
+                      enable_checkpointing=False,
+                      enable_progress_bar=False,
+                      enable_model_summary=False,
+                     )
+
+    partfield_model = Model(partfield_cfg, obj_path)
+    with partfield_checkpoint_scope():
+        output = trainer.predict(partfield_model, ckpt_path=partfield_ckpt)
+    part_planes, uid = output[0]
+    np.save(f'{output_dir}/part_feat_{uid}_batch_part_plane.npy', part_planes)
+
+    del partfield_model
+    gc.collect() # Free up memory
+    log.info("Extracted PartField feature planes in %.2fs", time.perf_counter() - start)
+
+def run(args, cfg=None, generation_pipeline=None):
+    run_start = time.perf_counter()
+    reset_run_state()
+    cfg = cfg or OmegaConf.load('config/default.yaml')
+    if args.texture_optim_steps is not None:
+        if args.texture_optim_steps < 2:
+            raise ValueError("--texture_optim_steps must be an integer at least 2")
+        cfg = copy.deepcopy(cfg)
+        cfg.sim_guidance.steps = int(args.texture_optim_steps)
+        log.info("Overriding texture optimization steps: %d", int(cfg.sim_guidance.steps))
+    if args.mixed_self_attn_boost is not None:
+        if args.mixed_self_attn_boost < 1.0:
+            raise ValueError("--mixed_self_attn_boost must be at least 1.0 (1.0 = stock attention)")
+        cfg = copy.deepcopy(cfg)
+        cfg.sim_guidance.mixed_self_attn_boost = float(args.mixed_self_attn_boost)
+        log.info("Overriding mixed self-attention boost: %.2f",
+                 float(cfg.sim_guidance.mixed_self_attn_boost))
+
+    common.ensure_dir(args.output_dir)
+
+    def copy_input_npz(src_path, output_name):
+        if not src_path:
+            return
+        dst_path = osp.join(args.output_dir, output_name)
+        shutil.copy2(src_path, dst_path)
+        log.info(f"Copied input superquadrics: {dst_path}")
+
+    copy_input_npz(args.shape_superquadric_path, 'input_superquadrics_all.npz')
+    copy_input_npz(args.shape_superquadric_high_control_path, 'input_superquadrics_high_control.npz')
+    copy_input_npz(args.low_control_superquadric_mask_path, 'input_superquadrics_low_control_bbox.npz')
+
+    # Generate spatial control mesh from superquadrics
+    spatial_control_mesh_path = osp.join(args.output_dir, 'spatial_control_mesh.ply')
+    control_mesh_start = time.perf_counter()
+    aabb, center, scale = load_superquadrics(args.shape_superquadric_path, spatial_control_mesh_path)
+    export_colored_superquadrics_glb(
+        args.shape_superquadric_path,
+        osp.join(args.output_dir, INPUT_SQ_COLORED_GLB),
+        aabb=aabb,
+        center=center,
+        scale=scale,
+    )
+
+
+    low_control_superquadric_mask_path = None
+    if args.shape_tau_high_control is not None:
+        assert args.shape_tau_high_control >= args.shape_tau, "shape_tau_high_control must be greater than or equal to shape_tau"
+        
+        print(f"Using high control tau: {args.shape_tau_high_control} and low control tau: {args.shape_tau}, with local tau mode: {args.local_tau_mode}")
+        high_control_spatial_control_mesh_path = osp.join(args.output_dir, 'high_control_spatial_control_mesh.ply')
+        load_superquadrics(args.shape_superquadric_high_control_path, high_control_spatial_control_mesh_path, aabb=aabb, center=center, scale=scale)
+        
+        if args.local_tau_mode == 'low_control_mask':
+            low_control_superquadric_mask_path = osp.join(args.output_dir, 'low_control_superquadric_mask.ply')
+            load_superquadrics(args.low_control_superquadric_mask_path, low_control_superquadric_mask_path, aabb=aabb, center=center, scale=scale)
+    log.info("Prepared all superquadric control meshes in %.2fs", time.perf_counter() - control_mesh_start)
+
+
+    # Load structure mesh
+    log.info("Creating structure mesh with SpaceControl code...")
+
+    owns_pipeline = generation_pipeline is None
+    if owns_pipeline:
+        pipeline = load_trellis_pipeline(args, cfg)
+    else:
+        pipeline = generation_pipeline
+        log.info(f"Reusing preloaded TRELLIS pipeline from: {trellis_pipeline_path_from_args(args, cfg)}")
+    cuda_start = time.perf_counter()
+    pipeline.cuda()
+    move_trellis_text_conditioner(pipeline, 'cuda')
+    log.info(f"Moved TRELLIS pipeline to CUDA in {time.perf_counter() - cuda_start:.2f}s")
+
+    text_prompt = args.text_prompt
+
+    # Sparse voxels
+    structure_start = time.perf_counter()
+    log.info("Generating TRELLIS sparse structure with seed %d", int(args.structure_seed))
+    coords = pipeline.gen_structure_v2(text_prompt, seed=int(args.structure_seed), vis_output_dir=None, sparse_structure_sampler_params={
+        "steps": STEPS_SHAPE_GEN,
+        "cfg_strength": CFG_SHAPE_GEN,
+        "t0_idx_value": args.shape_tau,
+        "spatial_control_mesh_path": spatial_control_mesh_path,
+        "high_control_spatial_control_mesh_path": high_control_spatial_control_mesh_path if args.shape_tau_high_control is not None else None,
+        "low_control_superquadric_mask_path": low_control_superquadric_mask_path if low_control_superquadric_mask_path is not None else None,
+        "t0_idx_value_high_control": args.shape_tau_high_control if args.shape_tau_high_control is not None else None,
+        "polyak_update_tau": args.polyak_update_tau,
+        "local_tau_mode": args.local_tau_mode,
+        "n_repaint_steps": args.n_repaint_steps,
+    })
+    log.info(
+        "Generated TRELLIS sparse structure in %.2fs (%d voxels)",
+        time.perf_counter() - structure_start,
+        coords.shape[0],
+    )
+
+    # Convert sparse voxels to mesh
+    log.info("Converting sparse voxels to mesh...")
+
+    coords_np = coords.detach().cpu().numpy()
+
+    filtered_coords = coords_np[:, 1:]
+    print(f"Sparse voxel tensor shape: {coords_np.shape}")
+    print(f"Number of valid voxels: {filtered_coords.shape[0]}")
+
+    sparse_voxels_to_glb(filtered_coords, grid_size=64, output_filename=osp.join(args.output_dir, "sample.glb"))
+
+    # glb.export("sample.glb")
+
+    # log.info("Loading generated mesh...")
+
+    mesh_export_start = time.perf_counter()
+    struct_mesh = trimesh.load(osp.join(args.output_dir, "sample.glb"), force='mesh')
+    # Generator / marching-cubes output in Y-up; keep a copy for debugging.
+    struct_mesh.export(osp.join(args.output_dir, 'struct_mesh.glb'))
+    log.info("Loaded and exported raw structure mesh in %.2fs", time.perf_counter() - mesh_export_start)
+
+    if args.full_pipeline and args.guidance_mode == 'similarity':
+        log.info("Keeping TRELLIS pipeline loaded for similarity guidance; offloading it until refinement.")
+        offload_trellis_pipeline(pipeline)
+    else:
+        if owns_pipeline:
+            del pipeline
+            gc.collect() # Free up memory
+        else:
+            offload_trellis_pipeline(pipeline)
+
+    # Canonical mesh for renders, voxels, and PartField must share one frame (Z-up if converting).
+    normalize_start = time.perf_counter()
+    if args.convert_yup_to_zup:
+        struct_mesh = pointcloud.convert_mesh_yup_to_zup(struct_mesh)
+    struct_mesh.export(osp.join(args.output_dir, 'struct_mesh_zup.glb'))
+    log.info("Prepared canonical structure mesh in %.2fs", time.perf_counter() - normalize_start)
+    struct_mesh_for_pipeline = osp.join(args.output_dir, 'struct_mesh_zup.glb')
+
+    struct_render_dir = osp.join(args.output_dir, 'struct_renders')
+    common.ensure_dir(struct_render_dir)
+    if args.full_pipeline and args.guidance_mode == 'similarity':
+        render_start = time.perf_counter()
+        log.info("Exporting Blender-normalized structure mesh without PNG renders...")
+        out_renderviews = render.export_normalized_mesh(struct_mesh_for_pipeline, struct_render_dir)
+        log.info("Exported Blender-normalized structure mesh in %.2fs", time.perf_counter() - render_start)
+    else:
+        render_start = time.perf_counter()
+        log.info(f"Rendering structure mesh for {cfg.num_views // 10} views...")
+        out_renderviews = render.render_all_views(struct_mesh_for_pipeline, struct_render_dir, num_views=cfg.num_views // 10)
+        log.info("Rendered structure mesh views in %.2fs", time.perf_counter() - render_start)
+
+    # struct_renders/mesh.ply is the Blender-normalized mesh; use it as the single source of truth
+    # for both voxelization and PartField feature extraction so that both operate in the same
+    # coordinate system (Blender's normalize_scene + GLTF→Blender axis convention).
+    struct_blender_ply = osp.join(struct_render_dir, 'mesh.ply')
+
+    voxel_dir = osp.join(args.output_dir, 'voxels')
+    common.ensure_dir(voxel_dir)
+    voxel_start = time.perf_counter()
+    log.info("Voxelizing structure mesh...")
+    pointcloud.voxelize_mesh(struct_blender_ply, save_path=osp.join(voxel_dir, 'struct_voxels.ply'))
+    log.info("Voxelized structure mesh in %.2fs", time.perf_counter() - voxel_start)
+
+    if not args.full_pipeline:
+        log.info("Structure-only mode complete. Pass --full_pipeline to continue into PartField and refinement.")
+        return
+
+    log.info("Extracting Structure Mesh PartField feature planes...")
+    partfield_dir = osp.join(args.output_dir, 'partfield')
+    common.ensure_dir(partfield_dir)
+    # Use the same Blender-normalized PLY so the PartField triplane canonical space
+    # matches the coordinate system of struct_voxels.ply.
+    partfield_start = time.perf_counter()
+    predict_part(struct_blender_ply, partfield_dir)
+    log.info("Completed Structure Mesh PartField extraction in %.2fs", time.perf_counter() - partfield_start)
+
+    log.info("Visualizing PartField clusters on structure mesh...")
+    from lib.util.visualization import visualize_and_save, map_voxel_labels_to_vertices
+    from lib.util.partfield import cluster_geoms
+    _sv = utils3d.io.read_ply(osp.join(voxel_dir, 'struct_voxels.ply'))[0]
+    _sc = torch.from_numpy(_sv).float().cuda()
+    _sc4d = torch.cat([torch.zeros(_sc.shape[0], 1, dtype=torch.long, device='cuda'),
+                       ((_sc + 0.5) * 64).long()], dim=1)
+    _planes = torch.from_numpy(np.load(
+        osp.join(partfield_dir, 'part_feat_mesh_batch_part_plane.npy'),
+        allow_pickle=True)).cuda()
+    _vlabels = cluster_geoms(_sc4d, _planes, num_clusters=cfg.sim_guidance.num_part_clusters)
+    _mesh_vis = trimesh.load(struct_blender_ply, force='mesh')
+    _vtx_labels = map_voxel_labels_to_vertices(_mesh_vis.vertices, _sv, _vlabels)
+    visualize_and_save(_mesh_vis, _vtx_labels, args.output_dir, output_name='partfield_clusters.mp4')
+    del _sv, _sc, _sc4d, _planes, _vlabels, _mesh_vis, _vtx_labels
+    gc.collect()
+
+    if not out_renderviews:
+        log.info("Structure rendering failed!")
+
+    log.info("Running similarity-guided optimization...")
+
+    if args.appearance_image:
+        app_type = 'image'
+        app = args.appearance_image
+
+        app_image = Image.open(args.appearance_image).convert('RGB')
+        app_image.save(osp.join(args.output_dir, 'app_image.png'))
+
+    else:
+        app_type = 'text'
+        app = args.appearance_text
+
+    log.info(f"Using {app_type} for self-similarity guidance...")
+
+    # Parse per-SQ local conditioning args
+    num_superquadrics = count_superquadrics(args.shape_superquadric_path)
+    local_text_prompts = (
+        parse_local_condition_list(args.local_text_prompts, num_superquadrics, "local_text_prompts")
+        if app_type == 'text'
+        else None
+    )
+    local_image_paths = (
+        parse_local_condition_list(
+            args.local_image_paths,
+            num_superquadrics,
+            "local_image_paths",
+            require_existing_files=True,
+        )
+        if app_type == 'image'
+        else None
+    )
+
+    local_prompts     = local_text_prompts if app_type == 'text' else local_image_paths
+    local_prompt_type = app_type if local_prompts else None
+    individual_sq_meshes = build_individual_sq_meshes_normalized(args.shape_superquadric_path) if local_prompts else None
+
+    # Self-Similarity Optimization
+    sim_start = time.perf_counter()
+    self_similarity.optimize_self_similarity(
+        cfg, app, app_type, args.output_dir,
+        local_prompts=local_prompts,
+        local_prompt_type=local_prompt_type,
+        individual_sq_meshes=individual_sq_meshes,
+        generation_pipeline=pipeline,
+        decode_texture=not args.geometry_only_decode,
+    )
+    log.info("Completed SpaceFlow run in %.2fs (similarity optimization %.2fs)",
+             time.perf_counter() - run_start,
+             time.perf_counter() - sim_start)
+
+
+def main(argv=None):
+    args = init_args(argv)
+    run(args)
+
+if __name__ == "__main__":
+    main()

third_party/PartField/LICENSE

@@ -0,0 +1,36 @@
+NVIDIA License
+
+1. Definitions
+
+“Licensor” means any person or entity that distributes its Work.
+“Work” means (a) the original work of authorship made available under this license, which may include software, documentation, or other files, and (b) any additions to or derivative works  thereof  that are made available under this license.
+The terms “reproduce,” “reproduction,” “derivative works,” and “distribution” have the meaning as provided under U.S. copyright law; provided, however, that 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.
+Works are “made available” under this license by including in or with the Work either (a) a copyright notice referencing the applicability of this license to the Work, or (b) a copy of this license.
+
+2. License Grant
+
+2.1 Copyright Grant. Subject to the terms and conditions of this license, each Licensor grants to you a perpetual, worldwide, non-exclusive, royalty-free, copyright license to use, reproduce, prepare derivative works of, publicly display, publicly perform, sublicense and distribute its Work and any resulting derivative works in any form.
+
+3. Limitations
+
+3.1 Redistribution. You may reproduce or distribute the Work only if (a) you do so under this license, (b) you include a complete copy of this license with your distribution, and (c) you retain without modification any copyright, patent, trademark, or attribution notices that are present in the Work.
+
+3.2 Derivative Works. You may specify that additional or different terms apply to the use, reproduction, and distribution of your derivative works of the Work (“Your Terms”) only if (a) Your Terms provide that the use limitation in Section 3.3 applies to your derivative works, and (b) you identify the specific derivative works that are subject to Your Terms. Notwithstanding Your Terms, this license (including the redistribution requirements in Section 3.1) will continue to apply to the Work itself.
+
+3.3 Use Limitation. The Work and any derivative works thereof only may be used or intended for use non-commercially. Notwithstanding the foregoing, NVIDIA Corporation and its affiliates may use the Work and any derivative works commercially. As used herein, “non-commercially” means for non-commercial research and educational purposes only.
+
+3.4 Patent Claims. If you bring or threaten to bring a patent claim against any Licensor (including any claim, cross-claim or counterclaim in a lawsuit) to enforce any patents that you allege are infringed by any Work, then your rights under this license from such Licensor (including the grant in Section 2.1) will terminate immediately.
+
+3.5 Trademarks. This license does not grant any rights to use any Licensor’s or its affiliates’ names, logos, or trademarks, except as necessary to reproduce the notices described in this license.
+
+3.6 Termination. If you violate any term of this license, then your rights under this license (including the grant in Section 2.1) will terminate immediately.
+
+4. Disclaimer of Warranty.
+
+THE WORK IS PROVIDED “AS IS” WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF 
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER THIS LICENSE. 
+
+5. Limitation of Liability.
+
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK (INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION, LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
+

third_party/PartField/config.yaml

@@ -0,0 +1,88 @@
+seed: 0
+output_dir: "results"
+result_name: "demo"
+
+triplet_sampling: "random"
+load_original_mesh: false
+
+num_pos: 64
+num_neg_random: 256
+num_neg_hard_pc: 128
+num_neg_hard_emb: 128
+
+vertex_feature: false  # if true, sample feature on vertices; if false, sample feature on faces
+n_point_per_face: 1000
+n_sample_each: 10000
+preprocess_mesh: false
+
+regress_2d_feat: false
+
+is_pc: false
+
+cut_manifold: false
+remesh_demo: false
+correspondence_demo: false
+
+save_every_epoch: 10
+training_epochs: 30
+continue_training: false
+
+continue_ckpt: models/model_objaverse.ckpt
+epoch_selected: "epoch=50.ckpt"
+
+triplane_resolution: 128
+triplane_channels_low: 128
+triplane_channels_high: 512
+lr: 1.0e-3
+train: true
+test: false
+
+inference_save_pred_sdf_to_mesh: true
+inference_save_feat_pca: true
+name: "test"
+test_subset: false
+test_corres: false
+test_partobjaversetiny: false
+
+dataset:
+  type: "Mix"
+  data_path: "objaverse_data"
+  train_batch_size: 1
+  val_batch_size: 1
+  train_num_workers: 8
+  val_num_workers: 32
+  all_files: []  # only used for correspondence demo
+
+voxel2triplane:
+  transformer_dim: 1024
+  transformer_layers: 6
+  transformer_heads: 8
+  triplane_low_res: 32
+  triplane_high_res: 256
+  triplane_dim: 64
+  normalize_vox_feat: false
+
+loss:
+  triplet: 1.0
+  sdf: 1.0
+  feat: 10.0
+  l1: 0.0
+
+use_pvcnn: false
+use_pvcnnonly: true
+
+pvcnn:
+  point_encoder_type: 'pvcnn'
+  use_point_scatter: true
+  z_triplane_channels: 256
+  z_triplane_resolution: 128
+  unet_cfg:
+    depth: 3
+    enabled: true
+    rolled: true
+    use_3d_aware: true
+    start_hidden_channels: 32
+    use_initial_conv: false
+
+use_2d_feat: false
+inference_metrics_only: false

third_party/PartField/partfield/config/__init__.py

@@ -0,0 +1,26 @@
+import argparse
+import os.path as osp
+from datetime import datetime
+import pytz
+
+def default_argument_parser(add_help=True, default_config_file=""):
+    parser = argparse.ArgumentParser(add_help=add_help)
+    parser.add_argument("--config-file", '-c', default=default_config_file, metavar="FILE", help="path to config file")
+    parser.add_argument(
+        "--opts",
+        help="Modify config options using the command-line",
+        default=None,
+        nargs=argparse.REMAINDER,
+    )
+    return parser
+
+def setup(args, freeze=True):
+    from .defaults import _C as cfg
+    cfg = cfg.clone()
+    cfg.merge_from_file(args.config_file)
+    cfg.merge_from_list(args.opts)
+    dt = datetime.now(pytz.timezone('America/Los_Angeles')).strftime('%y%m%d-%H%M%S')
+    cfg.output_dir = osp.join(cfg.output_dir, cfg.name, dt)
+    if freeze:
+        cfg.freeze()
+    return cfg

third_party/PartField/partfield/config/defaults.py

@@ -0,0 +1,92 @@
+from yacs.config import CfgNode as CN
+
+_C = CN()
+_C.seed = 0
+_C.output_dir = "results"
+_C.result_name = "test_all"
+
+_C.triplet_sampling = "random"
+_C.load_original_mesh = False
+
+_C.num_pos = 64
+_C.num_neg_random = 256
+_C.num_neg_hard_pc = 128
+_C.num_neg_hard_emb = 128
+
+_C.vertex_feature = False  # if true, sample feature on vertices; if false, sample feature on faces
+_C.n_point_per_face = 2000
+_C.n_sample_each = 10000
+_C.preprocess_mesh = False
+
+_C.regress_2d_feat = False
+
+_C.is_pc = False
+
+_C.cut_manifold = False
+_C.remesh_demo = False
+_C.correspondence_demo = False
+
+_C.save_every_epoch = 10
+_C.training_epochs = 30
+_C.continue_training = False
+
+_C.continue_ckpt = None
+_C.epoch_selected = "epoch=50.ckpt"
+
+_C.triplane_resolution = 128
+_C.triplane_channels_low = 128
+_C.triplane_channels_high = 512
+_C.lr = 1e-3
+_C.train = True
+_C.test = False
+
+_C.inference_save_pred_sdf_to_mesh=True
+_C.inference_save_feat_pca=True
+_C.name = "test"
+_C.test_subset = False
+_C.test_corres = False
+_C.test_partobjaversetiny = False
+
+_C.dataset = CN()
+_C.dataset.type = "Demo_Dataset"
+_C.dataset.data_path = "objaverse_data/"
+_C.dataset.train_num_workers = 64
+_C.dataset.val_num_workers = 32
+_C.dataset.train_batch_size = 2
+_C.dataset.val_batch_size = 2
+_C.dataset.all_files = []  # only used for correspondence demo
+
+_C.voxel2triplane = CN()
+_C.voxel2triplane.transformer_dim = 1024
+_C.voxel2triplane.transformer_layers = 6
+_C.voxel2triplane.transformer_heads = 8
+_C.voxel2triplane.triplane_low_res = 32
+_C.voxel2triplane.triplane_high_res = 256
+_C.voxel2triplane.triplane_dim = 64
+_C.voxel2triplane.normalize_vox_feat = False
+
+
+_C.loss = CN()
+_C.loss.triplet = 0.0
+_C.loss.sdf = 1.0
+_C.loss.feat = 10.0
+_C.loss.l1 = 0.0
+
+_C.use_pvcnn = False
+_C.use_pvcnnonly = True
+
+_C.pvcnn = CN()
+_C.pvcnn.point_encoder_type = 'pvcnn'
+_C.pvcnn.use_point_scatter = True
+_C.pvcnn.z_triplane_channels = 64
+_C.pvcnn.z_triplane_resolution = 256
+_C.pvcnn.unet_cfg = CN()
+_C.pvcnn.unet_cfg.depth = 3
+_C.pvcnn.unet_cfg.enabled = True
+_C.pvcnn.unet_cfg.rolled = True
+_C.pvcnn.unet_cfg.use_3d_aware = True
+_C.pvcnn.unet_cfg.start_hidden_channels = 32
+_C.pvcnn.unet_cfg.use_initial_conv = False
+
+_C.use_2d_feat = False
+_C.inference_metrics_only = False

third_party/PartField/partfield/dataloader.py

@@ -0,0 +1,78 @@
+import os
+import torch
+import trimesh
+import numpy as np
+import gc
+
+def quad_to_triangle_mesh(F):
+    """
+    Converts a quad-dominant mesh into a pure triangle mesh by splitting quads into two triangles.
+
+    Parameters:
+        quad_mesh (trimesh.Trimesh): Input mesh with quad faces.
+
+    Returns:
+        trimesh.Trimesh: A new mesh with only triangle faces.
+    """
+    faces = F
+
+    ### If already a triangle mesh -- skip
+    if len(faces[0]) == 3:
+        return F
+
+    new_faces = []
+
+    for face in faces:
+        if len(face) == 4:  # Quad face
+            # Split into two triangles
+            new_faces.append([face[0], face[1], face[2]])  # Triangle 1
+            new_faces.append([face[0], face[2], face[3]])  # Triangle 2
+        else:
+            print(f"Warning: Skipping non-triangle/non-quad face {face}")
+
+    new_faces = np.array(new_faces)
+
+    return new_faces
+
+class Demo_Dataset(torch.utils.data.Dataset):
+    def __init__(self, obj_path):
+        super().__init__()
+
+        self.obj_path = obj_path
+        self.pc_num_pts = 100000
+
+    
+    def __len__(self):
+        return 1
+
+    def get_model(self):
+        uid = os.path.basename(self.obj_path).split(".")[-2]
+        mesh = trimesh.load(self.obj_path, force='mesh', process=False)
+        vertices = mesh.vertices
+        faces = mesh.faces
+
+        bbmin = vertices.min(0)
+        bbmax = vertices.max(0)
+        center = (bbmin + bbmax) * 0.5
+        scale = 2.0 * 0.9 / (bbmax - bbmin).max()
+        vertices = (vertices - center) * scale
+        mesh.vertices = vertices
+
+        ### Make sure it is a triangle mesh -- just convert the quad
+        mesh.faces = quad_to_triangle_mesh(faces)
+        pc, _ = trimesh.sample.sample_surface(mesh, self.pc_num_pts) 
+
+        result = {
+                    'uid': uid
+                }
+
+        result['pc'] = torch.tensor(pc, dtype=torch.float32)
+        result['vertices'] = mesh.vertices
+        result['faces'] = mesh.faces
+
+        return result
+
+    def __getitem__(self, index):
+        gc.collect()
+
+        return self.get_model()

third_party/PartField/partfield/model/PVCNN/conv_pointnet.py

@@ -0,0 +1,251 @@
+"""
+Taken from gensdf
+https://github.com/princeton-computational-imaging/gensdf
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+# from dnnlib.util import printarr
+try:
+    from torch_scatter import scatter_mean, scatter_max
+except:
+    pass
+# from .unet import UNet
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# Resnet Blocks
+class ResnetBlockFC(nn.Module):
+    ''' Fully connected ResNet Block class.
+    Args:
+        size_in (int): input dimension
+        size_out (int): output dimension
+        size_h (int): hidden dimension
+    '''
+
+    def __init__(self, size_in, size_out=None, size_h=None):
+        super().__init__()
+        # Attributes
+        if size_out is None:
+            size_out = size_in
+
+        if size_h is None:
+            size_h = min(size_in, size_out)
+
+        self.size_in = size_in
+        self.size_h = size_h
+        self.size_out = size_out
+        # Submodules
+        self.fc_0 = nn.Linear(size_in, size_h)
+        self.fc_1 = nn.Linear(size_h, size_out)
+        self.actvn = nn.ReLU()
+
+        if size_in == size_out:
+            self.shortcut = None
+        else:
+            self.shortcut = nn.Linear(size_in, size_out, bias=False)
+        # Initialization
+        nn.init.zeros_(self.fc_1.weight)
+
+    def forward(self, x):
+        net = self.fc_0(self.actvn(x))
+        dx = self.fc_1(self.actvn(net))
+
+        if self.shortcut is not None:
+            x_s = self.shortcut(x)
+        else:
+            x_s = x
+
+        return x_s + dx
+
+
+class ConvPointnet(nn.Module):
+    ''' PointNet-based encoder network with ResNet blocks for each point.
+        Number of input points are fixed.
+    
+    Args:
+        c_dim (int): dimension of latent code c
+        dim (int): input points dimension
+        hidden_dim (int): hidden dimension of the network
+        scatter_type (str): feature aggregation when doing local pooling
+        unet (bool): weather to use U-Net
+        unet_kwargs (str): U-Net parameters
+        plane_resolution (int): defined resolution for plane feature
+        plane_type (str): feature type, 'xz' - 1-plane, ['xz', 'xy', 'yz'] - 3-plane, ['grid'] - 3D grid volume
+        padding (float): conventional padding paramter of ONet for unit cube, so [-0.5, 0.5] -> [-0.55, 0.55]
+        n_blocks (int): number of blocks ResNetBlockFC layers
+    '''
+
+    def __init__(self, c_dim=128, dim=3, hidden_dim=128, scatter_type='max', 
+                #  unet=False, unet_kwargs=None, 
+                 plane_resolution=None, plane_type=['xz', 'xy', 'yz'], padding=0.1, n_blocks=5):
+        super().__init__()
+        self.c_dim = c_dim
+
+        self.fc_pos = nn.Linear(dim, 2*hidden_dim)
+        self.blocks = nn.ModuleList([
+            ResnetBlockFC(2*hidden_dim, hidden_dim) for i in range(n_blocks)
+        ])
+        self.fc_c = nn.Linear(hidden_dim, c_dim)
+
+        self.actvn = nn.ReLU()
+        self.hidden_dim = hidden_dim
+
+        # if unet:
+        #     self.unet = UNet(c_dim, in_channels=c_dim, **unet_kwargs)
+        # else:
+        #     self.unet = None
+
+        self.reso_plane = plane_resolution
+        self.plane_type = plane_type
+        self.padding = padding
+
+        if scatter_type == 'max':
+            self.scatter = scatter_max
+        elif scatter_type == 'mean':
+            self.scatter = scatter_mean
+
+
+    # takes in "p": point cloud and "query": sdf_xyz 
+    # sample plane features for unlabeled_query as well 
+    def forward(self, p):#, query2):
+        batch_size, T, D = p.size()
+
+        # acquire the index for each point
+        coord = {}
+        index = {}
+        if 'xz' in self.plane_type:
+            coord['xz'] = self.normalize_coordinate(p.clone(), plane='xz', padding=self.padding)
+            index['xz'] = self.coordinate2index(coord['xz'], self.reso_plane)
+        if 'xy' in self.plane_type:
+            coord['xy'] = self.normalize_coordinate(p.clone(), plane='xy', padding=self.padding)
+            index['xy'] = self.coordinate2index(coord['xy'], self.reso_plane)
+        if 'yz' in self.plane_type:
+            coord['yz'] = self.normalize_coordinate(p.clone(), plane='yz', padding=self.padding)
+            index['yz'] = self.coordinate2index(coord['yz'], self.reso_plane)
+
+        
+        net = self.fc_pos(p)
+
+        net = self.blocks[0](net)
+        for block in self.blocks[1:]:
+            pooled = self.pool_local(coord, index, net)
+            net = torch.cat([net, pooled], dim=2)
+            net = block(net)
+
+        c = self.fc_c(net)
+
+        fea = {}
+        plane_feat_sum = 0
+        #second_sum = 0
+        if 'xz' in self.plane_type:
+            fea['xz'] = self.generate_plane_features(p, c, plane='xz') # shape: batch, latent size, resolution, resolution (e.g. 16, 256, 64, 64)
+            # plane_feat_sum += self.sample_plane_feature(query, fea['xz'], 'xz')
+            #second_sum += self.sample_plane_feature(query2, fea['xz'], 'xz')
+        if 'xy' in self.plane_type:
+            fea['xy'] = self.generate_plane_features(p, c, plane='xy')
+            # plane_feat_sum += self.sample_plane_feature(query, fea['xy'], 'xy')
+            #second_sum += self.sample_plane_feature(query2, fea['xy'], 'xy')
+        if 'yz' in self.plane_type:
+            fea['yz'] = self.generate_plane_features(p, c, plane='yz')
+            # plane_feat_sum += self.sample_plane_feature(query, fea['yz'], 'yz')
+            #second_sum += self.sample_plane_feature(query2, fea['yz'], 'yz')
+        return fea
+
+        # return plane_feat_sum.transpose(2,1)#, second_sum.transpose(2,1)
+
+
+    def normalize_coordinate(self, p, padding=0.1, plane='xz'):
+        ''' Normalize coordinate to [0, 1] for unit cube experiments
+
+        Args:
+            p (tensor): point
+            padding (float): conventional padding paramter of ONet for unit cube, so [-0.5, 0.5] -> [-0.55, 0.55]
+            plane (str): plane feature type, ['xz', 'xy', 'yz']
+        '''
+        if plane == 'xz':
+            xy = p[:, :, [0, 2]]
+        elif plane =='xy':
+            xy = p[:, :, [0, 1]]
+        else:
+            xy = p[:, :, [1, 2]]
+
+        xy_new = xy / (1 + padding + 10e-6) # (-0.5, 0.5)
+        xy_new = xy_new + 0.5 # range (0, 1)
+
+        # f there are outliers out of the range
+        if xy_new.max() >= 1:
+            xy_new[xy_new >= 1] = 1 - 10e-6
+        if xy_new.min() < 0:
+            xy_new[xy_new < 0] = 0.0
+        return xy_new
+
+
+    def coordinate2index(self, x, reso):
+        ''' Normalize coordinate to [0, 1] for unit cube experiments.
+            Corresponds to our 3D model
+
+        Args:
+            x (tensor): coordinate
+            reso (int): defined resolution
+            coord_type (str): coordinate type
+        '''
+        x = (x * reso).long()
+        index = x[:, :, 0] + reso * x[:, :, 1]
+        index = index[:, None, :]
+        return index
+
+
+    # xy is the normalized coordinates of the point cloud of each plane 
+    # I'm pretty sure the keys of xy are the same as those of index, so xy isn't needed here as input 
+    def pool_local(self, xy, index, c):
+        bs, fea_dim = c.size(0), c.size(2)
+        keys = xy.keys()
+
+        c_out = 0
+        for key in keys:
+            # scatter plane features from points
+            fea = self.scatter(c.permute(0, 2, 1), index[key], dim_size=self.reso_plane**2)
+            if self.scatter == scatter_max:
+                fea = fea[0]
+            # gather feature back to points
+            fea = fea.gather(dim=2, index=index[key].expand(-1, fea_dim, -1))
+            c_out += fea
+        return c_out.permute(0, 2, 1)
+
+
+    def generate_plane_features(self, p, c, plane='xz'):
+        # acquire indices of features in plane
+        xy = self.normalize_coordinate(p.clone(), plane=plane, padding=self.padding) # normalize to the range of (0, 1)
+        index = self.coordinate2index(xy, self.reso_plane)
+
+        # scatter plane features from points
+        fea_plane = c.new_zeros(p.size(0), self.c_dim, self.reso_plane**2)
+        c = c.permute(0, 2, 1) # B x 512 x T
+        fea_plane = scatter_mean(c, index, out=fea_plane) # B x 512 x reso^2
+        fea_plane = fea_plane.reshape(p.size(0), self.c_dim, self.reso_plane, self.reso_plane) # sparce matrix (B x 512 x reso x reso)
+        
+        # printarr(fea_plane, c, p, xy, index)
+        # import pdb; pdb.set_trace()
+
+        # process the plane features with UNet
+        # if self.unet is not None:
+        #     fea_plane = self.unet(fea_plane)
+
+        return fea_plane
+
+
+    # sample_plane_feature function copied from /src/conv_onet/models/decoder.py
+    # uses values from plane_feature and pixel locations from vgrid to interpolate feature
+    def sample_plane_feature(self, query, plane_feature, plane):
+        xy = self.normalize_coordinate(query.clone(), plane=plane, padding=self.padding)
+        xy = xy[:, :, None].float()
+        vgrid = 2.0 * xy - 1.0 # normalize to (-1, 1)
+        sampled_feat = F.grid_sample(plane_feature, vgrid, padding_mode='border', align_corners=True, mode='bilinear').squeeze(-1)
+        return sampled_feat
+
+
+    

third_party/PartField/partfield/model/PVCNN/dnnlib_util.py

@@ -0,0 +1,1074 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
+
+"""Miscellaneous utility classes and functions."""
+from collections import namedtuple
+import time
+import ctypes
+import fnmatch
+import importlib
+import inspect
+import numpy as np
+import json
+import os
+import shutil
+import sys
+import types
+import io
+import pickle
+import re
+# import requests
+import html
+import hashlib
+import glob
+import tempfile
+import urllib
+import urllib.request
+import uuid
+import boto3
+import threading
+from contextlib import ContextDecorator
+from contextlib import contextmanager, nullcontext
+
+from distutils.util import strtobool
+from typing import Any, List, Tuple, Union
+import importlib
+from loguru import logger
+# import wandb
+import torch
+import psutil
+import subprocess
+
+import random
+import string
+import pdb
+
+# Util classes
+# ------------------------------------------------------------------------------------------
+
+
+class EasyDict(dict):
+    """Convenience class that behaves like a dict but allows access with the attribute syntax."""
+
+    def __getattr__(self, name: str) -> Any:
+        try:
+            return self[name]
+        except KeyError:
+            raise AttributeError(name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        self[name] = value
+
+    def __delattr__(self, name: str) -> None:
+        del self[name]
+
+
+class Logger(object):
+    """Redirect stderr to stdout, optionally print stdout to a file, and optionally force flushing on both stdout and the file."""
+
+    def __init__(self, file_name: str = None, file_mode: str = "w", should_flush: bool = True):
+        self.file = None
+
+        if file_name is not None:
+            self.file = open(file_name, file_mode)
+
+        self.should_flush = should_flush
+        self.stdout = sys.stdout
+        self.stderr = sys.stderr
+
+        sys.stdout = self
+        sys.stderr = self
+
+    def __enter__(self) -> "Logger":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.close()
+
+    def write(self, text: Union[str, bytes]) -> None:
+        """Write text to stdout (and a file) and optionally flush."""
+        if isinstance(text, bytes):
+            text = text.decode()
+        if len(text) == 0:  # workaround for a bug in VSCode debugger: sys.stdout.write(''); sys.stdout.flush() => crash
+            return
+
+        if self.file is not None:
+            self.file.write(text)
+
+        self.stdout.write(text)
+
+        if self.should_flush:
+            self.flush()
+
+    def flush(self) -> None:
+        """Flush written text to both stdout and a file, if open."""
+        if self.file is not None:
+            self.file.flush()
+
+        self.stdout.flush()
+
+    def close(self) -> None:
+        """Flush, close possible files, and remove stdout/stderr mirroring."""
+        self.flush()
+
+        # if using multiple loggers, prevent closing in wrong order
+        if sys.stdout is self:
+            sys.stdout = self.stdout
+        if sys.stderr is self:
+            sys.stderr = self.stderr
+
+        if self.file is not None:
+            self.file.close()
+            self.file = None
+
+
+# Cache directories
+# ------------------------------------------------------------------------------------------
+
+_dnnlib_cache_dir = None
+
+
+def set_cache_dir(path: str) -> None:
+    global _dnnlib_cache_dir
+    _dnnlib_cache_dir = path
+
+
+def make_cache_dir_path(*paths: str) -> str:
+    if _dnnlib_cache_dir is not None:
+        return os.path.join(_dnnlib_cache_dir, *paths)
+    if 'DNNLIB_CACHE_DIR' in os.environ:
+        return os.path.join(os.environ['DNNLIB_CACHE_DIR'], *paths)
+    if 'HOME' in os.environ:
+        return os.path.join(os.environ['HOME'], '.cache', 'dnnlib', *paths)
+    if 'USERPROFILE' in os.environ:
+        return os.path.join(os.environ['USERPROFILE'], '.cache', 'dnnlib', *paths)
+    return os.path.join(tempfile.gettempdir(), '.cache', 'dnnlib', *paths)
+
+
+# Small util functions
+# ------------------------------------------------------------------------------------------
+
+
+def format_time(seconds: Union[int, float]) -> str:
+    """Convert the seconds to human readable string with days, hours, minutes and seconds."""
+    s = int(np.rint(seconds))
+
+    if s < 60:
+        return "{0}s".format(s)
+    elif s < 60 * 60:
+        return "{0}m {1:02}s".format(s // 60, s % 60)
+    elif s < 24 * 60 * 60:
+        return "{0}h {1:02}m {2:02}s".format(s // (60 * 60), (s // 60) % 60, s % 60)
+    else:
+        return "{0}d {1:02}h {2:02}m".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24, (s // 60) % 60)
+
+
+def format_time_brief(seconds: Union[int, float]) -> str:
+    """Convert the seconds to human readable string with days, hours, minutes and seconds."""
+    s = int(np.rint(seconds))
+
+    if s < 60:
+        return "{0}s".format(s)
+    elif s < 60 * 60:
+        return "{0}m {1:02}s".format(s // 60, s % 60)
+    elif s < 24 * 60 * 60:
+        return "{0}h {1:02}m".format(s // (60 * 60), (s // 60) % 60)
+    else:
+        return "{0}d {1:02}h".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24)
+
+
+def ask_yes_no(question: str) -> bool:
+    """Ask the user the question until the user inputs a valid answer."""
+    while True:
+        try:
+            print("{0} [y/n]".format(question))
+            return strtobool(input().lower())
+        except ValueError:
+            pass
+
+
+def tuple_product(t: Tuple) -> Any:
+    """Calculate the product of the tuple elements."""
+    result = 1
+
+    for v in t:
+        result *= v
+
+    return result
+
+
+_str_to_ctype = {
+    "uint8": ctypes.c_ubyte,
+    "uint16": ctypes.c_uint16,
+    "uint32": ctypes.c_uint32,
+    "uint64": ctypes.c_uint64,
+    "int8": ctypes.c_byte,
+    "int16": ctypes.c_int16,
+    "int32": ctypes.c_int32,
+    "int64": ctypes.c_int64,
+    "float32": ctypes.c_float,
+    "float64": ctypes.c_double
+}
+
+
+def get_dtype_and_ctype(type_obj: Any) -> Tuple[np.dtype, Any]:
+    """Given a type name string (or an object having a __name__ attribute), return matching Numpy and ctypes types that have the same size in bytes."""
+    type_str = None
+
+    if isinstance(type_obj, str):
+        type_str = type_obj
+    elif hasattr(type_obj, "__name__"):
+        type_str = type_obj.__name__
+    elif hasattr(type_obj, "name"):
+        type_str = type_obj.name
+    else:
+        raise RuntimeError("Cannot infer type name from input")
+
+    assert type_str in _str_to_ctype.keys()
+
+    my_dtype = np.dtype(type_str)
+    my_ctype = _str_to_ctype[type_str]
+
+    assert my_dtype.itemsize == ctypes.sizeof(my_ctype)
+
+    return my_dtype, my_ctype
+
+
+def is_pickleable(obj: Any) -> bool:
+    try:
+        with io.BytesIO() as stream:
+            pickle.dump(obj, stream)
+        return True
+    except:
+        return False
+
+
+# Functionality to import modules/objects by name, and call functions by name
+# ------------------------------------------------------------------------------------------
+
+def get_module_from_obj_name(obj_name: str) -> Tuple[types.ModuleType, str]:
+    """Searches for the underlying module behind the name to some python object.
+    Returns the module and the object name (original name with module part removed)."""
+
+    # allow convenience shorthands, substitute them by full names
+    obj_name = re.sub("^np.", "numpy.", obj_name)
+    obj_name = re.sub("^tf.", "tensorflow.", obj_name)
+
+    # list alternatives for (module_name, local_obj_name)
+    parts = obj_name.split(".")
+    name_pairs = [(".".join(parts[:i]), ".".join(parts[i:])) for i in range(len(parts), 0, -1)]
+
+    # try each alternative in turn
+    for module_name, local_obj_name in name_pairs:
+        try:
+            module = importlib.import_module(module_name)  # may raise ImportError
+            get_obj_from_module(module, local_obj_name)  # may raise AttributeError
+            return module, local_obj_name
+        except:
+            pass
+
+    # maybe some of the modules themselves contain errors?
+    for module_name, _local_obj_name in name_pairs:
+        try:
+            importlib.import_module(module_name)  # may raise ImportError
+        except ImportError:
+            if not str(sys.exc_info()[1]).startswith("No module named '" + module_name + "'"):
+                raise
+
+    # maybe the requested attribute is missing?
+    for module_name, local_obj_name in name_pairs:
+        try:
+            module = importlib.import_module(module_name)  # may raise ImportError
+            get_obj_from_module(module, local_obj_name)  # may raise AttributeError
+        except ImportError:
+            pass
+
+    # we are out of luck, but we have no idea why
+    raise ImportError(obj_name)
+
+
+def get_obj_from_module(module: types.ModuleType, obj_name: str) -> Any:
+    """Traverses the object name and returns the last (rightmost) python object."""
+    if obj_name == '':
+        return module
+    obj = module
+    for part in obj_name.split("."):
+        obj = getattr(obj, part)
+    return obj
+
+
+def get_obj_by_name(name: str) -> Any:
+    """Finds the python object with the given name."""
+    module, obj_name = get_module_from_obj_name(name)
+    return get_obj_from_module(module, obj_name)
+
+
+def call_func_by_name(*args, func_name: str = None, **kwargs) -> Any:
+    """Finds the python object with the given name and calls it as a function."""
+    assert func_name is not None
+    func_obj = get_obj_by_name(func_name)
+    assert callable(func_obj)
+    return func_obj(*args, **kwargs)
+
+
+def construct_class_by_name(*args, class_name: str = None, **kwargs) -> Any:
+    """Finds the python class with the given name and constructs it with the given arguments."""
+    return call_func_by_name(*args, func_name=class_name, **kwargs)
+
+
+def get_module_dir_by_obj_name(obj_name: str) -> str:
+    """Get the directory path of the module containing the given object name."""
+    module, _ = get_module_from_obj_name(obj_name)
+    return os.path.dirname(inspect.getfile(module))
+
+
+def is_top_level_function(obj: Any) -> bool:
+    """Determine whether the given object is a top-level function, i.e., defined at module scope using 'def'."""
+    return callable(obj) and obj.__name__ in sys.modules[obj.__module__].__dict__
+
+
+def get_top_level_function_name(obj: Any) -> str:
+    """Return the fully-qualified name of a top-level function."""
+    assert is_top_level_function(obj)
+    module = obj.__module__
+    if module == '__main__':
+        module = os.path.splitext(os.path.basename(sys.modules[module].__file__))[0]
+    return module + "." + obj.__name__
+
+
+# File system helpers
+# ------------------------------------------------------------------------------------------
+
+def list_dir_recursively_with_ignore(dir_path: str, ignores: List[str] = None, add_base_to_relative: bool = False) -> List[Tuple[str, str]]:
+    """List all files recursively in a given directory while ignoring given file and directory names.
+    Returns list of tuples containing both absolute and relative paths."""
+    assert os.path.isdir(dir_path)
+    base_name = os.path.basename(os.path.normpath(dir_path))
+
+    if ignores is None:
+        ignores = []
+
+    result = []
+
+    for root, dirs, files in os.walk(dir_path, topdown=True):
+        for ignore_ in ignores:
+            dirs_to_remove = [d for d in dirs if fnmatch.fnmatch(d, ignore_)]
+
+            # dirs need to be edited in-place
+            for d in dirs_to_remove:
+                dirs.remove(d)
+
+            files = [f for f in files if not fnmatch.fnmatch(f, ignore_)]
+
+        absolute_paths = [os.path.join(root, f) for f in files]
+        relative_paths = [os.path.relpath(p, dir_path) for p in absolute_paths]
+
+        if add_base_to_relative:
+            relative_paths = [os.path.join(base_name, p) for p in relative_paths]
+
+        assert len(absolute_paths) == len(relative_paths)
+        result += zip(absolute_paths, relative_paths)
+
+    return result
+
+
+def copy_files_and_create_dirs(files: List[Tuple[str, str]]) -> None:
+    """Takes in a list of tuples of (src, dst) paths and copies files.
+    Will create all necessary directories."""
+    for file in files:
+        target_dir_name = os.path.dirname(file[1])
+
+        # will create all intermediate-level directories
+        if not os.path.exists(target_dir_name):
+            os.makedirs(target_dir_name)
+
+        shutil.copyfile(file[0], file[1])
+
+
+# URL helpers
+# ------------------------------------------------------------------------------------------
+
+def is_url(obj: Any, allow_file_urls: bool = False) -> bool:
+    """Determine whether the given object is a valid URL string."""
+    if not isinstance(obj, str) or not "://" in obj:
+        return False
+    if allow_file_urls and obj.startswith('file://'):
+        return True
+    try:
+        res = requests.compat.urlparse(obj)
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+        res = requests.compat.urlparse(requests.compat.urljoin(obj, "/"))
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+    except:
+        return False
+    return True
+
+
+def open_url(url: str, cache_dir: str = None, num_attempts: int = 10, verbose: bool = True, return_filename: bool = False, cache: bool = True) -> Any:
+    """Download the given URL and return a binary-mode file object to access the data."""
+    assert num_attempts >= 1
+    assert not (return_filename and (not cache))
+
+    # Doesn't look like an URL scheme so interpret it as a local filename.
+    if not re.match('^[a-z]+://', url):
+        return url if return_filename else open(url, "rb")
+
+    # Handle file URLs.  This code handles unusual file:// patterns that
+    # arise on Windows:
+    #
+    # file:///c:/foo.txt
+    #
+    # which would translate to a local '/c:/foo.txt' filename that's
+    # invalid.  Drop the forward slash for such pathnames.
+    #
+    # If you touch this code path, you should test it on both Linux and
+    # Windows.
+    #
+    # Some internet resources suggest using urllib.request.url2pathname() but
+    # but that converts forward slashes to backslashes and this causes
+    # its own set of problems.
+    if url.startswith('file://'):
+        filename = urllib.parse.urlparse(url).path
+        if re.match(r'^/[a-zA-Z]:', filename):
+            filename = filename[1:]
+        return filename if return_filename else open(filename, "rb")
+
+    assert is_url(url)
+
+    # Lookup from cache.
+    if cache_dir is None:
+        cache_dir = make_cache_dir_path('downloads')
+
+    url_md5 = hashlib.md5(url.encode("utf-8")).hexdigest()
+    if cache:
+        cache_files = glob.glob(os.path.join(cache_dir, url_md5 + "_*"))
+        if len(cache_files) == 1:
+            filename = cache_files[0]
+            return filename if return_filename else open(filename, "rb")
+
+    # Download.
+    url_name = None
+    url_data = None
+    with requests.Session() as session:
+        if verbose:
+            print("Downloading %s ..." % url, end="", flush=True)
+        for attempts_left in reversed(range(num_attempts)):
+            try:
+                with session.get(url) as res:
+                    res.raise_for_status()
+                    if len(res.content) == 0:
+                        raise IOError("No data received")
+
+                    if len(res.content) < 8192:
+                        content_str = res.content.decode("utf-8")
+                        if "download_warning" in res.headers.get("Set-Cookie", ""):
+                            links = [html.unescape(link) for link in content_str.split('"') if "export=download" in link]
+                            if len(links) == 1:
+                                url = requests.compat.urljoin(url, links[0])
+                                raise IOError("Google Drive virus checker nag")
+                        if "Google Drive - Quota exceeded" in content_str:
+                            raise IOError("Google Drive download quota exceeded -- please try again later")
+
+                    match = re.search(r'filename="([^"]*)"', res.headers.get("Content-Disposition", ""))
+                    url_name = match[1] if match else url
+                    url_data = res.content
+                    if verbose:
+                        print(" done")
+                    break
+            except KeyboardInterrupt:
+                raise
+            except:
+                if not attempts_left:
+                    if verbose:
+                        print(" failed")
+                    raise
+                if verbose:
+                    print(".", end="", flush=True)
+
+    # Save to cache.
+    if cache:
+        safe_name = re.sub(r"[^0-9a-zA-Z-._]", "_", url_name)
+        cache_file = os.path.join(cache_dir, url_md5 + "_" + safe_name)
+        temp_file = os.path.join(cache_dir, "tmp_" + uuid.uuid4().hex + "_" + url_md5 + "_" + safe_name)
+        os.makedirs(cache_dir, exist_ok=True)
+        with open(temp_file, "wb") as f:
+            f.write(url_data)
+        os.replace(temp_file, cache_file)  # atomic
+        if return_filename:
+            return cache_file
+
+    # Return data as file object.
+    assert not return_filename
+    return io.BytesIO(url_data)
+
+# ------------------------------------------------------------------------------------------
+# util function modified from https://github.com/nv-tlabs/LION/blob/0467d2199076e95a7e88bafd99dcd7d48a04b4a7/utils/model_helper.py
+def import_class(model_str):
+    from torch_utils.dist_utils import is_rank0
+    if is_rank0():
+        logger.info('import: {}', model_str)
+    p, m = model_str.rsplit('.', 1)
+    mod = importlib.import_module(p)
+    Model = getattr(mod, m)
+    return Model
+
+class ScopedTorchProfiler(ContextDecorator):
+    """
+    Marks ranges for both nvtx profiling (with nsys) and torch autograd profiler
+    """
+    __global_counts = {}
+    enabled=False
+
+    def __init__(self, unique_name: str):
+        """
+        Names must be unique!
+        """
+        ScopedTorchProfiler.__global_counts[unique_name] = 0
+        self._name = unique_name
+        self._autograd_scope = torch.profiler.record_function(unique_name)
+    
+    def __enter__(self):
+        if ScopedTorchProfiler.enabled:
+            torch.cuda.nvtx.range_push(self._name)
+        self._autograd_scope.__enter__()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._autograd_scope.__exit__(exc_type, exc_value, traceback)
+        if ScopedTorchProfiler.enabled:
+            torch.cuda.nvtx.range_pop()
+
+class TimingsMonitor():
+    CUDATimer = namedtuple('CUDATimer', ['start', 'end'])
+    def __init__(self, device, enabled=True, timing_names:List[str]=[], cuda_timing_names:List[str]=[]):
+        """
+        Usage:
+            tmonitor = TimingsMonitor(device)
+            for i in range(n_iter):
+                # Record arbitrary scopes
+                with tmonitor.timing_scope('regular_scope_name'):
+                    ...
+                    with tmonitor.cuda_timing_scope('nested_scope_name'):
+                        ...
+                with tmonitor.cuda_timing_scope('cuda_scope_name'):
+                    ...
+                tmonitor.record_timing('duration_name', end_time - start_time)
+
+                # Gather timings
+                tmonitor.record_all_cuda_timings()
+                tmonitor.update_all_averages()
+                averages = tmonitor.get_average_timings()
+                all_timings = tmonitor.get_timings()
+
+        Two types of timers, standard report timing and cuda timings.
+        Cuda timing supports scoped context manager cuda_event_scope.
+        Args:
+            device: device to time on (needed for cuda timers)
+            # enabled: HACK to only report timings from rank 0, set enabled=(global_rank==0)
+            timing_names: timings to report optional (will auto add new names)
+            cuda_timing_names: cuda periods to time optional (will auto add new names)
+        """
+        self.enabled=enabled
+        self.device = device
+
+        # Normal timing
+        # self.all_timings_dict = {k:None for k in timing_names + cuda_timing_names}
+        self.all_timings_dict = {}
+        self.avg_meter_dict = {}
+
+        # Cuda event timers to measure time spent on pushing data to gpu and on training step
+        self.cuda_event_timers = {}
+
+        for k in timing_names:
+            self.add_new_timing(k)
+
+        for k in cuda_timing_names:
+            self.add_new_cuda_timing(k)
+
+        # Running averages
+        # self.avg_meter_dict = {k:AverageMeter() for k in self.all_timings_dict}
+    
+    def add_new_timing(self, name):
+        self.avg_meter_dict[name] = AverageMeter()
+        self.all_timings_dict[name] = None
+    
+    def add_new_cuda_timing(self, name):
+        start_event = torch.cuda.Event(enable_timing=True)
+        end_event = torch.cuda.Event(enable_timing=True)
+        self.cuda_event_timers[name] = self.CUDATimer(start=start_event, end=end_event)
+        self.add_new_timing(name)
+    
+    def clear_timings(self):
+        self.all_timings_dict = {k:None for k in self.all_timings_dict}
+
+    def get_timings(self):
+        return self.all_timings_dict
+
+    def get_average_timings(self):
+        return {k:v.avg for k,v in self.avg_meter_dict.items()}
+
+    def update_all_averages(self):
+        """
+        Once per iter, when timings have been finished recording, one should
+        call update_average_iter to keep running average of timings.
+        """
+        for k,v in self.all_timings_dict.items():
+            if v is None:
+                print("none_timing", k)
+                continue
+            self.avg_meter_dict[k].update(v)
+    
+    def record_timing(self, name, value):
+        if name not in self.all_timings_dict: self.add_new_timing(name)
+        # assert name in self.all_timings_dict
+        self.all_timings_dict[name] = value
+
+    def _record_cuda_event_start(self, name):
+        if name in self.cuda_event_timers:
+            self.cuda_event_timers[name].start.record(
+                torch.cuda.current_stream(self.device))
+
+    def _record_cuda_event_end(self, name):
+        if name in self.cuda_event_timers:
+            self.cuda_event_timers[name].end.record(
+                torch.cuda.current_stream(self.device))
+
+    @contextmanager
+    def cuda_timing_scope(self, name, profile=True):
+        if name not in self.all_timings_dict: self.add_new_cuda_timing(name)
+        with ScopedTorchProfiler(name) if profile else nullcontext():
+            self._record_cuda_event_start(name)
+            try:
+                yield
+            finally:
+                self._record_cuda_event_end(name)
+
+    @contextmanager
+    def timing_scope(self, name, profile=True):
+        if name not in self.all_timings_dict: self.add_new_timing(name)
+        with ScopedTorchProfiler(name) if profile else nullcontext():
+            start_time = time.time()
+            try:
+                yield
+            finally:
+                self.record_timing(name, time.time()-start_time)
+    
+    def record_all_cuda_timings(self):
+        """ After all the cuda events call this to synchronize and record down the cuda timings. """
+        for k, events in self.cuda_event_timers.items():
+            with torch.no_grad():
+                events.end.synchronize()
+                # Convert to seconds
+                time_elapsed = events.start.elapsed_time(events.end)/1000.
+                self.all_timings_dict[k] = time_elapsed
+
+def init_s3(config_file):
+    config = json.load(open(config_file, 'r'))
+    s3_client = boto3.client("s3", **config)
+    return s3_client
+
+def download_from_s3(file_path, target_path, cfg):
+    tic = time.time()
+    s3_client = init_s3(cfg.checkpoint.write_s3_config)  # use to test the s3_client can be init
+    bucket_name = file_path.split('/')[2]
+    file_key = file_path.split(bucket_name+'/')[-1]
+    print(bucket_name, file_key)
+    s3_client.download_file(bucket_name, file_key, target_path)
+    logger.info(f'finish download from ! s3://{bucket_name}/{file_key} to {target_path} %.1f sec'%(
+        time.time() - tic))
+
+def upload_to_s3(buffer, bucket_name, key, config_dict):
+    logger.info(f'start upload_to_s3! bucket_name={bucket_name}, key={key}')
+    tic = time.time()
+    s3 = boto3.client('s3', **config_dict)
+    s3.put_object(Bucket=bucket_name, Key=key, Body=buffer.getvalue())
+    logger.info(f'finish upload_to_s3! s3://{bucket_name}/{key} %.1f sec'%(time.time() - tic))
+
+def write_ckpt_to_s3(cfg, all_model_dict, ckpt_name):
+    buffer = io.BytesIO()
+    tic = time.time()
+    torch.save(all_model_dict, buffer)  # take ~0.25 sec
+    # logger.info('write ckpt to buffer: %.2f sec'%(time.time() - tic))
+    group, name = cfg.outdir.rstrip("/").split("/")[-2:]
+    key = f"checkpoints/{group}/{name}/ckpt/{ckpt_name}"
+    bucket_name = cfg.checkpoint.write_s3_bucket
+     
+    s3_client = init_s3(cfg.checkpoint.write_s3_config)  # use to test the s3_client can be init
+
+    config_dict = json.load(open(cfg.checkpoint.write_s3_config, 'r'))
+    upload_thread = threading.Thread(target=upload_to_s3, args=(buffer, bucket_name, key, config_dict))
+    upload_thread.start()
+    path = f"s3://{bucket_name}/{key}" 
+    return path
+
+def upload_file_to_s3(cfg, file_path, key_name=None):
+    # file_path is the local file path, can be a yaml file
+    # this function is used to upload the ckecpoint only
+    tic = time.time()
+    group, name = cfg.outdir.rstrip("/").split("/")[-2:]
+    if key_name is None:
+        key = os.path.basename(file_path)
+    key = f"checkpoints/{group}/{name}/{key}"
+    bucket_name = cfg.checkpoint.write_s3_bucket
+    s3_client = init_s3(cfg.checkpoint.write_s3_config)
+    # Upload the file
+    with open(file_path, 'rb') as f:
+        s3_client.upload_fileobj(f, bucket_name, key)
+    full_s3_path = f"s3://{bucket_name}/{key}"
+    logger.info(f'upload_to_s3: {file_path} {full_s3_path} | use time: {time.time()-tic}')
+
+    return full_s3_path
+
+
+def load_from_s3(file_path, cfg, load_fn):
+    """
+        ckpt_path example:
+            s3://xzeng/checkpoints/2023_0413/vae_kl_5e-1/ckpt/snapshot_epo000163_iter164000.pt
+    """
+    s3_client = init_s3(cfg.checkpoint.write_s3_config)  # use to test the s3_client can be init
+    bucket_name = file_path.split("s3://")[-1].split('/')[0]
+    key = file_path.split(f'{bucket_name}/')[-1]
+    # logger.info(f"-> try to load s3://{bucket_name}/{key} ")
+    tic = time.time()
+    for attemp in range(10):
+        try:
+            # Download the state dict from S3 into memory (as a binary stream)
+            with io.BytesIO() as buffer:
+                s3_client.download_fileobj(bucket_name, key, buffer)
+                buffer.seek(0)
+
+                # Load the state dict into a PyTorch model
+                # out = torch.load(buffer, map_location=torch.device("cpu"))
+                out = load_fn(buffer)
+            break
+        except:
+            logger.info(f"fail to load s3://{bucket_name}/{key} attemp: {attemp}")
+    from torch_utils.dist_utils import is_rank0
+    if is_rank0():
+        logger.info(f'loaded {file_path} | use time: {time.time()-tic:.1f} sec')
+    return out
+
+def load_torch_dict_from_s3(ckpt_path, cfg):
+    """
+        ckpt_path example:
+            s3://xzeng/checkpoints/2023_0413/vae_kl_5e-1/ckpt/snapshot_epo000163_iter164000.pt
+    """
+    s3_client = init_s3(cfg.checkpoint.write_s3_config)  # use to test the s3_client can be init
+    bucket_name = ckpt_path.split("s3://")[-1].split('/')[0]
+    key = ckpt_path.split(f'{bucket_name}/')[-1]
+    for attemp in range(10):
+        try:
+            # Download the state dict from S3 into memory (as a binary stream)
+            with io.BytesIO() as buffer:
+                s3_client.download_fileobj(bucket_name, key, buffer)
+                buffer.seek(0)
+
+                # Load the state dict into a PyTorch model
+                out = torch.load(buffer, map_location=torch.device("cpu"))
+            break
+        except:
+            logger.info(f"fail to load s3://{bucket_name}/{key} attemp: {attemp}")
+    return out
+
+def count_parameters_in_M(model):
+    return np.sum(np.prod(v.size()) for name, v in model.named_parameters() if "auxiliary" not in name) / 1e6
+
+def printarr(*arrs, float_width=6, **kwargs):
+    """
+    Print a pretty table giving name, shape, dtype, type, and content information for input tensors or scalars.
+
+    Call like: printarr(my_arr, some_other_arr, maybe_a_scalar). Accepts a variable number of arguments.
+
+    Inputs can be:
+        - Numpy tensor arrays
+        - Pytorch tensor arrays
+        - Jax tensor arrays
+        - Python ints / floats
+        - None
+
+    It may also work with other array-like types, but they have not been tested.
+
+    Use the `float_width` option specify the precision to which floating point types are printed.
+
+    Author: Nicholas Sharp (nmwsharp.com)
+    Canonical source: https://gist.github.com/nmwsharp/54d04af87872a4988809f128e1a1d233
+    License: This snippet may be used under an MIT license, and it is also released into the public domain. 
+             Please retain this docstring as a reference.
+    """
+    
+    frame = inspect.currentframe().f_back
+    default_name = "[temporary]"
+
+    ## helpers to gather data about each array
+    def name_from_outer_scope(a):
+        if a is None:
+            return '[None]'
+        name = default_name
+        for k, v in frame.f_locals.items():
+            if v is a:
+                name = k
+                break
+        return name
+
+    def type_strip(type_str):
+        return type_str.lstrip('<class ').rstrip('>').replace('torch.', '').strip("'")
+
+    def dtype_str(a):
+        if a is None:
+            return 'None'
+        if isinstance(a, int):
+            return 'int'
+        if isinstance(a, float):
+            return 'float'
+        if isinstance(a, list) and len(a)>0:
+            return type_strip(str(type(a[0])))
+        if hasattr(a, 'dtype'):
+            return type_strip(str(a.dtype))
+        else:
+            return ''
+    def shape_str(a):
+        if a is None:
+            return 'N/A'
+        if isinstance(a, int):
+            return 'scalar'
+        if isinstance(a, float):
+            return 'scalar'
+        if isinstance(a, list):
+            return f"[{shape_str(a[0]) if len(a)>0 else '?'}]*{len(a)}"
+        if hasattr(a, 'shape'):
+            return str(tuple(a.shape))
+        else:
+            return ''
+    def type_str(a):
+        return type_strip(str(type(a))) # TODO this is is weird... what's the better way?
+    def device_str(a):
+        if hasattr(a, 'device'):
+            device_str = str(a.device)
+            if len(device_str) < 10:
+                # heuristic: jax returns some goofy long string we don't want, ignore it
+                return device_str
+        return ""
+    def format_float(x):
+        return f"{x:{float_width}g}"
+    def minmaxmean_str(a):
+        if a is None:
+            return ('N/A', 'N/A', 'N/A', 'N/A')
+        if isinstance(a, int) or isinstance(a, float): 
+            return (format_float(a),)*4
+
+        # compute min/max/mean. if anything goes wrong, just print 'N/A'
+        min_str = "N/A"
+        try: min_str = format_float(a.min())
+        except: pass
+        max_str = "N/A"
+        try: max_str = format_float(a.max())
+        except: pass
+        mean_str = "N/A"
+        try: mean_str = format_float(a.mean())
+        except: pass
+        try: median_str = format_float(a.median())
+        except:
+            try: median_str = format_float(np.median(np.array(a)))
+            except: median_str = 'N/A'
+        return (min_str, max_str, mean_str, median_str)
+
+    def get_prop_dict(a,k=None):
+        minmaxmean = minmaxmean_str(a)
+        props = {
+            'name' : name_from_outer_scope(a) if k is None else k,
+            # 'type' : str(type(a)).replace('torch.',''),
+            'dtype' : dtype_str(a),
+            'shape' : shape_str(a),
+            'type' : type_str(a),
+            'device' : device_str(a),
+            'min' : minmaxmean[0],
+            'max' : minmaxmean[1],
+            'mean' : minmaxmean[2],
+            'median': minmaxmean[3]
+        }
+        return props
+
+    try:
+
+        props = ['name', 'type', 'dtype', 'shape', 'device', 'min', 'max', 'mean', 'median']
+
+        # precompute all of the properties for each input
+        str_props = []
+        for a in arrs:
+            str_props.append(get_prop_dict(a))
+        for k,a in kwargs.items():
+            str_props.append(get_prop_dict(a, k=k))
+
+        # for each property, compute its length
+        maxlen = {}
+        for p in props: maxlen[p] = 0
+        for sp in str_props:
+            for p in props:
+                maxlen[p] = max(maxlen[p], len(sp[p]))
+
+        # if any property got all empty strings, don't bother printing it, remove if from the list
+        props = [p for p in props if maxlen[p] > 0]
+
+        # print a header
+        header_str = ""
+        for p in props:
+            prefix =  "" if p == 'name' else " | "
+            fmt_key = ">" if p == 'name' else "<"
+            header_str += f"{prefix}{p:{fmt_key}{maxlen[p]}}"
+        print(header_str)
+        print("-"*len(header_str))
+            
+        # now print the acual arrays
+        for strp in str_props:
+            for p in props:
+                prefix =  "" if p == 'name' else " | "
+                fmt_key = ">" if p == 'name' else "<"
+                print(f"{prefix}{strp[p]:{fmt_key}{maxlen[p]}}", end='')
+            print("")
+
+    finally:
+        del frame
+
+def debug_print_all_tensor_sizes(min_tot_size = 0):
+    import gc
+    print("---------------------------------------"*3)
+    for obj in gc.get_objects():
+        try:
+            if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
+                if np.prod(obj.size())>=min_tot_size:
+                    print(type(obj), obj.size())
+        except:
+            pass
+def print_cpu_usage():
+    
+    # Get current CPU usage as a percentage
+    cpu_usage = psutil.cpu_percent()
+    
+    # Get current memory usage
+    memory_usage = psutil.virtual_memory().used
+    
+    # Convert memory usage to a human-readable format
+    memory_usage_str = psutil._common.bytes2human(memory_usage)
+    
+    # Print CPU and memory usage
+    msg = f"Current CPU usage: {cpu_usage}% | "
+    msg += f"Current memory usage: {memory_usage_str}"
+    return msg
+
+def calmsize(num_bytes):
+    if math.isnan(num_bytes):
+        return ''
+    for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
+        if abs(num_bytes) < 1024.0:
+            return "{:.1f}{}B".format(num_bytes, unit)
+        num_bytes /= 1024.0
+    return "{:.1f}{}B".format(num_bytes, 'Y')
+
+def readable_size(num_bytes: int) -> str: 
+    return calmsize(num_bytes) ## '' if math.isnan(num_bytes) else '{:.1f}'.format(calmsize(num_bytes))
+
+def get_gpu_memory():
+    """
+    Get the current GPU memory usage for each device as a dictionary
+    """
+    output = subprocess.check_output(["nvidia-smi", "--query-gpu=memory.used", "--format=csv"])
+    output = output.decode("utf-8")
+    gpu_memory_values = output.split("\n")[1:-1]
+    gpu_memory_values = [int(x.strip().split()[0]) for x in gpu_memory_values]
+    gpu_memory = dict(zip(range(len(gpu_memory_values)), gpu_memory_values))
+    return gpu_memory
+
+def get_gpu_util():
+    """
+    Get the current GPU memory usage for each device as a dictionary
+    """
+    output = subprocess.check_output(["nvidia-smi", "--query-gpu=utilization.gpu", "--format=csv"])
+    output = output.decode("utf-8")
+    gpu_memory_values = output.split("\n")[1:-1]
+    gpu_memory_values = [int(x.strip().split()[0]) for x in gpu_memory_values]
+    gpu_util = dict(zip(range(len(gpu_memory_values)), gpu_memory_values))
+    return gpu_util
+
+
+def print_gpu_usage():
+    useage = get_gpu_memory()
+    msg = f" | GPU usage: "
+    for k, v in useage.items():
+        msg += f"{k}: {v} MB "
+    # utilization = get_gpu_util()
+    # msg + ' | util '
+    # for k, v in utilization.items():
+    #     msg += f"{k}: {v} % "
+    return msg
+
+class AverageMeter(object):
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.avg = 0
+        self.sum = 0
+        self.cnt = 0
+
+    def update(self, val, n=1):
+        self.sum += val * n
+        self.cnt += n
+        self.avg = self.sum / self.cnt
+
+
+def generate_random_string(length):
+    # This script will generate a string of 10 random ASCII letters (both lowercase and uppercase).
+    # You can adjust the length parameter to fit your needs.
+    letters = string.ascii_letters
+    return ''.join(random.choice(letters) for _ in range(length))
+
+
+class ForkedPdb(pdb.Pdb):
+    """
+    PDB Subclass for debugging multi-processed code
+    Suggested in: https://stackoverflow.com/questions/4716533/how-to-attach-debugger-to-a-python-subproccess
+    """
+    def interaction(self, *args, **kwargs):
+        _stdin = sys.stdin
+        try:
+            sys.stdin = open('/dev/stdin')
+            pdb.Pdb.interaction(self, *args, **kwargs)
+        finally:
+            sys.stdin = _stdin
+
+def check_exist_in_s3(file_path, s3_config):
+    s3 = init_s3(s3_config)
+    bucket_name, object_name = s3path_to_bucket_key(file_path)
+
+    try:
+        s3.head_object(Bucket=bucket_name, Key=object_name)
+        return 1
+    except:
+        logger.info(f'file not found: s3://{bucket_name}/{object_name}')
+        return 0
+
+def s3path_to_bucket_key(file_path):
+    bucket_name = file_path.split('/')[2]
+    object_name = file_path.split(bucket_name + '/')[-1]
+    return bucket_name, object_name
+
+def copy_file_to_s3(cfg, file_path_local, file_path_s3):
+    # work similar as upload_file_to_s3, but not trying to parse the file path
+    # file_path_s3: s3://{bucket}/{key}
+    bucket_name, key = s3path_to_bucket_key(file_path_s3)
+    tic = time.time()
+    s3_client = init_s3(cfg.checkpoint.write_s3_config)
+
+    # Upload the file
+    with open(file_path_local, 'rb') as f:
+        s3_client.upload_fileobj(f, bucket_name, key)
+    full_s3_path = f"s3://{bucket_name}/{key}"
+    logger.info(f'copy file: {file_path_local} {full_s3_path} | use time: {time.time()-tic}')
+    return full_s3_path

third_party/PartField/partfield/model/PVCNN/encoder_pc.py

@@ -0,0 +1,243 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
+
+from ast import Dict
+import math
+
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch_scatter import scatter_mean #, scatter_max
+
+from .unet_3daware import setup_unet #UNetTriplane3dAware
+from .conv_pointnet import ConvPointnet
+
+from .pc_encoder import PVCNNEncoder #PointNet
+
+import einops
+
+from .dnnlib_util import ScopedTorchProfiler, printarr
+
+def generate_plane_features(p, c, resolution, plane='xz'):
+    """
+    Args:
+        p: (B,3,n_p)
+        c: (B,C,n_p)
+    """
+    padding = 0.
+    c_dim = c.size(1)
+    # acquire indices of features in plane
+    xy = normalize_coordinate(p.clone(), plane=plane, padding=padding) # normalize to the range of (0, 1)
+    index = coordinate2index(xy, resolution)
+
+    # scatter plane features from points
+    fea_plane = c.new_zeros(p.size(0), c_dim, resolution**2)
+    fea_plane = scatter_mean(c, index, out=fea_plane) # B x 512 x reso^2
+    fea_plane = fea_plane.reshape(p.size(0), c_dim, resolution, resolution) # sparce matrix (B x 512 x reso x reso)
+    return fea_plane
+
+def normalize_coordinate(p, padding=0.1, plane='xz'):
+    ''' Normalize coordinate to [0, 1] for unit cube experiments
+
+    Args:
+        p (tensor): point
+        padding (float): conventional padding paramter of ONet for unit cube, so [-0.5, 0.5] -> [-0.55, 0.55]
+        plane (str): plane feature type, ['xz', 'xy', 'yz']
+    '''
+    if plane == 'xz':
+        xy = p[:, :, [0, 2]]
+    elif plane =='xy':
+        xy = p[:, :, [0, 1]]
+    else:
+        xy = p[:, :, [1, 2]]
+
+    xy_new = xy / (1 + padding + 10e-6) # (-0.5, 0.5)
+    xy_new = xy_new + 0.5 # range (0, 1)
+
+    # if there are outliers out of the range
+    if xy_new.max() >= 1:
+        xy_new[xy_new >= 1] = 1 - 10e-6
+    if xy_new.min() < 0:
+        xy_new[xy_new < 0] = 0.0
+    return xy_new
+
+
+def coordinate2index(x, resolution):
+    ''' Normalize coordinate to [0, 1] for unit cube experiments.
+        Corresponds to our 3D model
+
+    Args:
+        x (tensor): coordinate
+        reso (int): defined resolution
+        coord_type (str): coordinate type
+    '''
+    x = (x * resolution).long()
+    index = x[:, :, 0] + resolution * x[:, :, 1]
+    index = index[:, None, :]
+    return index
+
+def softclip(x, min, max, hardness=5):
+    # Soft clipping for the logsigma
+    x = min + F.softplus(hardness*(x - min))/hardness
+    x = max - F.softplus(-hardness*(x - max))/hardness
+    return x
+
+
+def sample_triplane_feat(feature_triplane, normalized_pos):
+    '''
+        normalized_pos [-1, 1]
+    '''
+    tri_plane = torch.unbind(feature_triplane, dim=1)
+
+    x_feat = F.grid_sample(
+        tri_plane[0],
+        torch.cat(
+            [normalized_pos[:, :, 0:1], normalized_pos[:, :, 1:2]],
+            dim=-1).unsqueeze(dim=1), padding_mode='border',
+        align_corners=True)
+    y_feat = F.grid_sample(
+        tri_plane[1],
+        torch.cat(
+            [normalized_pos[:, :, 1:2], normalized_pos[:, :, 2:3]],
+            dim=-1).unsqueeze(dim=1), padding_mode='border',
+        align_corners=True)
+
+    z_feat = F.grid_sample(
+        tri_plane[2],
+        torch.cat(
+            [normalized_pos[:, :, 0:1], normalized_pos[:, :, 2:3]],
+            dim=-1).unsqueeze(dim=1), padding_mode='border',
+        align_corners=True)
+    final_feat = (x_feat + y_feat + z_feat)
+    final_feat = final_feat.squeeze(dim=2).permute(0, 2, 1)  # 32dimension
+    return final_feat
+
+
+# @persistence.persistent_class
+class TriPlanePC2Encoder(torch.nn.Module):
+    # Encoder that encode point cloud to triplane feature vector similar to ConvOccNet
+    def __init__(
+            self,
+            cfg,
+            device='cuda',
+            shape_min=-1.0,
+            shape_length=2.0,
+            use_2d_feat=False,
+            # point_encoder='pvcnn',
+            # use_point_scatter=False
+    ):
+        """
+        Outputs latent triplane from PC input
+        Configs:
+            max_logsigma: (float) Soft clip upper range for logsigm
+            min_logsigma: (float)
+            point_encoder_type: (str) one of ['pvcnn', 'pointnet']
+            pvcnn_flatten_voxels: (bool) for pvcnn whether to reduce voxel 
+                features (instead of scattering point features)
+            unet_cfg: (dict)
+            z_triplane_channels: (int) output latent triplane
+            z_triplane_resolution: (int)
+        Args:
+
+        """
+        # assert img_resolution >= 4 and img_resolution & (img_resolution - 1) == 0
+        super().__init__()
+        self.device = device
+
+        self.cfg = cfg
+
+        self.shape_min = shape_min
+        self.shape_length = shape_length
+
+        self.z_triplane_resolution = cfg.z_triplane_resolution
+        z_triplane_channels = cfg.z_triplane_channels
+
+        point_encoder_out_dim = z_triplane_channels #* 2
+
+        in_channels = 6
+        # self.resample_filter=[1, 3, 3, 1]
+        if cfg.point_encoder_type == 'pvcnn':
+            self.pc_encoder = PVCNNEncoder(point_encoder_out_dim, 
+            device=self.device, in_channels=in_channels, use_2d_feat=use_2d_feat)  # Encode it to a volume vector.
+        elif cfg.point_encoder_type == 'pointnet':
+            # TODO the pointnet was buggy, investigate
+            self.pc_encoder = ConvPointnet(c_dim=point_encoder_out_dim, 
+                                           dim=in_channels, hidden_dim=32, 
+                                           plane_resolution=self.z_triplane_resolution, 
+                                           padding=0)
+        else:
+            raise NotImplementedError(f"Point encoder {cfg.point_encoder_type} not implemented")
+
+        if cfg.unet_cfg.enabled:
+            self.unet_encoder = setup_unet(
+                output_channels=point_encoder_out_dim, 
+                input_channels=point_encoder_out_dim, 
+                unet_cfg=cfg.unet_cfg)
+        else:
+            self.unet_encoder = None
+
+    # @ScopedTorchProfiler('encode')
+    def encode(self, point_cloud_xyz, point_cloud_feature, mv_feat=None, pc2pc_idx=None) -> Dict:
+        # output = AttrDict()
+        point_cloud_xyz = (point_cloud_xyz - self.shape_min) / self.shape_length # [0, 1]
+        point_cloud_xyz = point_cloud_xyz - 0.5 # [-0.5, 0.5]
+        point_cloud = torch.cat([point_cloud_xyz, point_cloud_feature], dim=-1)
+
+        if self.cfg.point_encoder_type == 'pvcnn':
+            if mv_feat is not None:
+                pc_feat, points_feat = self.pc_encoder(point_cloud, mv_feat, pc2pc_idx)
+            else:
+                pc_feat, points_feat = self.pc_encoder(point_cloud)  # 3D feature volume: BxDx32x32x32
+            if self.cfg.use_point_scatter:
+                # Scattering from PVCNN point features
+                points_feat_ = points_feat[0]
+                # shape: batch, latent size, resolution, resolution (e.g. 16, 256, 64, 64)
+                pc_feat_1 = generate_plane_features(point_cloud_xyz, points_feat_, 
+                                                    resolution=self.z_triplane_resolution, plane='xy') 
+                pc_feat_2 = generate_plane_features(point_cloud_xyz, points_feat_,
+                                                    resolution=self.z_triplane_resolution, plane='yz') 
+                pc_feat_3 = generate_plane_features(point_cloud_xyz, points_feat_,
+                                                    resolution=self.z_triplane_resolution, plane='xz') 
+                pc_feat = pc_feat[0]
+
+            else:
+                pc_feat = pc_feat[0]
+                sf = self.z_triplane_resolution//32 # 32 is PVCNN's voxel dim
+
+                pc_feat_1 = torch.mean(pc_feat, dim=-1) #xy_plane, normalize in z plane
+                pc_feat_2 = torch.mean(pc_feat, dim=-3) #yz_plane, normalize in x plane
+                pc_feat_3 = torch.mean(pc_feat, dim=-2) #xz_plane, normalize in y plane
+
+                # nearest upsample
+                pc_feat_1 = einops.repeat(pc_feat_1, 'b c h w -> b c (h hm ) (w wm)', hm = sf, wm = sf)
+                pc_feat_2 = einops.repeat(pc_feat_2, 'b c h w -> b c (h hm) (w wm)', hm = sf, wm = sf)
+                pc_feat_3 = einops.repeat(pc_feat_3, 'b c h w -> b c (h hm) (w wm)', hm = sf, wm = sf)
+        elif self.cfg.point_encoder_type == 'pointnet':
+            assert self.cfg.use_point_scatter
+            # Run ConvPointnet
+            pc_feat = self.pc_encoder(point_cloud)
+            pc_feat_1 = pc_feat['xy'] #
+            pc_feat_2 = pc_feat['yz']
+            pc_feat_3 = pc_feat['xz']
+        else:
+            raise NotImplementedError()
+
+        if self.unet_encoder is not None:
+            # TODO eval adding a skip connection
+            # Unet expects B, 3, C, H, W
+            pc_feat_tri_plane_stack_pre = torch.stack([pc_feat_1, pc_feat_2, pc_feat_3], dim=1)
+            # dpc_feat_tri_plane_stack = self.unet_encoder(pc_feat_tri_plane_stack_pre)
+            # pc_feat_tri_plane_stack = pc_feat_tri_plane_stack_pre + dpc_feat_tri_plane_stack
+            pc_feat_tri_plane_stack = self.unet_encoder(pc_feat_tri_plane_stack_pre)
+            pc_feat_1, pc_feat_2, pc_feat_3 = torch.unbind(pc_feat_tri_plane_stack, dim=1)
+
+        return torch.stack([pc_feat_1, pc_feat_2, pc_feat_3], dim=1)
+        
+    def forward(self, point_cloud_xyz, point_cloud_feature=None, mv_feat=None, pc2pc_idx=None):
+        return self.encode(point_cloud_xyz, point_cloud_feature=point_cloud_feature, mv_feat=mv_feat, pc2pc_idx=pc2pc_idx)

third_party/PartField/partfield/model/PVCNN/pc_encoder.py

@@ -0,0 +1,90 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import functools
+
+from .pv_module import SharedMLP, PVConv 
+
+def create_pointnet_components(
+        blocks, in_channels, with_se=False, normalize=True, eps=0,
+        width_multiplier=1, voxel_resolution_multiplier=1, scale_pvcnn=False, device='cuda'):
+    r, vr = width_multiplier, voxel_resolution_multiplier
+    layers, concat_channels = [], 0
+    for out_channels, num_blocks, voxel_resolution in blocks:
+        out_channels = int(r * out_channels)
+        if voxel_resolution is None:
+            block = functools.partial(SharedMLP, device=device)
+        else:
+            block = functools.partial(
+                PVConv, kernel_size=3, resolution=int(vr * voxel_resolution),
+                with_se=with_se, normalize=normalize, eps=eps, scale_pvcnn=scale_pvcnn, device=device)
+        for _ in range(num_blocks):
+            layers.append(block(in_channels, out_channels))
+            in_channels = out_channels
+            concat_channels += out_channels
+    return layers, in_channels, concat_channels
+
+class PCMerger(nn.Module):
+# merge surface sampled PC and rendering backprojected PC (w/ 2D features):
+    def __init__(self, in_channels=204, device="cuda"):
+        super(PCMerger, self).__init__()
+        self.mlp_normal = SharedMLP(3, [128, 128], device=device)
+        self.mlp_rgb = SharedMLP(3, [128, 128], device=device)
+        self.mlp_sam = SharedMLP(204 - 6, [128, 128], device=device)
+
+    def forward(self, feat, mv_feat, pc2pc_idx):
+        mv_feat_normal = self.mlp_normal(mv_feat[:, :3, :])
+        mv_feat_rgb = self.mlp_rgb(mv_feat[:, 3:6, :])
+        mv_feat_sam = self.mlp_sam(mv_feat[:, 6:, :])
+
+        mv_feat_normal = mv_feat_normal.permute(0, 2, 1)
+        mv_feat_rgb = mv_feat_rgb.permute(0, 2, 1)
+        mv_feat_sam = mv_feat_sam.permute(0, 2, 1)
+        feat = feat.permute(0, 2, 1)
+        
+        for i in range(mv_feat.shape[0]):
+            mask = (pc2pc_idx[i] != -1).reshape(-1)
+            idx = pc2pc_idx[i][mask].reshape(-1)
+            feat[i][mask] += mv_feat_normal[i][idx] + mv_feat_rgb[i][idx] + mv_feat_sam[i][idx]
+            
+        return feat.permute(0, 2, 1)
+
+
+class PVCNNEncoder(nn.Module):
+    def __init__(self, pvcnn_feat_dim, device='cuda', in_channels=3, use_2d_feat=False):
+        super(PVCNNEncoder, self).__init__()
+        self.device = device
+        self.blocks = ((pvcnn_feat_dim, 1, 32), (128, 2, 16), (256, 1, 8))
+        self.use_2d_feat=use_2d_feat
+        if in_channels == 6:
+            self.append_channel = 2
+        elif in_channels == 3:
+            self.append_channel = 1
+        else:
+            raise NotImplementedError
+        layers, channels_point, concat_channels_point = create_pointnet_components(
+            blocks=self.blocks, in_channels=in_channels + self.append_channel, with_se=False, normalize=False,
+            width_multiplier=1, voxel_resolution_multiplier=1, scale_pvcnn=True,
+            device=device
+        )
+        self.encoder = nn.ModuleList(layers)#.to(self.device)
+        if self.use_2d_feat:
+            self.merger = PCMerger()
+
+        
+
+    def forward(self, input_pc, mv_feat=None, pc2pc_idx=None):
+        features = input_pc.permute(0, 2, 1) * 2  # make point cloud [-1, 1]
+        coords = features[:, :3, :]
+        out_features_list = []
+        voxel_feature_list = []
+        zero_padding = torch.zeros(features.shape[0], self.append_channel, features.shape[-1], device=features.device, dtype=torch.float)
+        features = torch.cat([features, zero_padding], dim=1)##################
+
+        for i in range(len(self.encoder)):
+            features, _, voxel_feature = self.encoder[i]((features, coords))
+            if i == 0 and mv_feat is not None:
+               features = self.merger(features, mv_feat.permute(0, 2, 1), pc2pc_idx)
+            out_features_list.append(features)
+            voxel_feature_list.append(voxel_feature)
+        return voxel_feature_list, out_features_list

third_party/PartField/partfield/model/PVCNN/pv_module/__init__.py

@@ -0,0 +1,2 @@
+from .pvconv import PVConv
+from .shared_mlp import SharedMLP

third_party/PartField/partfield/model/PVCNN/pv_module/ball_query.py

@@ -0,0 +1,34 @@
+import torch
+import torch.nn as nn
+
+from . import functional as F
+
+__all__ = ['BallQuery']
+
+
+class BallQuery(nn.Module):
+    def __init__(self, radius, num_neighbors, include_coordinates=True):
+        super().__init__()
+        self.radius = radius
+        self.num_neighbors = num_neighbors
+        self.include_coordinates = include_coordinates
+
+    def forward(self, points_coords, centers_coords, points_features=None):
+        points_coords = points_coords.contiguous()
+        centers_coords = centers_coords.contiguous()
+        neighbor_indices = F.ball_query(centers_coords, points_coords, self.radius, self.num_neighbors)
+        neighbor_coordinates = F.grouping(points_coords, neighbor_indices)
+        neighbor_coordinates = neighbor_coordinates - centers_coords.unsqueeze(-1)
+
+        if points_features is None:
+            assert self.include_coordinates, 'No Features For Grouping'
+            neighbor_features = neighbor_coordinates
+        else:
+            neighbor_features = F.grouping(points_features, neighbor_indices)
+            if self.include_coordinates:
+                neighbor_features = torch.cat([neighbor_coordinates, neighbor_features], dim=1)
+        return neighbor_features
+
+    def extra_repr(self):
+        return 'radius={}, num_neighbors={}{}'.format(
+            self.radius, self.num_neighbors, ', include coordinates' if self.include_coordinates else '')

third_party/PartField/partfield/model/PVCNN/pv_module/frustum.py

@@ -0,0 +1,141 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from . import functional as PF
+
+__all__ = ['FrustumPointNetLoss', 'get_box_corners_3d']
+
+
+class FrustumPointNetLoss(nn.Module):
+    def __init__(
+            self, num_heading_angle_bins, num_size_templates, size_templates, box_loss_weight=1.0,
+            corners_loss_weight=10.0, heading_residual_loss_weight=20.0, size_residual_loss_weight=20.0):
+        super().__init__()
+        self.box_loss_weight = box_loss_weight
+        self.corners_loss_weight = corners_loss_weight
+        self.heading_residual_loss_weight = heading_residual_loss_weight
+        self.size_residual_loss_weight = size_residual_loss_weight
+
+        self.num_heading_angle_bins = num_heading_angle_bins
+        self.num_size_templates = num_size_templates
+        self.register_buffer('size_templates', size_templates.view(self.num_size_templates, 3))
+        self.register_buffer(
+            'heading_angle_bin_centers', torch.arange(0, 2 * np.pi, 2 * np.pi / self.num_heading_angle_bins)
+        )
+
+    def forward(self, inputs, targets):
+        mask_logits = inputs['mask_logits']  # (B, 2, N)
+        center_reg = inputs['center_reg']  # (B, 3)
+        center = inputs['center']  # (B, 3)
+        heading_scores = inputs['heading_scores']  # (B, NH)
+        heading_residuals_normalized = inputs['heading_residuals_normalized']  # (B, NH)
+        heading_residuals = inputs['heading_residuals']  # (B, NH)
+        size_scores = inputs['size_scores']  # (B, NS)
+        size_residuals_normalized = inputs['size_residuals_normalized']  # (B, NS, 3)
+        size_residuals = inputs['size_residuals']  # (B, NS, 3)
+
+        mask_logits_target = targets['mask_logits']  # (B, N)
+        center_target = targets['center']  # (B, 3)
+        heading_bin_id_target = targets['heading_bin_id']  # (B, )
+        heading_residual_target = targets['heading_residual']  # (B, )
+        size_template_id_target = targets['size_template_id']  # (B, )
+        size_residual_target = targets['size_residual']  # (B, 3)
+
+        batch_size = center.size(0)
+        batch_id = torch.arange(batch_size, device=center.device)
+
+        # Basic Classification and Regression losses
+        mask_loss = F.cross_entropy(mask_logits, mask_logits_target)
+        heading_loss = F.cross_entropy(heading_scores, heading_bin_id_target)
+        size_loss = F.cross_entropy(size_scores, size_template_id_target)
+        center_loss = PF.huber_loss(torch.norm(center_target - center, dim=-1), delta=2.0)
+        center_reg_loss = PF.huber_loss(torch.norm(center_target - center_reg, dim=-1), delta=1.0)
+
+        # Refinement losses for size/heading
+        heading_residuals_normalized = heading_residuals_normalized[batch_id, heading_bin_id_target]  # (B, )
+        heading_residual_normalized_target = heading_residual_target / (np.pi / self.num_heading_angle_bins)
+        heading_residual_normalized_loss = PF.huber_loss(
+            heading_residuals_normalized - heading_residual_normalized_target, delta=1.0
+        )
+        size_residuals_normalized = size_residuals_normalized[batch_id, size_template_id_target]  # (B, 3)
+        size_residual_normalized_target = size_residual_target / self.size_templates[size_template_id_target]
+        size_residual_normalized_loss = PF.huber_loss(
+            torch.norm(size_residual_normalized_target - size_residuals_normalized, dim=-1), delta=1.0
+        )
+
+        # Bounding box losses
+        heading = (heading_residuals[batch_id, heading_bin_id_target]
+                   + self.heading_angle_bin_centers[heading_bin_id_target])  # (B, )
+        # Warning: in origin code, size_residuals are added twice (issue #43 and #49 in charlesq34/frustum-pointnets)
+        size = (size_residuals[batch_id, size_template_id_target]
+                + self.size_templates[size_template_id_target])  # (B, 3)
+        corners = get_box_corners_3d(centers=center, headings=heading, sizes=size, with_flip=False)  # (B, 3, 8)
+        heading_target = self.heading_angle_bin_centers[heading_bin_id_target] + heading_residual_target  # (B, )
+        size_target = self.size_templates[size_template_id_target] + size_residual_target  # (B, 3)
+        corners_target, corners_target_flip = get_box_corners_3d(
+            centers=center_target, headings=heading_target,
+            sizes=size_target, with_flip=True)  # (B, 3, 8)
+        corners_loss = PF.huber_loss(
+            torch.min(
+                torch.norm(corners - corners_target, dim=1), torch.norm(corners - corners_target_flip, dim=1)
+            ), delta=1.0)
+        # Summing up
+        loss = mask_loss + self.box_loss_weight * (
+                center_loss + center_reg_loss + heading_loss + size_loss
+                + self.heading_residual_loss_weight * heading_residual_normalized_loss
+                + self.size_residual_loss_weight * size_residual_normalized_loss
+                + self.corners_loss_weight * corners_loss
+        )
+
+        return loss
+
+
+def get_box_corners_3d(centers, headings, sizes, with_flip=False):
+    """
+    :param centers: coords of box centers, FloatTensor[N, 3]
+    :param headings: heading angles, FloatTensor[N, ]
+    :param sizes: box sizes, FloatTensor[N, 3]
+    :param with_flip: bool, whether to return flipped box (headings + np.pi)
+    :return:
+        coords of box corners, FloatTensor[N, 3, 8]
+        NOTE: corner points are in counter clockwise order, e.g.,
+          2--1
+        3--0 5
+        7--4
+    """
+    l = sizes[:, 0]  # (N,)
+    w = sizes[:, 1]  # (N,)
+    h = sizes[:, 2]  # (N,)
+    x_corners = torch.stack([l / 2, l / 2, -l / 2, -l / 2, l / 2, l / 2, -l / 2, -l / 2], dim=1)  # (N, 8)
+    y_corners = torch.stack([h / 2, h / 2, h / 2, h / 2, -h / 2, -h / 2, -h / 2, -h / 2], dim=1)  # (N, 8)
+    z_corners = torch.stack([w / 2, -w / 2, -w / 2, w / 2, w / 2, -w / 2, -w / 2, w / 2], dim=1)  # (N, 8)
+
+    c = torch.cos(headings)  # (N,)
+    s = torch.sin(headings)  # (N,)
+    o = torch.ones_like(headings)  # (N,)
+    z = torch.zeros_like(headings)  # (N,)
+
+    centers = centers.unsqueeze(-1)  # (B, 3, 1)
+    corners = torch.stack([x_corners, y_corners, z_corners], dim=1)  # (N, 3, 8)
+    R = torch.stack([c, z, s, z, o, z, -s, z, c], dim=1).view(-1, 3, 3)  # roty matrix: (N, 3, 3)
+    if with_flip:
+        R_flip = torch.stack([-c, z, -s, z, o, z, s, z, -c], dim=1).view(-1, 3, 3)
+        return torch.matmul(R, corners) + centers, torch.matmul(R_flip, corners) + centers
+    else:
+        return torch.matmul(R, corners) + centers
+
+    # centers = centers.unsqueeze(1)  # (B, 1, 3)
+    # corners = torch.stack([x_corners, y_corners, z_corners], dim=-1)  # (N, 8, 3)
+    # RT = torch.stack([c, z, -s, z, o, z, s, z, c], dim=1).view(-1, 3, 3)  # (N, 3, 3)
+    # if with_flip:
+    #     RT_flip = torch.stack([-c, z, s, z, o, z, -s, z, -c], dim=1).view(-1, 3, 3)  # (N, 3, 3)
+    #     return torch.matmul(corners, RT) + centers, torch.matmul(corners, RT_flip) + centers  # (N, 8, 3)
+    # else:
+    #     return torch.matmul(corners, RT) + centers  # (N, 8, 3)
+
+    # corners = torch.stack([x_corners, y_corners, z_corners], dim=1)  # (N, 3, 8)
+    # R = torch.stack([c, z, s, z, o, z, -s, z, c], dim=1).view(-1, 3, 3)  # (N, 3, 3)
+    # corners = torch.matmul(R, corners) + centers.unsqueeze(2)  # (N, 3, 8)
+    # corners = corners.transpose(1, 2)  # (N, 8, 3)

third_party/PartField/partfield/model/PVCNN/pv_module/functional/__init__.py

@@ -0,0 +1 @@
+from .devoxelization import trilinear_devoxelize

third_party/PartField/partfield/model/PVCNN/pv_module/functional/devoxelization.py

@@ -0,0 +1,12 @@
+from torch.autograd import Function
+import torch
+import torch.nn.functional as F
+
+__all__ = ['trilinear_devoxelize']
+
+def trilinear_devoxelize(c, coords, r, training=None):
+    coords = (coords * 2 + 1.0) / r - 1.0
+    coords = coords.permute(0, 2, 1).reshape(c.shape[0], 1, 1, -1, 3)
+    f = F.grid_sample(input=c, grid=coords, padding_mode='border', align_corners=False)
+    f = f.squeeze(dim=2).squeeze(dim=2)
+    return f

third_party/PartField/partfield/model/PVCNN/pv_module/loss.py

@@ -0,0 +1,10 @@
+import torch.nn as nn
+
+from . import functional as F
+
+__all__ = ['KLLoss']
+
+
+class KLLoss(nn.Module):
+    def forward(self, x, y):
+        return F.kl_loss(x, y)

third_party/PartField/partfield/model/PVCNN/pv_module/pointnet.py

@@ -0,0 +1,113 @@
+import torch
+import torch.nn as nn
+
+from . import functional as F
+from .ball_query import BallQuery
+from .shared_mlp import SharedMLP
+
+__all__ = ['PointNetAModule', 'PointNetSAModule', 'PointNetFPModule']
+
+
+class PointNetAModule(nn.Module):
+    def __init__(self, in_channels, out_channels, include_coordinates=True):
+        super().__init__()
+        if not isinstance(out_channels, (list, tuple)):
+            out_channels = [[out_channels]]
+        elif not isinstance(out_channels[0], (list, tuple)):
+            out_channels = [out_channels]
+
+        mlps = []
+        total_out_channels = 0
+        for _out_channels in out_channels:
+            mlps.append(
+                SharedMLP(
+                    in_channels=in_channels + (3 if include_coordinates else 0),
+                    out_channels=_out_channels, dim=1)
+            )
+            total_out_channels += _out_channels[-1]
+
+        self.include_coordinates = include_coordinates
+        self.out_channels = total_out_channels
+        self.mlps = nn.ModuleList(mlps)
+
+    def forward(self, inputs):
+        features, coords = inputs
+        if self.include_coordinates:
+            features = torch.cat([features, coords], dim=1)
+        coords = torch.zeros((coords.size(0), 3, 1), device=coords.device)
+        if len(self.mlps) > 1:
+            features_list = []
+            for mlp in self.mlps:
+                features_list.append(mlp(features).max(dim=-1, keepdim=True).values)
+            return torch.cat(features_list, dim=1), coords
+        else:
+            return self.mlps[0](features).max(dim=-1, keepdim=True).values, coords
+
+    def extra_repr(self):
+        return f'out_channels={self.out_channels}, include_coordinates={self.include_coordinates}'
+
+
+class PointNetSAModule(nn.Module):
+    def __init__(self, num_centers, radius, num_neighbors, in_channels, out_channels, include_coordinates=True):
+        super().__init__()
+        if not isinstance(radius, (list, tuple)):
+            radius = [radius]
+        if not isinstance(num_neighbors, (list, tuple)):
+            num_neighbors = [num_neighbors] * len(radius)
+        assert len(radius) == len(num_neighbors)
+        if not isinstance(out_channels, (list, tuple)):
+            out_channels = [[out_channels]] * len(radius)
+        elif not isinstance(out_channels[0], (list, tuple)):
+            out_channels = [out_channels] * len(radius)
+        assert len(radius) == len(out_channels)
+
+        groupers, mlps = [], []
+        total_out_channels = 0
+        for _radius, _out_channels, _num_neighbors in zip(radius, out_channels, num_neighbors):
+            groupers.append(
+                BallQuery(radius=_radius, num_neighbors=_num_neighbors, include_coordinates=include_coordinates)
+            )
+            mlps.append(
+                SharedMLP(
+                    in_channels=in_channels + (3 if include_coordinates else 0),
+                    out_channels=_out_channels, dim=2)
+            )
+            total_out_channels += _out_channels[-1]
+
+        self.num_centers = num_centers
+        self.out_channels = total_out_channels
+        self.groupers = nn.ModuleList(groupers)
+        self.mlps = nn.ModuleList(mlps)
+
+    def forward(self, inputs):
+        features, coords = inputs
+        centers_coords = F.furthest_point_sample(coords, self.num_centers)
+        features_list = []
+        for grouper, mlp in zip(self.groupers, self.mlps):
+            features_list.append(mlp(grouper(coords, centers_coords, features)).max(dim=-1).values)
+        if len(features_list) > 1:
+            return torch.cat(features_list, dim=1), centers_coords
+        else:
+            return features_list[0], centers_coords
+
+    def extra_repr(self):
+        return f'num_centers={self.num_centers}, out_channels={self.out_channels}'
+
+
+class PointNetFPModule(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.mlp = SharedMLP(in_channels=in_channels, out_channels=out_channels, dim=1)
+
+    def forward(self, inputs):
+        if len(inputs) == 3:
+            points_coords, centers_coords, centers_features = inputs
+            points_features = None
+        else:
+            points_coords, centers_coords, centers_features, points_features = inputs
+        interpolated_features = F.nearest_neighbor_interpolate(points_coords, centers_coords, centers_features)
+        if points_features is not None:
+            interpolated_features = torch.cat(
+                [interpolated_features, points_features], dim=1
+            )
+        return self.mlp(interpolated_features), points_coords

third_party/PartField/partfield/model/PVCNN/pv_module/pvconv.py

@@ -0,0 +1,38 @@
+import torch.nn as nn
+
+from . import functional as F
+from .voxelization import Voxelization
+from .shared_mlp import SharedMLP
+import torch
+
+__all__ = ['PVConv']
+
+
+class PVConv(nn.Module):
+    def __init__(
+            self, in_channels, out_channels, kernel_size, resolution, with_se=False, normalize=True, eps=0, scale_pvcnn=False,
+            device='cuda'):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.resolution = resolution
+        self.voxelization = Voxelization(resolution, normalize=normalize, eps=eps, scale_pvcnn=scale_pvcnn)
+        voxel_layers = [
+            nn.Conv3d(in_channels, out_channels, kernel_size, stride=1, padding=kernel_size // 2, device=device),
+            nn.InstanceNorm3d(out_channels, eps=1e-4, device=device),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv3d(out_channels, out_channels, kernel_size, stride=1, padding=kernel_size // 2, device=device),
+            nn.InstanceNorm3d(out_channels, eps=1e-4, device=device),
+            nn.LeakyReLU(0.1, True),
+        ]
+        self.voxel_layers = nn.Sequential(*voxel_layers)
+        self.point_features = SharedMLP(in_channels, out_channels, device=device)
+
+    def forward(self, inputs):
+        features, coords = inputs
+        voxel_features, voxel_coords = self.voxelization(features, coords)
+        voxel_features = self.voxel_layers(voxel_features)
+        devoxel_features = F.trilinear_devoxelize(voxel_features, voxel_coords, self.resolution, self.training)
+        fused_features = devoxel_features + self.point_features(features)
+        return fused_features, coords, voxel_features

third_party/PartField/partfield/model/PVCNN/pv_module/shared_mlp.py

@@ -0,0 +1,35 @@
+import torch.nn as nn
+
+__all__ = ['SharedMLP']
+
+
+class SharedMLP(nn.Module):
+    def __init__(self, in_channels, out_channels, dim=1, device='cuda'):
+        super().__init__()
+        # print('==> SharedMLP device: ', device)
+        if dim == 1:
+            conv = nn.Conv1d
+            bn = nn.InstanceNorm1d
+        elif dim == 2:
+            conv = nn.Conv2d
+            bn = nn.InstanceNorm1d
+        else:
+            raise ValueError
+        if not isinstance(out_channels, (list, tuple)):
+            out_channels = [out_channels]
+        layers = []
+        for oc in out_channels:
+            layers.extend(
+                [
+                    conv(in_channels, oc, 1, device=device),
+                    bn(oc, device=device),
+                    nn.ReLU(True),
+                ])
+            in_channels = oc
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, inputs):
+        if isinstance(inputs, (list, tuple)):
+            return (self.layers(inputs[0]), *inputs[1:])
+        else:
+            return self.layers(inputs)

third_party/PartField/partfield/model/PVCNN/pv_module/voxelization.py

@@ -0,0 +1,50 @@
+import torch
+import torch.nn as nn
+
+from . import functional as F
+
+__all__ = ['Voxelization']
+
+
+def my_voxelization(features, coords, resolution):
+    b, c, _ = features.shape
+    result = torch.zeros(b, c + 1, resolution * resolution * resolution, device=features.device, dtype=torch.float)
+    r = resolution
+    r2 = resolution * resolution
+    indices = coords[:, 0] * r2 + coords[:, 1] * r + coords[:, 2]
+    indices = indices.unsqueeze(dim=1).expand(-1, result.shape[1], -1)
+    features = torch.cat([features, torch.ones(features.shape[0], 1, features.shape[2], device=features.device, dtype=features.dtype)], dim=1)
+    out_feature = result.scatter_(index=indices.long(), src=features, dim=2, reduce='add')
+    cnt = out_feature[:, -1:, :]
+    zero_mask = (cnt == 0).float()
+    cnt = cnt * (1 - zero_mask) + zero_mask * 1e-5
+    vox_feature = out_feature[:, :-1, :] / cnt
+    return vox_feature.view(b, c, resolution, resolution, resolution)
+
+class Voxelization(nn.Module):
+    def __init__(self, resolution, normalize=True, eps=0, scale_pvcnn=False):
+        super().__init__()
+        self.r = int(resolution)
+        self.normalize = normalize
+        self.eps = eps
+        self.scale_pvcnn = scale_pvcnn
+        assert not normalize
+
+    def forward(self, features, coords):
+        with torch.no_grad():
+            coords = coords.detach()
+
+            if self.normalize:
+                norm_coords = norm_coords / (norm_coords.norm(dim=1, keepdim=True).max(dim=2, keepdim=True).values * 2.0 + self.eps) + 0.5
+            else:
+                if self.scale_pvcnn:
+                    norm_coords = (coords + 1) / 2.0 # [0, 1]
+                else:
+                    norm_coords = (norm_coords + 1) / 2.0
+            norm_coords = torch.clamp(norm_coords * self.r, 0, self.r - 1)
+            vox_coords = torch.round(norm_coords)
+        new_vox_feat = my_voxelization(features, vox_coords, self.r)
+        return new_vox_feat, norm_coords
+
+    def extra_repr(self):
+        return 'resolution={}{}'.format(self.r, ', normalized eps = {}'.format(self.eps) if self.normalize else '')

third_party/PartField/partfield/model/PVCNN/unet_3daware.py

@@ -0,0 +1,427 @@
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import init
+
+import einops
+
+def conv3x3(in_channels, out_channels, stride=1, 
+            padding=1, bias=True, groups=1):    
+    return nn.Conv2d(
+        in_channels,
+        out_channels,
+        kernel_size=3,
+        stride=stride,
+        padding=padding,
+        bias=bias,
+        groups=groups)
+
+def upconv2x2(in_channels, out_channels, mode='transpose'):
+    if mode == 'transpose':
+        return nn.ConvTranspose2d(
+            in_channels,
+            out_channels,
+            kernel_size=2,
+            stride=2)
+    else:
+        # out_channels is always going to be the same
+        # as in_channels
+        return nn.Sequential(
+            nn.Upsample(mode='bilinear', scale_factor=2),
+            conv1x1(in_channels, out_channels))
+
+def conv1x1(in_channels, out_channels, groups=1):
+    return nn.Conv2d(
+        in_channels,
+        out_channels,
+        kernel_size=1,
+        groups=groups,
+        stride=1)
+
+class ConvTriplane3dAware(nn.Module):
+    """ 3D aware triplane conv (as described in RODIN) """
+    def __init__(self, internal_conv_f, in_channels, out_channels, order='xz'):
+        """
+        Args:
+            internal_conv_f: function that should return a 2D convolution Module 
+                given in and out channels
+            order: if triplane input is in 'xz' order
+        """
+        super(ConvTriplane3dAware, self).__init__()
+        # Need 3 seperate convolutions
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        assert order in ['xz', 'zx']
+        self.order = order
+        # Going to stack from other planes
+        self.plane_convs =  nn.ModuleList([
+            internal_conv_f(3*self.in_channels, self.out_channels) for _ in range(3)])
+    
+    def forward(self, triplanes_list):
+        """
+        Args:
+            triplanes_list: [(B,Ci,H,W)]*3 in xy,yz,(zx or xz) depending on order
+        Returns:
+            out_triplanes_list: [(B,Co,H,W)]*3 in xy,yz,(zx or xz) depending on order
+        """
+        inps = list(triplanes_list)
+        xp = 1 #(yz)
+        yp = 2 #(zx)
+        zp = 0 #(xy)
+
+        if self.order == 'xz':
+            # get into zx order
+            inps[yp] = einops.rearrange(inps[yp], 'b c x z -> b c z x')
+
+
+        oplanes = [None]*3 
+        # order shouldn't matter
+        for iplane in [zp, xp, yp]:
+            # i_plane -> (j,k)
+
+            # need to average out i and convert to (j,k)
+            # j_plane -> (k,i)
+            # k_plane -> (i,j)
+            jplane = (iplane+1)%3
+            kplane = (iplane+2)%3
+
+            ifeat = inps[iplane]
+            # need to average out nonshared dim
+            # Average pool across
+
+            # j_plane -> (k,i) -> (k,1) -> (1,k) -> (j,k)
+            # b c k i -> b c k 1
+            jpool = torch.mean(inps[jplane], dim=3 ,keepdim=True)
+            jpool = einops.rearrange(jpool, 'b c k 1 -> b c 1 k')
+            jpool = einops.repeat(jpool, 'b c 1 k -> b c j k', j=ifeat.size(2))
+
+            # k_plane -> (i,j) -> (1,j) -> (j,1) -> (j,k)
+            # b c i j -> b c 1 j
+            kpool = torch.mean(inps[kplane], dim=2 ,keepdim=True)
+            kpool = einops.rearrange(kpool, 'b c 1 j -> b c j 1')
+            kpool = einops.repeat(kpool, 'b c j 1 -> b c j k', k=ifeat.size(3))
+
+            # b c h w
+            # jpool = jpool.expand_as(ifeat)
+            # kpool = kpool.expand_as(ifeat)
+
+            # concat and conv on feature dim
+            catfeat = torch.cat([ifeat, jpool, kpool], dim=1)
+            oplane = self.plane_convs[iplane](catfeat)
+            oplanes[iplane] = oplane
+
+        if self.order == 'xz':
+            # get back into xz order
+            oplanes[yp] = einops.rearrange(oplanes[yp], 'b c z x -> b c x z')
+
+        return oplanes
+
+def roll_triplanes(triplanes_list):
+    # B, C, tri, h, w
+    tristack = torch.stack((triplanes_list),dim=2)
+    return einops.rearrange(tristack, 'b c tri h w -> b c (tri h) w', tri=3)
+
+def unroll_triplanes(rolled_triplane):
+    # B, C, tri*h, w
+    tristack = einops.rearrange(rolled_triplane, 'b c (tri h) w -> b c tri h w', tri=3)
+    return torch.unbind(tristack, dim=2)
+
+def conv1x1triplane3daware(in_channels, out_channels, order='xz', **kwargs):    
+    return ConvTriplane3dAware(lambda inp, out: conv1x1(inp,out,**kwargs), 
+                               in_channels, out_channels,order=order)
+
+def Normalize(in_channels, num_groups=32):
+    num_groups = min(in_channels, num_groups)  # avoid error if in_channels < 32
+    return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+
+def nonlinearity(x):
+    # return F.relu(x)
+    # Swish
+    return x*torch.sigmoid(x)
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x):
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            # no asymmetric padding in torch conv, must do it ourselves
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=3,
+                                        stride=2,
+                                        padding=0)
+
+    def forward(self, x):
+        if self.with_conv:
+            pad = (0,1,0,1)
+            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+class ResnetBlock3dAware(nn.Module):
+    def __init__(self, in_channels, out_channels=None):
+            #, conv_shortcut=False):
+        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.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = conv3x3(self.in_channels, self.out_channels)
+
+        self.norm_mid = Normalize(out_channels)
+        self.conv_3daware = conv1x1triplane3daware(self.out_channels, self.out_channels)
+
+        self.norm2 = Normalize(out_channels)
+        self.conv2 = conv3x3(self.out_channels, self.out_channels)
+
+        if self.in_channels != self.out_channels:
+            self.nin_shortcut = torch.nn.Conv2d(in_channels,
+                                                out_channels,
+                                                kernel_size=1,
+                                                stride=1,
+                                                padding=0)
+
+    def forward(self, x):
+        # 3x3 plane comm
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        # 1x1 3d aware, crossplane comm
+        h = self.norm_mid(h)
+        h = nonlinearity(h)
+        h = unroll_triplanes(h)
+        h = self.conv_3daware(h)
+        h = roll_triplanes(h)
+
+        # 3x3 plane comm
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+
+        return x+h
+
+class DownConv3dAware(nn.Module):
+    """
+    A helper Module that performs 2 convolutions and 1 MaxPool.
+    A ReLU activation follows each convolution.
+    """
+    def __init__(self, in_channels, out_channels, downsample=True, with_conv=False):
+        super(DownConv3dAware, self).__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.block = ResnetBlock3dAware(in_channels=in_channels, 
+            out_channels=out_channels)
+
+        self.do_downsample = downsample
+        self.downsample = Downsample(out_channels, with_conv=with_conv)
+
+    def forward(self, x):
+        """
+        rolled input, rolled output
+        Args:
+            x: rolled (b c (tri*h) w)
+        """
+        x = self.block(x)
+        before_pool = x
+        # if self.pooling:
+        #     x = self.pool(x)
+        if self.do_downsample:
+            # unroll and cat channel-wise (to prevent pooling across triplane boundaries)
+            x = einops.rearrange(x, 'b c (tri h) w -> b (c tri) h w', tri=3)
+            x = self.downsample(x)
+            # undo
+            x = einops.rearrange(x, 'b (c tri) h w -> b c (tri h) w', tri=3)
+        return x, before_pool
+
+class UpConv3dAware(nn.Module):
+    """
+    A helper Module that performs 2 convolutions and 1 UpConvolution.
+    A ReLU activation follows each convolution.
+    """
+    def __init__(self, in_channels, out_channels, 
+                 merge_mode='concat', with_conv=False): #up_mode='transpose', ):
+        super(UpConv3dAware, self).__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.merge_mode = merge_mode
+            
+        self.upsample = Upsample(in_channels, with_conv)
+
+        if self.merge_mode == 'concat':
+            self.norm1 = Normalize(in_channels+out_channels)
+            self.block = ResnetBlock3dAware(in_channels=in_channels+out_channels, 
+                out_channels=out_channels)
+        else:
+            self.norm1 = Normalize(in_channels)
+            self.block = ResnetBlock3dAware(in_channels=in_channels, 
+                out_channels=out_channels)
+        
+
+    def forward(self, from_down, from_up):
+        """ Forward pass
+        rolled inputs, rolled output
+        rolled (b c (tri*h) w)
+        Arguments:
+            from_down: tensor from the encoder pathway
+            from_up: upconv'd tensor from the decoder pathway
+        """
+        # from_up = self.upconv(from_up)
+        from_up = self.upsample(from_up)
+        if self.merge_mode == 'concat':
+            x = torch.cat((from_up, from_down), 1)
+        else:
+            x = from_up + from_down
+
+        x = self.norm1(x)
+        x = self.block(x)
+        return x
+
+class UNetTriplane3dAware(nn.Module):
+    def __init__(self, out_channels, in_channels=3, depth=5, 
+                 start_filts=64,# up_mode='transpose', 
+                 use_initial_conv=False,
+                 merge_mode='concat', **kwargs):
+        """
+        Arguments:
+            in_channels: int, number of channels in the input tensor.
+                Default is 3 for RGB images.
+            depth: int, number of MaxPools in the U-Net.
+            start_filts: int, number of convolutional filters for the 
+                first conv.
+        """
+        super(UNetTriplane3dAware, self).__init__()
+    
+
+        self.out_channels = out_channels 
+        self.in_channels = in_channels
+        self.start_filts = start_filts
+        self.depth = depth
+
+        self.use_initial_conv = use_initial_conv
+        if use_initial_conv:
+            self.conv_initial = conv1x1(self.in_channels, self.start_filts)
+
+        self.down_convs = []
+        self.up_convs = []
+
+        # create the encoder pathway and add to a list
+        for i in range(depth):
+            if i == 0:
+                ins = self.start_filts if use_initial_conv else self.in_channels
+            else:
+                ins = outs
+            outs = self.start_filts*(2**i)
+            downsamp_it = True if i < depth-1 else False
+
+            down_conv = DownConv3dAware(ins, outs, downsample = downsamp_it)
+            self.down_convs.append(down_conv)
+
+        for i in range(depth-1):
+            ins = outs
+            outs = ins // 2
+            up_conv = UpConv3dAware(ins, outs,
+                merge_mode=merge_mode)
+            self.up_convs.append(up_conv)
+
+        # add the list of modules to current module
+        self.down_convs = nn.ModuleList(self.down_convs)
+        self.up_convs = nn.ModuleList(self.up_convs)
+
+        self.norm_out = Normalize(outs)
+        self.conv_final = conv1x1(outs, self.out_channels)
+
+        self.reset_params()
+
+    @staticmethod
+    def weight_init(m):
+        if isinstance(m, nn.Conv2d):
+            # init.xavier_normal_(m.weight, gain=0.1)
+            init.xavier_normal_(m.weight)
+            init.constant_(m.bias, 0)
+
+
+    def reset_params(self):
+        for i, m in enumerate(self.modules()):
+            self.weight_init(m)
+
+
+    def forward(self, x):
+        """
+        Args:
+            x: Stacked triplane expected to be in (B,3,C,H,W)
+        """
+        # Roll
+        x = einops.rearrange(x, 'b tri c h w -> b c (tri h) w', tri=3)
+
+        if self.use_initial_conv:
+            x = self.conv_initial(x)
+
+        encoder_outs = []
+        # encoder pathway, save outputs for merging
+        for i, module in enumerate(self.down_convs):
+            x, before_pool = module(x)
+            encoder_outs.append(before_pool)
+        
+        # Spend a block in the middle
+        # x = self.block_mid(x)
+
+        for i, module in enumerate(self.up_convs):
+            before_pool = encoder_outs[-(i+2)]
+            x = module(before_pool, x)
+        
+        x = self.norm_out(x)
+        
+        # No softmax is used. This means you need to use
+        # nn.CrossEntropyLoss is your training script,
+        # as this module includes a softmax already.
+        x = self.conv_final(nonlinearity(x))
+
+        # Unroll
+        x = einops.rearrange(x, 'b c (tri h) w -> b tri c h w', tri=3)
+        return x
+
+    
+def setup_unet(output_channels, input_channels, unet_cfg):
+    if unet_cfg['use_3d_aware']:
+        assert(unet_cfg['rolled'])
+        unet = UNetTriplane3dAware(
+                                        out_channels=output_channels, 
+                                        in_channels=input_channels, 
+                                        depth=unet_cfg['depth'], 
+                                        use_initial_conv=unet_cfg['use_initial_conv'],
+                                        start_filts=unet_cfg['start_hidden_channels'],)
+    else:
+        raise NotImplementedError
+    return unet
+

third_party/PartField/partfield/model/UNet/buildingblocks.py

@@ -0,0 +1,546 @@
+#https://github.com/wolny/pytorch-3dunet/blob/master/pytorch3dunet/unet3d/buildingblocks.py
+# MIT License
+
+# Copyright (c) 2018 Adrian Wolny
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from functools import partial
+
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+# from pytorch3dunet.unet3d.se import ChannelSELayer3D, ChannelSpatialSELayer3D, SpatialSELayer3D
+
+
+def create_conv(in_channels, out_channels, kernel_size, order, num_groups, padding,
+                dropout_prob, is3d):
+    """
+    Create a list of modules with together constitute a single conv layer with non-linearity
+    and optional batchnorm/groupnorm.
+
+    Args:
+        in_channels (int): number of input channels
+        out_channels (int): number of output channels
+        kernel_size(int or tuple): size of the convolving kernel
+        order (string): order of things, e.g.
+            'cr' -> conv + ReLU
+            'gcr' -> groupnorm + conv + ReLU
+            'cl' -> conv + LeakyReLU
+            'ce' -> conv + ELU
+            'bcr' -> batchnorm + conv + ReLU
+            'cbrd' -> conv + batchnorm + ReLU + dropout
+            'cbrD' -> conv + batchnorm + ReLU + dropout2d
+        num_groups (int): number of groups for the GroupNorm
+        padding (int or tuple): add zero-padding added to all three sides of the input
+        dropout_prob (float): dropout probability
+        is3d (bool): is3d (bool): if True use Conv3d, otherwise use Conv2d
+    Return:
+        list of tuple (name, module)
+    """
+    assert 'c' in order, "Conv layer MUST be present"
+    assert order[0] not in 'rle', 'Non-linearity cannot be the first operation in the layer'
+
+    modules = []
+    for i, char in enumerate(order):
+        if char == 'r':
+            modules.append(('ReLU', nn.ReLU(inplace=True)))
+        elif char == 'l':
+            modules.append(('LeakyReLU', nn.LeakyReLU(inplace=True)))
+        elif char == 'e':
+            modules.append(('ELU', nn.ELU(inplace=True)))
+        elif char == 'c':
+            # add learnable bias only in the absence of batchnorm/groupnorm
+            bias = not ('g' in order or 'b' in order)
+            if is3d:
+                conv = nn.Conv3d(in_channels, out_channels, kernel_size, padding=padding, bias=bias)
+            else:
+                conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding, bias=bias)
+
+            modules.append(('conv', conv))
+        elif char == 'g':
+            is_before_conv = i < order.index('c')
+            if is_before_conv:
+                num_channels = in_channels
+            else:
+                num_channels = out_channels
+
+            # use only one group if the given number of groups is greater than the number of channels
+            if num_channels < num_groups:
+                num_groups = 1
+
+            assert num_channels % num_groups == 0, f'Expected number of channels in input to be divisible by num_groups. num_channels={num_channels}, num_groups={num_groups}'
+            modules.append(('groupnorm', nn.GroupNorm(num_groups=num_groups, num_channels=num_channels)))
+        elif char == 'b':
+            is_before_conv = i < order.index('c')
+            if is3d:
+                bn = nn.BatchNorm3d
+            else:
+                bn = nn.BatchNorm2d
+
+            if is_before_conv:
+                modules.append(('batchnorm', bn(in_channels)))
+            else:
+                modules.append(('batchnorm', bn(out_channels)))
+        elif char == 'd':
+            modules.append(('dropout', nn.Dropout(p=dropout_prob)))
+        elif char == 'D':
+            modules.append(('dropout2d', nn.Dropout2d(p=dropout_prob)))
+        else:
+            raise ValueError(f"Unsupported layer type '{char}'. MUST be one of ['b', 'g', 'r', 'l', 'e', 'c', 'd', 'D']")
+
+    return modules
+
+
+class SingleConv(nn.Sequential):
+    """
+    Basic convolutional module consisting of a Conv3d, non-linearity and optional batchnorm/groupnorm. The order
+    of operations can be specified via the `order` parameter
+
+    Args:
+        in_channels (int): number of input channels
+        out_channels (int): number of output channels
+        kernel_size (int or tuple): size of the convolving kernel
+        order (string): determines the order of layers, e.g.
+            'cr' -> conv + ReLU
+            'crg' -> conv + ReLU + groupnorm
+            'cl' -> conv + LeakyReLU
+            'ce' -> conv + ELU
+        num_groups (int): number of groups for the GroupNorm
+        padding (int or tuple): add zero-padding
+        dropout_prob (float): dropout probability, default 0.1
+        is3d (bool): if True use Conv3d, otherwise use Conv2d
+    """
+
+    def __init__(self, in_channels, out_channels, kernel_size=3, order='gcr', num_groups=8,
+                 padding=1, dropout_prob=0.1, is3d=True):
+        super(SingleConv, self).__init__()
+
+        for name, module in create_conv(in_channels, out_channels, kernel_size, order,
+                                        num_groups, padding, dropout_prob, is3d):
+            self.add_module(name, module)
+
+
+class DoubleConv(nn.Sequential):
+    """
+    A module consisting of two consecutive convolution layers (e.g. BatchNorm3d+ReLU+Conv3d).
+    We use (Conv3d+ReLU+GroupNorm3d) by default.
+    This can be changed however by providing the 'order' argument, e.g. in order
+    to change to Conv3d+BatchNorm3d+ELU use order='cbe'.
+    Use padded convolutions to make sure that the output (H_out, W_out) is the same
+    as (H_in, W_in), so that you don't have to crop in the decoder path.
+
+    Args:
+        in_channels (int): number of input channels
+        out_channels (int): number of output channels
+        encoder (bool): if True we're in the encoder path, otherwise we're in the decoder
+        kernel_size (int or tuple): size of the convolving kernel
+        order (string): determines the order of layers, e.g.
+            'cr' -> conv + ReLU
+            'crg' -> conv + ReLU + groupnorm
+            'cl' -> conv + LeakyReLU
+            'ce' -> conv + ELU
+        num_groups (int): number of groups for the GroupNorm
+        padding (int or tuple): add zero-padding added to all three sides of the input
+        upscale (int): number of the convolution to upscale in encoder if DoubleConv, default: 2
+        dropout_prob (float or tuple): dropout probability for each convolution, default 0.1
+        is3d (bool): if True use Conv3d instead of Conv2d layers
+    """
+
+    def __init__(self, in_channels, out_channels, encoder, kernel_size=3, order='gcr',
+                 num_groups=8, padding=1, upscale=2, dropout_prob=0.1, is3d=True):
+        super(DoubleConv, self).__init__()
+        if encoder:
+            # we're in the encoder path
+            conv1_in_channels = in_channels
+            if upscale == 1:
+                conv1_out_channels = out_channels
+            else:
+                conv1_out_channels = out_channels // 2
+            if conv1_out_channels < in_channels:
+                conv1_out_channels = in_channels
+            conv2_in_channels, conv2_out_channels = conv1_out_channels, out_channels
+        else:
+            # we're in the decoder path, decrease the number of channels in the 1st convolution
+            conv1_in_channels, conv1_out_channels = in_channels, out_channels
+            conv2_in_channels, conv2_out_channels = out_channels, out_channels
+
+        # check if dropout_prob is a tuple and if so
+        # split it for different dropout probabilities for each convolution.
+        if isinstance(dropout_prob, list) or isinstance(dropout_prob, tuple):
+            dropout_prob1 = dropout_prob[0]
+            dropout_prob2 = dropout_prob[1]
+        else:
+            dropout_prob1 = dropout_prob2 = dropout_prob
+
+        # conv1
+        self.add_module('SingleConv1',
+                        SingleConv(conv1_in_channels, conv1_out_channels, kernel_size, order, num_groups,
+                                   padding=padding, dropout_prob=dropout_prob1, is3d=is3d))
+        # conv2
+        self.add_module('SingleConv2',
+                        SingleConv(conv2_in_channels, conv2_out_channels, kernel_size, order, num_groups,
+                                   padding=padding, dropout_prob=dropout_prob2, is3d=is3d))
+
+
+class ResNetBlock(nn.Module):
+    """
+    Residual block that can be used instead of standard DoubleConv in the Encoder module.
+    Motivated by: https://arxiv.org/pdf/1706.00120.pdf
+
+    Notice we use ELU instead of ReLU (order='cge') and put non-linearity after the groupnorm.
+    """
+
+    def __init__(self, in_channels, out_channels, kernel_size=3, order='cge', num_groups=8, is3d=True, **kwargs):
+        super(ResNetBlock, self).__init__()
+
+        if in_channels != out_channels:
+            # conv1x1 for increasing the number of channels
+            if is3d:
+                self.conv1 = nn.Conv3d(in_channels, out_channels, 1)
+            else:
+                self.conv1 = nn.Conv2d(in_channels, out_channels, 1)
+        else:
+            self.conv1 = nn.Identity()
+
+        self.conv2 = SingleConv(in_channels, out_channels, kernel_size=kernel_size, order=order, num_groups=num_groups,
+                                is3d=is3d)
+        # remove non-linearity from the 3rd convolution since it's going to be applied after adding the residual
+        n_order = order
+        for c in 'rel':
+            n_order = n_order.replace(c, '')
+        self.conv3 = SingleConv(out_channels, out_channels, kernel_size=kernel_size, order=n_order,
+                                num_groups=num_groups, is3d=is3d)
+
+        # create non-linearity separately
+        if 'l' in order:
+            self.non_linearity = nn.LeakyReLU(negative_slope=0.1, inplace=True)
+        elif 'e' in order:
+            self.non_linearity = nn.ELU(inplace=True)
+        else:
+            self.non_linearity = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        # apply first convolution to bring the number of channels to out_channels
+        residual = self.conv1(x)
+
+        out = self.conv2(x)
+        out = self.conv3(out)
+
+        out += residual
+        out = self.non_linearity(out)
+
+        return out
+
+class Encoder(nn.Module):
+    """
+    A single module from the encoder path consisting of the optional max
+    pooling layer (one may specify the MaxPool kernel_size to be different
+    from the standard (2,2,2), e.g. if the volumetric data is anisotropic
+    (make sure to use complementary scale_factor in the decoder path) followed by
+    a basic module (DoubleConv or ResNetBlock).
+
+    Args:
+        in_channels (int): number of input channels
+        out_channels (int): number of output channels
+        conv_kernel_size (int or tuple): size of the convolving kernel
+        apply_pooling (bool): if True use MaxPool3d before DoubleConv
+        pool_kernel_size (int or tuple): the size of the window
+        pool_type (str): pooling layer: 'max' or 'avg'
+        basic_module(nn.Module): either ResNetBlock or DoubleConv
+        conv_layer_order (string): determines the order of layers
+            in `DoubleConv` module. See `DoubleConv` for more info.
+        num_groups (int): number of groups for the GroupNorm
+        padding (int or tuple): add zero-padding added to all three sides of the input
+        upscale (int): number of the convolution to upscale in encoder if DoubleConv, default: 2
+        dropout_prob (float or tuple): dropout probability, default 0.1
+        is3d (bool): use 3d or 2d convolutions/pooling operation
+    """
+
+    def __init__(self, in_channels, out_channels, conv_kernel_size=3, apply_pooling=True,
+                 pool_kernel_size=2, pool_type='max', basic_module=DoubleConv, conv_layer_order='gcr',
+                 num_groups=8, padding=1, upscale=2, dropout_prob=0.1, is3d=True):
+        super(Encoder, self).__init__()
+        assert pool_type in ['max', 'avg']
+        if apply_pooling:
+            if pool_type == 'max':
+                if is3d:
+                    self.pooling = nn.MaxPool3d(kernel_size=pool_kernel_size)
+                else:
+                    self.pooling = nn.MaxPool2d(kernel_size=pool_kernel_size)
+            else:
+                if is3d:
+                    self.pooling = nn.AvgPool3d(kernel_size=pool_kernel_size)
+                else:
+                    self.pooling = nn.AvgPool2d(kernel_size=pool_kernel_size)
+        else:
+            self.pooling = None
+
+        self.basic_module = basic_module(in_channels, out_channels,
+                                         encoder=True,
+                                         kernel_size=conv_kernel_size,
+                                         order=conv_layer_order,
+                                         num_groups=num_groups,
+                                         padding=padding,
+                                         upscale=upscale,
+                                         dropout_prob=dropout_prob,
+                                         is3d=is3d)
+
+    def forward(self, x):
+        if self.pooling is not None:
+            x = self.pooling(x)
+        x = self.basic_module(x)
+        return x
+
+
+class Decoder(nn.Module):
+    """
+    A single module for decoder path consisting of the upsampling layer
+    (either learned ConvTranspose3d or nearest neighbor interpolation)
+    followed by a basic module (DoubleConv or ResNetBlock).
+
+    Args:
+        in_channels (int): number of input channels
+        out_channels (int): number of output channels
+        conv_kernel_size (int or tuple): size of the convolving kernel
+        scale_factor (int or tuple): used as the multiplier for the image H/W/D in
+            case of nn.Upsample or as stride in case of ConvTranspose3d, must reverse the MaxPool3d operation
+            from the corresponding encoder
+        basic_module(nn.Module): either ResNetBlock or DoubleConv
+        conv_layer_order (string): determines the order of layers
+            in `DoubleConv` module. See `DoubleConv` for more info.
+        num_groups (int): number of groups for the GroupNorm
+        padding (int or tuple): add zero-padding added to all three sides of the input
+        upsample (str): algorithm used for upsampling:
+            InterpolateUpsampling:   'nearest' | 'linear' | 'bilinear' | 'trilinear' | 'area'
+            TransposeConvUpsampling: 'deconv'
+            No upsampling:           None
+            Default: 'default' (chooses automatically)
+        dropout_prob (float or tuple): dropout probability, default 0.1
+    """
+
+    def __init__(self, in_channels, out_channels, conv_kernel_size=3, scale_factor=2, basic_module=DoubleConv,
+                 conv_layer_order='gcr', num_groups=8, padding=1, upsample='default',
+                 dropout_prob=0.1, is3d=True):
+        super(Decoder, self).__init__()
+
+        # perform concat joining per default
+        concat = True
+
+        # don't adapt channels after join operation
+        adapt_channels = False
+
+        if upsample is not None and upsample != 'none':
+            if upsample == 'default':
+                if basic_module == DoubleConv:
+                    upsample = 'nearest'  # use nearest neighbor interpolation for upsampling
+                    concat = True  # use concat joining
+                    adapt_channels = False  # don't adapt channels
+                elif basic_module == ResNetBlock: #or basic_module == ResNetBlockSE:
+                    upsample = 'deconv'  # use deconvolution upsampling
+                    concat = False  # use summation joining
+                    adapt_channels = True  # adapt channels after joining
+
+            # perform deconvolution upsampling if mode is deconv
+            if upsample == 'deconv':
+                self.upsampling = TransposeConvUpsampling(in_channels=in_channels, out_channels=out_channels,
+                                                          kernel_size=conv_kernel_size, scale_factor=scale_factor,
+                                                          is3d=is3d)
+            else:
+                self.upsampling = InterpolateUpsampling(mode=upsample)
+        else:
+            # no upsampling
+            self.upsampling = NoUpsampling()
+            # concat joining
+            self.joining = partial(self._joining, concat=True)
+
+        # perform joining operation
+        self.joining = partial(self._joining, concat=concat)
+
+        # adapt the number of in_channels for the ResNetBlock
+        if adapt_channels is True:
+            in_channels = out_channels
+
+        self.basic_module = basic_module(in_channels, out_channels,
+                                         encoder=False,
+                                         kernel_size=conv_kernel_size,
+                                         order=conv_layer_order,
+                                         num_groups=num_groups,
+                                         padding=padding,
+                                         dropout_prob=dropout_prob,
+                                         is3d=is3d)
+
+    def forward(self, encoder_features, x):
+        x = self.upsampling(encoder_features=encoder_features, x=x)
+        x = self.joining(encoder_features, x)
+        x = self.basic_module(x)
+        return x
+
+    @staticmethod
+    def _joining(encoder_features, x, concat):
+        if concat:
+            return torch.cat((encoder_features, x), dim=1)
+        else:
+            return encoder_features + x
+
+
+def create_encoders(in_channels, f_maps, basic_module, conv_kernel_size, conv_padding,
+                    conv_upscale, dropout_prob,
+                    layer_order, num_groups, pool_kernel_size, is3d):
+    # create encoder path consisting of Encoder modules. Depth of the encoder is equal to `len(f_maps)`
+    encoders = []
+    for i, out_feature_num in enumerate(f_maps):
+        if i == 0:
+            # apply conv_coord only in the first encoder if any
+            encoder = Encoder(in_channels, out_feature_num,
+                              apply_pooling=False,  # skip pooling in the firs encoder
+                              basic_module=basic_module,
+                              conv_layer_order=layer_order,
+                              conv_kernel_size=conv_kernel_size,
+                              num_groups=num_groups,
+                              padding=conv_padding,
+                              upscale=conv_upscale,
+                              dropout_prob=dropout_prob,
+                              is3d=is3d)
+        else:
+            encoder = Encoder(f_maps[i - 1], out_feature_num,
+                              basic_module=basic_module,
+                              conv_layer_order=layer_order,
+                              conv_kernel_size=conv_kernel_size,
+                              num_groups=num_groups,
+                              pool_kernel_size=pool_kernel_size,
+                              padding=conv_padding,
+                              upscale=conv_upscale,
+                              dropout_prob=dropout_prob,
+                              is3d=is3d)
+
+        encoders.append(encoder)
+
+    return nn.ModuleList(encoders)
+
+
+def create_decoders(f_maps, basic_module, conv_kernel_size, conv_padding, layer_order,
+                    num_groups, upsample, dropout_prob, is3d):
+    # create decoder path consisting of the Decoder modules. The length of the decoder list is equal to `len(f_maps) - 1`
+    decoders = []
+    reversed_f_maps = list(reversed(f_maps[1:]))
+    for i in range(len(reversed_f_maps) - 1):
+        if basic_module == DoubleConv and upsample != 'deconv':
+            in_feature_num = reversed_f_maps[i] + reversed_f_maps[i + 1]
+        else:
+            in_feature_num = reversed_f_maps[i]
+
+        out_feature_num = reversed_f_maps[i + 1]
+
+        decoder = Decoder(in_feature_num, out_feature_num,
+                          basic_module=basic_module,
+                          conv_layer_order=layer_order,
+                          conv_kernel_size=conv_kernel_size,
+                          num_groups=num_groups,
+                          padding=conv_padding,
+                          upsample=upsample,
+                          dropout_prob=dropout_prob,
+                          is3d=is3d)
+        decoders.append(decoder)
+    return nn.ModuleList(decoders)
+
+
+class AbstractUpsampling(nn.Module):
+    """
+    Abstract class for upsampling. A given implementation should upsample a given 5D input tensor using either
+    interpolation or learned transposed convolution.
+    """
+
+    def __init__(self, upsample):
+        super(AbstractUpsampling, self).__init__()
+        self.upsample = upsample
+
+    def forward(self, encoder_features, x):
+        # get the spatial dimensions of the output given the encoder_features
+        output_size = encoder_features.size()[2:]
+        # upsample the input and return
+        return self.upsample(x, output_size)
+
+
+class InterpolateUpsampling(AbstractUpsampling):
+    """
+    Args:
+        mode (str): algorithm used for upsampling:
+            'nearest' | 'linear' | 'bilinear' | 'trilinear' | 'area'. Default: 'nearest'
+            used only if transposed_conv is False
+    """
+
+    def __init__(self, mode='nearest'):
+        upsample = partial(self._interpolate, mode=mode)
+        super().__init__(upsample)
+
+    @staticmethod
+    def _interpolate(x, size, mode):
+        return F.interpolate(x, size=size, mode=mode)
+
+
+class TransposeConvUpsampling(AbstractUpsampling):
+    """
+    Args:
+        in_channels (int): number of input channels for transposed conv
+            used only if transposed_conv is True
+        out_channels (int): number of output channels for transpose conv
+            used only if transposed_conv is True
+        kernel_size (int or tuple): size of the convolving kernel
+            used only if transposed_conv is True
+        scale_factor (int or tuple): stride of the convolution
+            used only if transposed_conv is True
+        is3d (bool): if True use ConvTranspose3d, otherwise use ConvTranspose2d
+    """
+
+    class Upsample(nn.Module):
+        """
+        Workaround the 'ValueError: requested an output size...' in the `_output_padding` method in
+        transposed convolution. It performs transposed conv followed by the interpolation to the correct size if necessary.
+        """
+
+        def __init__(self, conv_transposed, is3d):
+            super().__init__()
+            self.conv_transposed = conv_transposed
+            self.is3d = is3d
+
+        def forward(self, x, size):
+            x = self.conv_transposed(x)
+            return F.interpolate(x, size=size)
+
+    def __init__(self, in_channels, out_channels, kernel_size=3, scale_factor=2, is3d=True):
+        # make sure that the output size reverses the MaxPool3d from the corresponding encoder
+        if is3d is True:
+            conv_transposed = nn.ConvTranspose3d(in_channels, out_channels, kernel_size=kernel_size,
+                                                 stride=scale_factor, padding=1, bias=False)
+        else:
+            conv_transposed = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=kernel_size,
+                                                 stride=scale_factor, padding=1, bias=False)
+        upsample = self.Upsample(conv_transposed, is3d)
+        super().__init__(upsample)
+
+
+class NoUpsampling(AbstractUpsampling):
+    def __init__(self):
+        super().__init__(self._no_upsampling)
+
+    @staticmethod
+    def _no_upsampling(x, size):
+        return x

third_party/PartField/partfield/model/UNet/model.py

@@ -0,0 +1,170 @@
+# https://github.com/wolny/pytorch-3dunet/blob/master/pytorch3dunet/unet3d/buildingblocks.py
+# MIT License
+
+# Copyright (c) 2018 Adrian Wolny
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import torch.nn as nn
+
+from partfield.model.UNet.buildingblocks import DoubleConv, ResNetBlock, \
+    create_decoders, create_encoders
+
+def number_of_features_per_level(init_channel_number, num_levels):
+    return [init_channel_number * 2 ** k for k in range(num_levels)]
+
+class AbstractUNet(nn.Module):
+    """
+    Base class for standard and residual UNet.
+
+    Args:
+        in_channels (int): number of input channels
+        out_channels (int): number of output segmentation masks;
+            Note that the of out_channels might correspond to either
+            different semantic classes or to different binary segmentation mask.
+            It's up to the user of the class to interpret the out_channels and
+            use the proper loss criterion during training (i.e. CrossEntropyLoss (multi-class)
+            or BCEWithLogitsLoss (two-class) respectively)
+        f_maps (int, tuple): number of feature maps at each level of the encoder; if it's an integer the number
+            of feature maps is given by the geometric progression: f_maps ^ k, k=1,2,3,4
+        final_sigmoid (bool): if True apply element-wise nn.Sigmoid after the final 1x1 convolution,
+            otherwise apply nn.Softmax. In effect only if `self.training == False`, i.e. during validation/testing
+        basic_module: basic model for the encoder/decoder (DoubleConv, ResNetBlock, ....)
+        layer_order (string): determines the order of layers in `SingleConv` module.
+            E.g. 'crg' stands for GroupNorm3d+Conv3d+ReLU. See `SingleConv` for more info
+        num_groups (int): number of groups for the GroupNorm
+        num_levels (int): number of levels in the encoder/decoder path (applied only if f_maps is an int)
+            default: 4
+        is_segmentation (bool): if True and the model is in eval mode, Sigmoid/Softmax normalization is applied
+            after the final convolution; if False (regression problem) the normalization layer is skipped
+        conv_kernel_size (int or tuple): size of the convolving kernel in the basic_module
+        pool_kernel_size (int or tuple): the size of the window
+        conv_padding (int or tuple): add zero-padding added to all three sides of the input
+        conv_upscale (int): number of the convolution to upscale in encoder if DoubleConv, default: 2
+        upsample (str): algorithm used for decoder upsampling:
+            InterpolateUpsampling:   'nearest' | 'linear' | 'bilinear' | 'trilinear' | 'area'
+            TransposeConvUpsampling: 'deconv'
+            No upsampling:           None
+            Default: 'default' (chooses automatically)
+        dropout_prob (float or tuple): dropout probability, default: 0.1
+        is3d (bool): if True the model is 3D, otherwise 2D, default: True
+    """
+
+    def __init__(self, in_channels, out_channels, final_sigmoid, basic_module, f_maps=64, layer_order='gcr',
+                 num_groups=8, num_levels=4, is_segmentation=False, conv_kernel_size=3, pool_kernel_size=2,
+                 conv_padding=1, conv_upscale=2, upsample='default', dropout_prob=0.1, is3d=True, encoder_only=False):
+        super(AbstractUNet, self).__init__()
+
+        if isinstance(f_maps, int):
+            f_maps = number_of_features_per_level(f_maps, num_levels=num_levels)
+
+        assert isinstance(f_maps, list) or isinstance(f_maps, tuple)
+        assert len(f_maps) > 1, "Required at least 2 levels in the U-Net"
+        if 'g' in layer_order:
+            assert num_groups is not None, "num_groups must be specified if GroupNorm is used"
+
+        # create encoder path
+        self.encoders = create_encoders(in_channels, f_maps, basic_module, conv_kernel_size,
+                                        conv_padding, conv_upscale, dropout_prob,
+                                        layer_order, num_groups, pool_kernel_size, is3d)
+
+        self.encoder_only = encoder_only
+        
+        if encoder_only == False:
+            # create decoder path
+            self.decoders = create_decoders(f_maps, basic_module, conv_kernel_size, conv_padding,
+                                            layer_order, num_groups, upsample, dropout_prob,
+                                            is3d)
+
+            # in the last layer a 1×1 convolution reduces the number of output channels to the number of labels
+            if is3d:
+                self.final_conv = nn.Conv3d(f_maps[1], out_channels, 1)
+            else:
+                self.final_conv = nn.Conv2d(f_maps[1], out_channels, 1)
+
+            if is_segmentation:
+                # semantic segmentation problem
+                if final_sigmoid:
+                    self.final_activation = nn.Sigmoid()
+                else:
+                    self.final_activation = nn.Softmax(dim=1)
+            else:
+                # regression problem
+                self.final_activation = None
+
+    def forward(self, x, return_bottleneck_feat=False):
+        # encoder part
+        encoders_features = []
+        for encoder in self.encoders:
+            x = encoder(x)
+            # reverse the encoder outputs to be aligned with the decoder
+            encoders_features.insert(0, x)
+
+        # remove the last encoder's output from the list
+        # !!remember: it's the 1st in the list
+        bottleneck_feat = encoders_features[0]
+        if self.encoder_only:
+            return bottleneck_feat
+        else:
+            encoders_features = encoders_features[1:]
+
+            # decoder part
+            for decoder, encoder_features in zip(self.decoders, encoders_features):
+                # pass the output from the corresponding encoder and the output
+                # of the previous decoder
+                x = decoder(encoder_features, x)
+
+            x = self.final_conv(x)
+            # During training the network outputs logits
+            if self.final_activation is not None:
+                x = self.final_activation(x)
+
+            if return_bottleneck_feat:
+                return x, bottleneck_feat
+            else:
+                return x
+
+class ResidualUNet3D(AbstractUNet):
+    """
+    Residual 3DUnet model implementation based on https://arxiv.org/pdf/1706.00120.pdf.
+    Uses ResNetBlock as a basic building block, summation joining instead
+    of concatenation joining and transposed convolutions for upsampling (watch out for block artifacts).
+    Since the model effectively becomes a residual net, in theory it allows for deeper UNet.
+    """
+
+    def __init__(self, in_channels, out_channels, final_sigmoid=True, f_maps=(8, 16, 64, 256, 1024), layer_order='gcr',
+                 num_groups=8, num_levels=5, is_segmentation=True, conv_padding=1,
+                 conv_upscale=2, upsample='default', dropout_prob=0.1, encoder_only=False,  **kwargs):
+        super(ResidualUNet3D, self).__init__(in_channels=in_channels,
+                                             out_channels=out_channels,
+                                             final_sigmoid=final_sigmoid,
+                                             basic_module=ResNetBlock,
+                                             f_maps=f_maps,
+                                             layer_order=layer_order,
+                                             num_groups=num_groups,
+                                             num_levels=num_levels,
+                                             is_segmentation=is_segmentation,
+                                             conv_padding=conv_padding,
+                                             conv_upscale=conv_upscale,
+                                             upsample=upsample,
+                                             dropout_prob=dropout_prob,
+                                             encoder_only=encoder_only,
+                                             is3d=True)
+
+

third_party/PartField/partfield/model/model_utils.py

@@ -0,0 +1,54 @@
+import torch
+import torch.nn as nn
+
+class VanillaMLP(nn.Module):
+    def __init__(self, input_dim, output_dim, out_activation, n_hidden_layers=4, n_neurons=64, activation="ReLU"):
+        super().__init__()
+        self.n_neurons = n_neurons
+        self.n_hidden_layers = n_hidden_layers
+        self.activation = activation
+        self.out_activation = out_activation
+        layers = [
+            self.make_linear(input_dim, self.n_neurons, is_first=True, is_last=False),
+            self.make_activation(),
+        ]
+        for i in range(self.n_hidden_layers - 1):
+            layers += [
+                self.make_linear(
+                    self.n_neurons, self.n_neurons, is_first=False, is_last=False
+                ),
+                self.make_activation(),
+            ]
+        layers += [
+            self.make_linear(self.n_neurons, output_dim, is_first=False, is_last=True)
+        ]
+        if self.out_activation == "sigmoid":
+            layers += [nn.Sigmoid()]
+        elif self.out_activation == "tanh":
+            layers += [nn.Tanh()]
+        elif self.out_activation == "hardtanh":
+            layers += [nn.Hardtanh()]
+        elif self.out_activation == "GELU":
+            layers += [nn.GELU()]            
+        elif self.out_activation == "RELU":
+            layers += [nn.ReLU()]  
+        else:
+            raise NotImplementedError
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x, split_size=100000):
+        with torch.cuda.amp.autocast(enabled=False):
+            out = self.layers(x)
+        return out
+
+    def make_linear(self, dim_in, dim_out, is_first, is_last):
+        layer = nn.Linear(dim_in, dim_out, bias=False)
+        return layer
+
+    def make_activation(self):
+        if self.activation == "ReLU":
+            return nn.ReLU(inplace=True)
+        elif self.activation == "GELU":
+            return nn.GELU()
+        else:
+            raise NotImplementedError