diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..7d715d1b374db322c19ecebf9d3cf95999f91b40 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +wheels/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl filter=lfs diff=lfs merge=lfs -text diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..261eeb9e9f8b2b4b0d119366dda99c6fd7d35c64 --- /dev/null +++ b/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. diff --git a/README.md b/README.md index 750fda748f0b4941cbfa185e6da712aa2b43c4eb..00adfa5d42ae57f01b7dd742765d2e991d65d453 100644 --- a/README.md +++ b/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`). diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md new file mode 100644 index 0000000000000000000000000000000000000000..d4dea237fdbe0947cbe391bfae8f1cd829842f9b --- /dev/null +++ b/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. diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..f30676fdef3ace0128c030417ab102b1b9aceeb6 --- /dev/null +++ b/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 ": " (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, ': ', 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) diff --git a/config/default.yaml b/config/default.yaml new file mode 100644 index 0000000000000000000000000000000000000000..54e6b8112158bbb0589b6fc881361d743584aafe --- /dev/null +++ b/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 diff --git a/config/trellis_pipeline/pipeline.json b/config/trellis_pipeline/pipeline.json new file mode 100644 index 0000000000000000000000000000000000000000..23b0f1509415a234df711aa27f03f916727819f5 --- /dev/null +++ b/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" + } +} \ No newline at end of file diff --git a/examples/chair.npz b/examples/chair.npz new file mode 100644 index 0000000000000000000000000000000000000000..57545d99f9a16866e73f5dde1f147bf54aa0b61e --- /dev/null +++ b/examples/chair.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dd111d52e0194ae0628dbda65f49cb938c71332709269b592df405f735a3a7d8 +size 4148 diff --git a/examples/hot_air_balloon.npz b/examples/hot_air_balloon.npz new file mode 100644 index 0000000000000000000000000000000000000000..937a9a07a0e62d3c3c1f5bb4d4be45c1b957c644 --- /dev/null +++ b/examples/hot_air_balloon.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1015eb538c4e613c84ac62bb7b8f3ec202d5852ab36f30d4b32ef2989cd06a88 +size 3910 diff --git a/examples/sailboat.npz b/examples/sailboat.npz new file mode 100644 index 0000000000000000000000000000000000000000..9f73ac522c03ee9d936215d247b2fc34741f52b3 --- /dev/null +++ b/examples/sailboat.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9710205f9549302bfd70de5b6f26ec2c461f2197edc89bc341f0c67f9943119b +size 3503 diff --git a/examples/toy_elephant.npz b/examples/toy_elephant.npz new file mode 100644 index 0000000000000000000000000000000000000000..81f003fd494de04f967b2226d76c3c86b8cc3b14 --- /dev/null +++ b/examples/toy_elephant.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b9c5b641efdfc1e6b16ce947b4775cb86606429dde70f1ee19a746beb57942a0 +size 2930 diff --git a/examples/toy_rocket.npz b/examples/toy_rocket.npz new file mode 100644 index 0000000000000000000000000000000000000000..f4687ee3290cb85938009f590763a4866329378b --- /dev/null +++ b/examples/toy_rocket.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0523710cdcc30a0bef3491d6239bea355a93e11a13ac752d729477bae309f6aa +size 2326 diff --git a/examples/windmill.npz b/examples/windmill.npz new file mode 100644 index 0000000000000000000000000000000000000000..2b805ae3071ee00e7e3582c22d574ccd754f3a56 --- /dev/null +++ b/examples/windmill.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:427e77c4072cbd0703f6f64fd8151bf6cba8a587a0d684f05cf580c19bcee7b9 +size 4558 diff --git a/lib/opt/self_similarity.py b/lib/opt/self_similarity.py new file mode 100644 index 0000000000000000000000000000000000000000..6050ce4636964308b89d1101646e5a47d771cfb5 --- /dev/null +++ b/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) diff --git a/lib/util/checkpoint_compat.py b/lib/util/checkpoint_compat.py new file mode 100644 index 0000000000000000000000000000000000000000..f1f86d01adc57e407166ffa37e86eaf95ae9cdd1 --- /dev/null +++ b/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 diff --git a/lib/util/common.py b/lib/util/common.py new file mode 100644 index 0000000000000000000000000000000000000000..cd06a548cbe68005f146406b0c9fef21809e910c --- /dev/null +++ b/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 \ No newline at end of file diff --git a/lib/util/generation.py b/lib/util/generation.py new file mode 100644 index 0000000000000000000000000000000000000000..617df89dc9b0e70ea62e1e4d9db51dacd2e7303a --- /dev/null +++ b/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) diff --git a/lib/util/model_revisions.py b/lib/util/model_revisions.py new file mode 100644 index 0000000000000000000000000000000000000000..9d5f5cb1ee352d3b0489cb11ea86500526508ddc --- /dev/null +++ b/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']}" diff --git a/lib/util/partfield.py b/lib/util/partfield.py new file mode 100644 index 0000000000000000000000000000000000000000..6fbda796f969dfbd8aeca1daa51def69ee6d75cb --- /dev/null +++ b/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 \ No newline at end of file diff --git a/lib/util/pipeline_compat.py b/lib/util/pipeline_compat.py new file mode 100644 index 0000000000000000000000000000000000000000..81b14ca6116aca636b6cd76c4ffadf6eb8c6b2d1 --- /dev/null +++ b/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 diff --git a/lib/util/pointcloud.py b/lib/util/pointcloud.py new file mode 100644 index 0000000000000000000000000000000000000000..ef5d72a929ba80da7b949aac093e6754bb9fe3e1 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/lib/util/render.py b/lib/util/render.py new file mode 100644 index 0000000000000000000000000000000000000000..287bbcf8722e0a27991d98281aa34d9d8fe60de2 --- /dev/null +++ b/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] diff --git a/lib/util/visualization.py b/lib/util/visualization.py new file mode 100644 index 0000000000000000000000000000000000000000..3408c9c55d9f3465113b9e6479cb0b199f3e11d0 --- /dev/null +++ b/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}") diff --git a/packages.txt b/packages.txt new file mode 100644 index 0000000000000000000000000000000000000000..444976eeb785bbb9b3b2df8de4198ddd470a1eff --- /dev/null +++ b/packages.txt @@ -0,0 +1,4 @@ +libgl1 +libglib2.0-0 +libgomp1 +libxrender1 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..1e24ff2b6328c00ddc82ccbbee2c5af50055dbbd --- /dev/null +++ b/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 diff --git a/requirements/model-revisions.json b/requirements/model-revisions.json new file mode 100644 index 0000000000000000000000000000000000000000..b7cae08cef29aaa793f8ef10381ba431dce107ed --- /dev/null +++ b/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." +} diff --git a/run_local_tau.py b/run_local_tau.py new file mode 100644 index 0000000000000000000000000000000000000000..7e4b6a2f012637568a9a1c5406f95cb8d8493eab --- /dev/null +++ b/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 /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() diff --git a/third_party/PartField/LICENSE b/third_party/PartField/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..8388f722de77ca12cac5a752999e51882096af68 --- /dev/null +++ b/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. + diff --git a/third_party/PartField/config.yaml b/third_party/PartField/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c72eb70f2836038d312e7559175865bdbf22de52 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/third_party/PartField/partfield/config/__init__.py b/third_party/PartField/partfield/config/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..39582506b85759ab473acedb5d0f15b7d7f26594 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/third_party/PartField/partfield/config/defaults.py b/third_party/PartField/partfield/config/defaults.py new file mode 100644 index 0000000000000000000000000000000000000000..be505f9655a3606d9aece9bb42e78f646d41889b --- /dev/null +++ b/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 \ No newline at end of file diff --git a/third_party/PartField/partfield/dataloader.py b/third_party/PartField/partfield/dataloader.py new file mode 100644 index 0000000000000000000000000000000000000000..3ae1256dbb42ea9656b6e80d39122d21cea1022f --- /dev/null +++ b/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() \ No newline at end of file diff --git a/third_party/PartField/partfield/model/PVCNN/conv_pointnet.py b/third_party/PartField/partfield/model/PVCNN/conv_pointnet.py new file mode 100644 index 0000000000000000000000000000000000000000..8c5c806f1e725ed9a75aeb752f3e2ae4a5c606a1 --- /dev/null +++ b/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 + + + \ No newline at end of file diff --git a/third_party/PartField/partfield/model/PVCNN/dnnlib_util.py b/third_party/PartField/partfield/model/PVCNN/dnnlib_util.py new file mode 100644 index 0000000000000000000000000000000000000000..9514fe685275a66fc83bf78fb0cf3c94952678dd --- /dev/null +++ b/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('').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 \ No newline at end of file diff --git a/third_party/PartField/partfield/model/PVCNN/encoder_pc.py b/third_party/PartField/partfield/model/PVCNN/encoder_pc.py new file mode 100644 index 0000000000000000000000000000000000000000..25a384ec0b7b8ce19d0336f1edd2463b8b3acd4e --- /dev/null +++ b/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) \ No newline at end of file diff --git a/third_party/PartField/partfield/model/PVCNN/pc_encoder.py b/third_party/PartField/partfield/model/PVCNN/pc_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..adeaeba96147a24515eb289b09da77bc4716869c --- /dev/null +++ b/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 \ No newline at end of file diff --git a/third_party/PartField/partfield/model/PVCNN/pv_module/__init__.py b/third_party/PartField/partfield/model/PVCNN/pv_module/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7fd32e598f709503f4e35171e09fbbedec05f9c3 --- /dev/null +++ b/third_party/PartField/partfield/model/PVCNN/pv_module/__init__.py @@ -0,0 +1,2 @@ +from .pvconv import PVConv +from .shared_mlp import SharedMLP \ No newline at end of file diff --git a/third_party/PartField/partfield/model/PVCNN/pv_module/ball_query.py b/third_party/PartField/partfield/model/PVCNN/pv_module/ball_query.py new file mode 100644 index 0000000000000000000000000000000000000000..ea2a8203baa3fc1c94959f1ba852839365bf38ce --- /dev/null +++ b/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 '') diff --git a/third_party/PartField/partfield/model/PVCNN/pv_module/frustum.py b/third_party/PartField/partfield/model/PVCNN/pv_module/frustum.py new file mode 100644 index 0000000000000000000000000000000000000000..fb302963a6472f949f4ab69ed42575d79b68b4ea --- /dev/null +++ b/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) diff --git a/third_party/PartField/partfield/model/PVCNN/pv_module/functional/__init__.py b/third_party/PartField/partfield/model/PVCNN/pv_module/functional/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..993d1d12511dce369d781b60be75e79a71762e47 --- /dev/null +++ b/third_party/PartField/partfield/model/PVCNN/pv_module/functional/__init__.py @@ -0,0 +1 @@ +from .devoxelization import trilinear_devoxelize diff --git a/third_party/PartField/partfield/model/PVCNN/pv_module/functional/devoxelization.py b/third_party/PartField/partfield/model/PVCNN/pv_module/functional/devoxelization.py new file mode 100644 index 0000000000000000000000000000000000000000..c60dab12d804ec3b41b53f7da3eecb20917077fc --- /dev/null +++ b/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 diff --git a/third_party/PartField/partfield/model/PVCNN/pv_module/loss.py b/third_party/PartField/partfield/model/PVCNN/pv_module/loss.py new file mode 100644 index 0000000000000000000000000000000000000000..a35cdd8a0fe83c8ca6b1d7040b66d142e76471df --- /dev/null +++ b/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) diff --git a/third_party/PartField/partfield/model/PVCNN/pv_module/pointnet.py b/third_party/PartField/partfield/model/PVCNN/pv_module/pointnet.py new file mode 100644 index 0000000000000000000000000000000000000000..e58e01cc2f84925d4817a8a04aefdbc3d36d484e --- /dev/null +++ b/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 diff --git a/third_party/PartField/partfield/model/PVCNN/pv_module/pvconv.py b/third_party/PartField/partfield/model/PVCNN/pv_module/pvconv.py new file mode 100644 index 0000000000000000000000000000000000000000..a64705da194cf2d32ff641025fad7b92d71dc67b --- /dev/null +++ b/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 diff --git a/third_party/PartField/partfield/model/PVCNN/pv_module/shared_mlp.py b/third_party/PartField/partfield/model/PVCNN/pv_module/shared_mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..e1d4ff864c05b894194ef11ac4b629ec72c4952b --- /dev/null +++ b/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) diff --git a/third_party/PartField/partfield/model/PVCNN/pv_module/voxelization.py b/third_party/PartField/partfield/model/PVCNN/pv_module/voxelization.py new file mode 100644 index 0000000000000000000000000000000000000000..15535abc63f7adebcd12bc436e84710c7c9862d2 --- /dev/null +++ b/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 '') diff --git a/third_party/PartField/partfield/model/PVCNN/unet_3daware.py b/third_party/PartField/partfield/model/PVCNN/unet_3daware.py new file mode 100644 index 0000000000000000000000000000000000000000..b0084f0c1d6989ae4ad103f364401c2f2bd5e361 --- /dev/null +++ b/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 + diff --git a/third_party/PartField/partfield/model/UNet/buildingblocks.py b/third_party/PartField/partfield/model/UNet/buildingblocks.py new file mode 100644 index 0000000000000000000000000000000000000000..e97f501d1813b03555dbec5658d024e06d761443 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/third_party/PartField/partfield/model/UNet/model.py b/third_party/PartField/partfield/model/UNet/model.py new file mode 100644 index 0000000000000000000000000000000000000000..db20b2f5de3d37a52f7465450f915e003ef412d6 --- /dev/null +++ b/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) + + diff --git a/third_party/PartField/partfield/model/model_utils.py b/third_party/PartField/partfield/model/model_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c1cc16bd96d21b69f21998e577f7f5f970d25cf3 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/third_party/PartField/partfield/model/triplane.py b/third_party/PartField/partfield/model/triplane.py new file mode 100644 index 0000000000000000000000000000000000000000..6274a8398d248d3ba4a5a4734c7b1bc90d596b10 --- /dev/null +++ b/third_party/PartField/partfield/model/triplane.py @@ -0,0 +1,331 @@ +#https://github.com/3DTopia/OpenLRM/blob/main/openlrm/models/modeling_lrm.py +# Copyright (c) 2023-2024, Zexin He +# +# 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 +# +# https://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. + +import torch +import torch.nn as nn +from functools import partial + +def project_onto_planes(planes, coordinates): + """ + Does a projection of a 3D point onto a batch of 2D planes, + returning 2D plane coordinates. + + Takes plane axes of shape n_planes, 3, 3 + # Takes coordinates of shape N, M, 3 + # returns projections of shape N*n_planes, M, 2 + """ + N, M, C = coordinates.shape + n_planes, _, _ = planes.shape + coordinates = coordinates.unsqueeze(1).expand(-1, n_planes, -1, -1).reshape(N*n_planes, M, 3) + inv_planes = torch.linalg.inv(planes).unsqueeze(0).expand(N, -1, -1, -1).reshape(N*n_planes, 3, 3) + projections = torch.bmm(coordinates, inv_planes) + return projections[..., :2] + +def sample_from_planes(plane_features, coordinates, mode='bilinear', padding_mode='zeros', box_warp=None): + plane_axes = torch.tensor([[[1, 0, 0], + [0, 1, 0], + [0, 0, 1]], + [[1, 0, 0], + [0, 0, 1], + [0, 1, 0]], + [[0, 0, 1], + [0, 1, 0], + [1, 0, 0]]], dtype=torch.float32).cuda() + + assert padding_mode == 'zeros' + N, n_planes, C, H, W = plane_features.shape + _, M, _ = coordinates.shape + plane_features = plane_features.view(N*n_planes, C, H, W) + + projected_coordinates = project_onto_planes(plane_axes, coordinates).unsqueeze(1) + output_features = torch.nn.functional.grid_sample(plane_features, projected_coordinates.float(), mode=mode, padding_mode=padding_mode, align_corners=False).permute(0, 3, 2, 1).reshape(N, n_planes, M, C) + return output_features + +def get_grid_coord(grid_size = 256, align_corners=False): + if align_corners == False: + coords = torch.linspace(-1 + 1/(grid_size), 1 - 1/(grid_size), steps=grid_size) + else: + coords = torch.linspace(-1, 1, steps=grid_size) + i, j, k = torch.meshgrid(coords, coords, coords, indexing='ij') + coordinates = torch.stack((i, j, k), dim=-1).reshape(-1, 3) + return coordinates + +class BasicBlock(nn.Module): + """ + Transformer block that is in its simplest form. + Designed for PF-LRM architecture. + """ + # Block contains a self-attention layer and an MLP + def __init__(self, inner_dim: int, num_heads: int, eps: float, + attn_drop: float = 0., attn_bias: bool = False, + mlp_ratio: float = 4., mlp_drop: float = 0.): + super().__init__() + self.norm1 = nn.LayerNorm(inner_dim, eps=eps) + self.self_attn = nn.MultiheadAttention( + embed_dim=inner_dim, num_heads=num_heads, + dropout=attn_drop, bias=attn_bias, batch_first=True) + self.norm2 = nn.LayerNorm(inner_dim, eps=eps) + self.mlp = nn.Sequential( + nn.Linear(inner_dim, int(inner_dim * mlp_ratio)), + nn.GELU(), + nn.Dropout(mlp_drop), + nn.Linear(int(inner_dim * mlp_ratio), inner_dim), + nn.Dropout(mlp_drop), + ) + + def forward(self, x): + # x: [N, L, D] + before_sa = self.norm1(x) + x = x + self.self_attn(before_sa, before_sa, before_sa, need_weights=False)[0] + x = x + self.mlp(self.norm2(x)) + return x + +class ConditionBlock(nn.Module): + """ + Transformer block that takes in a cross-attention condition. + Designed for SparseLRM architecture. + """ + # Block contains a cross-attention layer, a self-attention layer, and an MLP + def __init__(self, inner_dim: int, cond_dim: int, num_heads: int, eps: float, + attn_drop: float = 0., attn_bias: bool = False, + mlp_ratio: float = 4., mlp_drop: float = 0.): + super().__init__() + self.norm1 = nn.LayerNorm(inner_dim, eps=eps) + self.cross_attn = nn.MultiheadAttention( + embed_dim=inner_dim, num_heads=num_heads, kdim=cond_dim, vdim=cond_dim, + dropout=attn_drop, bias=attn_bias, batch_first=True) + self.norm2 = nn.LayerNorm(inner_dim, eps=eps) + self.self_attn = nn.MultiheadAttention( + embed_dim=inner_dim, num_heads=num_heads, + dropout=attn_drop, bias=attn_bias, batch_first=True) + self.norm3 = nn.LayerNorm(inner_dim, eps=eps) + self.mlp = nn.Sequential( + nn.Linear(inner_dim, int(inner_dim * mlp_ratio)), + nn.GELU(), + nn.Dropout(mlp_drop), + nn.Linear(int(inner_dim * mlp_ratio), inner_dim), + nn.Dropout(mlp_drop), + ) + + def forward(self, x, cond): + # x: [N, L, D] + # cond: [N, L_cond, D_cond] + x = x + self.cross_attn(self.norm1(x), cond, cond, need_weights=False)[0] + before_sa = self.norm2(x) + x = x + self.self_attn(before_sa, before_sa, before_sa, need_weights=False)[0] + x = x + self.mlp(self.norm3(x)) + return x + +class TransformerDecoder(nn.Module): + def __init__(self, block_type: str, + num_layers: int, num_heads: int, + inner_dim: int, cond_dim: int = None, + eps: float = 1e-6): + super().__init__() + self.block_type = block_type + self.layers = nn.ModuleList([ + self._block_fn(inner_dim, cond_dim)( + num_heads=num_heads, + eps=eps, + ) + for _ in range(num_layers) + ]) + self.norm = nn.LayerNorm(inner_dim, eps=eps) + + @property + def block_type(self): + return self._block_type + + @block_type.setter + def block_type(self, block_type): + assert block_type in ['cond', 'basic'], \ + f"Unsupported block type: {block_type}" + self._block_type = block_type + + def _block_fn(self, inner_dim, cond_dim): + assert inner_dim is not None, f"inner_dim must always be specified" + if self.block_type == 'basic': + return partial(BasicBlock, inner_dim=inner_dim) + elif self.block_type == 'cond': + assert cond_dim is not None, f"Condition dimension must be specified for ConditionBlock" + return partial(ConditionBlock, inner_dim=inner_dim, cond_dim=cond_dim) + else: + raise ValueError(f"Unsupported block type during runtime: {self.block_type}") + + + def forward_layer(self, layer: nn.Module, x: torch.Tensor, cond: torch.Tensor,): + if self.block_type == 'basic': + return layer(x) + elif self.block_type == 'cond': + return layer(x, cond) + else: + raise NotImplementedError + + def forward(self, x: torch.Tensor, cond: torch.Tensor = None): + # x: [N, L, D] + # cond: [N, L_cond, D_cond] or None + for layer in self.layers: + x = self.forward_layer(layer, x, cond) + x = self.norm(x) + return x + +class Voxel2Triplane(nn.Module): + """ + Full model of the basic single-view large reconstruction model. + """ + def __init__(self, transformer_dim: int, transformer_layers: int, transformer_heads: int, + triplane_low_res: int, triplane_high_res: int, triplane_dim: int, voxel_feat_dim: int, normalize_vox_feat=False, voxel_dim=16): + super().__init__() + + # attributes + self.triplane_low_res = triplane_low_res + self.triplane_high_res = triplane_high_res + self.triplane_dim = triplane_dim + self.voxel_feat_dim = voxel_feat_dim + + # initialize pos_embed with 1/sqrt(dim) * N(0, 1) + self.pos_embed = nn.Parameter(torch.randn(1, 3*triplane_low_res**2, transformer_dim) * (1. / transformer_dim) ** 0.5) + self.transformer = TransformerDecoder( + block_type='cond', + num_layers=transformer_layers, num_heads=transformer_heads, + inner_dim=transformer_dim, cond_dim=voxel_feat_dim + ) + self.upsampler = nn.ConvTranspose2d(transformer_dim, triplane_dim, kernel_size=8, stride=8, padding=0) + + self.normalize_vox_feat = normalize_vox_feat + if normalize_vox_feat: + self.vox_norm = nn.LayerNorm(voxel_feat_dim, eps=1e-6) + self.vox_pos_embed = nn.Parameter(torch.randn(1, voxel_dim * voxel_dim * voxel_dim, voxel_feat_dim) * (1. / voxel_feat_dim) ** 0.5) + + def forward_transformer(self, voxel_feats): + N = voxel_feats.shape[0] + x = self.pos_embed.repeat(N, 1, 1) # [N, L, D] + if self.normalize_vox_feat: + vox_pos_embed = self.vox_pos_embed.repeat(N, 1, 1) # [N, L, D] + voxel_feats = self.vox_norm(voxel_feats + vox_pos_embed) + x = self.transformer( + x, + cond=voxel_feats + ) + return x + + def reshape_upsample(self, tokens): + N = tokens.shape[0] + H = W = self.triplane_low_res + x = tokens.view(N, 3, H, W, -1) + x = torch.einsum('nihwd->indhw', x) # [3, N, D, H, W] + x = x.contiguous().view(3*N, -1, H, W) # [3*N, D, H, W] + x = self.upsampler(x) # [3*N, D', H', W'] + x = x.view(3, N, *x.shape[-3:]) # [3, N, D', H', W'] + x = torch.einsum('indhw->nidhw', x) # [N, 3, D', H', W'] + x = x.contiguous() + return x + + def forward(self, voxel_feats): + N = voxel_feats.shape[0] + + # encode image + assert voxel_feats.shape[-1] == self.voxel_feat_dim, \ + f"Feature dimension mismatch: {voxel_feats.shape[-1]} vs {self.voxel_feat_dim}" + + # transformer generating planes + tokens = self.forward_transformer(voxel_feats) + planes = self.reshape_upsample(tokens) + assert planes.shape[0] == N, "Batch size mismatch for planes" + assert planes.shape[1] == 3, "Planes should have 3 channels" + + return planes + + +class TriplaneTransformer(nn.Module): + """ + Full model of the basic single-view large reconstruction model. + """ + def __init__(self, input_dim: int, transformer_dim: int, transformer_layers: int, transformer_heads: int, + triplane_low_res: int, triplane_high_res: int, triplane_dim: int): + super().__init__() + + # attributes + self.triplane_low_res = triplane_low_res + self.triplane_high_res = triplane_high_res + self.triplane_dim = triplane_dim + + # initialize pos_embed with 1/sqrt(dim) * N(0, 1) + self.pos_embed = nn.Parameter(torch.randn(1, 3*triplane_low_res**2, transformer_dim) * (1. / transformer_dim) ** 0.5) + self.transformer = TransformerDecoder( + block_type='basic', + num_layers=transformer_layers, num_heads=transformer_heads, + inner_dim=transformer_dim, + ) + + self.downsampler = nn.Sequential( + nn.Conv2d(input_dim, transformer_dim, kernel_size=3, stride=1, padding=1), + nn.ReLU(), + nn.MaxPool2d(kernel_size=2, stride=2), # Reduces size from 128x128 to 64x64 + + nn.Conv2d(transformer_dim, transformer_dim, kernel_size=3, stride=1, padding=1), + nn.ReLU(), + nn.MaxPool2d(kernel_size=2, stride=2), # Reduces size from 64x64 to 32x32 + ) + + self.upsampler = nn.ConvTranspose2d(transformer_dim, triplane_dim, kernel_size=4, stride=4, padding=0) + + self.mlp = nn.Sequential( + nn.Linear(input_dim, triplane_dim), + nn.ReLU(), + nn.Linear(triplane_dim, triplane_dim) + ) + + def forward_transformer(self, triplanes): + N = triplanes.shape[0] + tokens = torch.einsum('nidhw->nihwd', triplanes).reshape(N, self.pos_embed.shape[1], -1) # [N, L, D] + x = self.pos_embed.repeat(N, 1, 1) + tokens # [N, L, D] + x = self.transformer(x) + return x + + def reshape_downsample(self, triplanes): + N = triplanes.shape[0] + H = W = self.triplane_high_res + x = triplanes.view(N, 3, -1, H, W) + x = torch.einsum('nidhw->indhw', x) # [3, N, D, H, W] + x = x.contiguous().view(3*N, -1, H, W) # [3*N, D, H, W] + x = self.downsampler(x) # [3*N, D', H', W'] + x = x.view(3, N, *x.shape[-3:]) # [3, N, D', H', W'] + x = torch.einsum('indhw->nidhw', x) # [N, 3, D', H', W'] + x = x.contiguous() + return x + + def reshape_upsample(self, tokens): + N = tokens.shape[0] + H = W = self.triplane_low_res + x = tokens.view(N, 3, H, W, -1) + x = torch.einsum('nihwd->indhw', x) # [3, N, D, H, W] + x = x.contiguous().view(3*N, -1, H, W) # [3*N, D, H, W] + x = self.upsampler(x) # [3*N, D', H', W'] + x = x.view(3, N, *x.shape[-3:]) # [3, N, D', H', W'] + x = torch.einsum('indhw->nidhw', x) # [N, 3, D', H', W'] + x = x.contiguous() + return x + + def forward(self, triplanes): + downsampled_triplanes = self.reshape_downsample(triplanes) + tokens = self.forward_transformer(downsampled_triplanes) + residual = self.reshape_upsample(tokens) + + triplanes = triplanes.permute(0, 1, 3, 4, 2).contiguous() + triplanes = self.mlp(triplanes) + triplanes = triplanes.permute(0, 1, 4, 2, 3).contiguous() + planes = triplanes + residual + return planes diff --git a/third_party/PartField/partfield/model_trainer_pvcnn_only_demo.py b/third_party/PartField/partfield/model_trainer_pvcnn_only_demo.py new file mode 100644 index 0000000000000000000000000000000000000000..f78a0c3b965117bbdf853d49c2c0a3aa4b16398c --- /dev/null +++ b/third_party/PartField/partfield/model_trainer_pvcnn_only_demo.py @@ -0,0 +1,117 @@ +import torch +import lightning.pytorch as pl +import torch.nn as nn +import os +import trimesh +import numpy as np + +from torch.utils.data import DataLoader + +from third_party.PartField.partfield.model.PVCNN.encoder_pc import TriPlanePC2Encoder, sample_triplane_feat +from third_party.PartField.partfield.model.triplane import TriplaneTransformer, get_grid_coord #, sample_from_planes, Voxel2Triplane +from third_party.PartField.partfield.model.model_utils import VanillaMLP +from third_party.PartField.partfield.dataloader import Demo_Dataset + +class Model(pl.LightningModule): + def __init__(self, cfg, obj_path): + super().__init__() + self.obj_path = obj_path + + self.save_hyperparameters() + self.cfg = cfg + self.automatic_optimization = False + self.triplane_resolution = cfg.triplane_resolution + self.triplane_channels_low = cfg.triplane_channels_low + self.triplane_transformer = TriplaneTransformer( + input_dim=cfg.triplane_channels_low * 2, + transformer_dim=1024, + transformer_layers=6, + transformer_heads=8, + triplane_low_res=32, + triplane_high_res=128, + triplane_dim=cfg.triplane_channels_high, + ) + self.sdf_decoder = VanillaMLP(input_dim=64, + output_dim=1, + out_activation="tanh", + n_neurons=64, #64 + n_hidden_layers=6) #6 + self.use_pvcnn = cfg.use_pvcnnonly + self.use_2d_feat = cfg.use_2d_feat + if self.use_pvcnn: + self.pvcnn = TriPlanePC2Encoder( + cfg.pvcnn, + device="cuda", + shape_min=-1, + shape_length=2, + use_2d_feat=self.use_2d_feat) #.cuda() + self.logit_scale = nn.Parameter(torch.tensor([1.0], requires_grad=True)) + self.grid_coord = get_grid_coord(256) + self.mse_loss = torch.nn.MSELoss() + self.l1_loss = torch.nn.L1Loss(reduction='none') + + if cfg.regress_2d_feat: + self.feat_decoder = VanillaMLP(input_dim=64, + output_dim=192, + out_activation="GELU", + n_neurons=64, #64 + n_hidden_layers=6) #6 + + def predict_dataloader(self): + dataset = Demo_Dataset(self.obj_path) + dataloader = DataLoader(dataset, + num_workers=self.cfg.dataset.val_num_workers, + batch_size=self.cfg.dataset.val_batch_size, + shuffle=False, + pin_memory=True, + drop_last=False) + + return dataloader + + @torch.no_grad() + def predict_step(self, batch, batch_idx): + N = batch['pc'].shape[0] + + assert N == 1 + + pc_feat = self.pvcnn(batch['pc'], batch['pc']) + + planes = pc_feat + planes = self.triplane_transformer(planes) + sdf_planes, part_planes = torch.split(planes, [64, planes.shape[2] - 64], dim=2) + + def sample_points(vertices, faces, n_point_per_face): + # Generate random barycentric coordinates + # borrowed from Kaolin https://github.com/NVIDIAGameWorks/kaolin/blob/master/kaolin/ops/mesh/trianglemesh.py#L43 + n_f = faces.shape[0] + u = torch.sqrt(torch.rand((n_f, n_point_per_face, 1), + device=vertices.device, + dtype=vertices.dtype)) + v = torch.rand((n_f, n_point_per_face, 1), + device=vertices.device, + dtype=vertices.dtype) + w0 = 1 - u + w1 = u * (1 - v) + w2 = u * v + + face_v_0 = torch.index_select(vertices, 0, faces[:, 0].reshape(-1)) + face_v_1 = torch.index_select(vertices, 0, faces[:, 1].reshape(-1)) + face_v_2 = torch.index_select(vertices, 0, faces[:, 2].reshape(-1)) + points = w0 * face_v_0.unsqueeze(dim=1) + w1 * face_v_1.unsqueeze(dim=1) + w2 * face_v_2.unsqueeze(dim=1) + return points + + def sample_and_mean_memory_save_version(part_planes, tensor_vertices, n_point_per_face): + n_sample_each = self.cfg.n_sample_each # we iterate over this to avoid OOM + n_v = tensor_vertices.shape[1] + n_sample = n_v // n_sample_each + 1 + all_sample = [] + for i_sample in range(n_sample): + sampled_feature = sample_triplane_feat(part_planes, tensor_vertices[:, i_sample * n_sample_each: i_sample * n_sample_each + n_sample_each,]) + assert sampled_feature.shape[1] % n_point_per_face == 0 + sampled_feature = sampled_feature.reshape(1, -1, n_point_per_face, sampled_feature.shape[-1]) + sampled_feature = torch.mean(sampled_feature, axis=-2) + all_sample.append(sampled_feature) + return torch.cat(all_sample, dim=1) + + part_planes = part_planes.cpu().numpy() + return part_planes, batch['uid'][0] \ No newline at end of file diff --git a/third_party/PartField/partfield/utils.py b/third_party/PartField/partfield/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8dc176434f91bd753f2daee6cc7d3c4fd7ce163b --- /dev/null +++ b/third_party/PartField/partfield/utils.py @@ -0,0 +1,5 @@ +import trimesh + +def load_mesh_util(input_fname): + mesh = trimesh.load(input_fname, force='mesh', process=False) + return mesh \ No newline at end of file diff --git a/third_party/TRELLIS/LICENSE b/third_party/TRELLIS/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..9e841e7a26e4eb057b24511e7b92d42b257a80e5 --- /dev/null +++ b/third_party/TRELLIS/LICENSE @@ -0,0 +1,21 @@ + MIT License + + Copyright (c) Microsoft Corporation. + + Permission is hereby granted, free of charge, to any person obtaining a copy + of this software and associated documentation files (the "Software"), to deal + in the Software without restriction, including without limitation the rights + to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + copies of the Software, and to permit persons to whom the Software is + furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in all + copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE diff --git a/third_party/TRELLIS/trellis/__init__.py b/third_party/TRELLIS/trellis/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..20d240afc9c26a21aee76954628b3d4ef9a1ccbd --- /dev/null +++ b/third_party/TRELLIS/trellis/__init__.py @@ -0,0 +1,6 @@ +from . import models +from . import modules +from . import pipelines +from . import renderers +from . import representations +from . import utils diff --git a/third_party/TRELLIS/trellis/datasets/__init__.py b/third_party/TRELLIS/trellis/datasets/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6798ca08a4bd855ba91fad3a5f85aef0146ff919 --- /dev/null +++ b/third_party/TRELLIS/trellis/datasets/__init__.py @@ -0,0 +1,58 @@ +import importlib + +__attributes = { + 'SparseStructure': 'sparse_structure', + + 'SparseFeat2Render': 'sparse_feat2render', + 'SLat2Render':'structured_latent2render', + 'Slat2RenderGeo':'structured_latent2render', + + 'SparseStructureLatent': 'sparse_structure_latent', + 'TextConditionedSparseStructureLatent': 'sparse_structure_latent', + 'ImageConditionedSparseStructureLatent': 'sparse_structure_latent', + + 'SLat': 'structured_latent', + 'TextConditionedSLat': 'structured_latent', + 'ImageConditionedSLat': 'structured_latent', +} + +__submodules = [] + +__all__ = list(__attributes.keys()) + __submodules + +def __getattr__(name): + if name not in globals(): + if name in __attributes: + module_name = __attributes[name] + module = importlib.import_module(f".{module_name}", __name__) + globals()[name] = getattr(module, name) + elif name in __submodules: + module = importlib.import_module(f".{name}", __name__) + globals()[name] = module + else: + raise AttributeError(f"module {__name__} has no attribute {name}") + return globals()[name] + + +# For Pylance +if __name__ == '__main__': + from .sparse_structure import SparseStructure + + from .sparse_feat2render import SparseFeat2Render + from .structured_latent2render import ( + SLat2Render, + Slat2RenderGeo, + ) + + from .sparse_structure_latent import ( + SparseStructureLatent, + TextConditionedSparseStructureLatent, + ImageConditionedSparseStructureLatent, + ) + + from .structured_latent import ( + SLat, + TextConditionedSLat, + ImageConditionedSLat, + ) + \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/datasets/components.py b/third_party/TRELLIS/trellis/datasets/components.py new file mode 100644 index 0000000000000000000000000000000000000000..9c863d6243324ec5acf17aac4f50d55b1335b3a1 --- /dev/null +++ b/third_party/TRELLIS/trellis/datasets/components.py @@ -0,0 +1,137 @@ +from typing import * +from abc import abstractmethod +import os +import json +import torch +import numpy as np +import pandas as pd +from PIL import Image +from torch.utils.data import Dataset + + +class StandardDatasetBase(Dataset): + """ + Base class for standard datasets. + + Args: + roots (str): paths to the dataset + """ + + def __init__(self, + roots: str, + ): + super().__init__() + self.roots = roots.split(',') + self.instances = [] + self.metadata = pd.DataFrame() + + self._stats = {} + for root in self.roots: + key = os.path.basename(root) + self._stats[key] = {} + metadata = pd.read_csv(os.path.join(root, 'metadata.csv')) + self._stats[key]['Total'] = len(metadata) + metadata, stats = self.filter_metadata(metadata) + self._stats[key].update(stats) + self.instances.extend([(root, sha256) for sha256 in metadata['sha256'].values]) + metadata.set_index('sha256', inplace=True) + self.metadata = pd.concat([self.metadata, metadata]) + + @abstractmethod + def filter_metadata(self, metadata: pd.DataFrame) -> Tuple[pd.DataFrame, Dict[str, int]]: + pass + + @abstractmethod + def get_instance(self, root: str, instance: str) -> Dict[str, Any]: + pass + + def __len__(self): + return len(self.instances) + + def __getitem__(self, index) -> Dict[str, Any]: + try: + root, instance = self.instances[index] + return self.get_instance(root, instance) + except Exception as e: + print(e) + return self.__getitem__(np.random.randint(0, len(self))) + + def __str__(self): + lines = [] + lines.append(self.__class__.__name__) + lines.append(f' - Total instances: {len(self)}') + lines.append(f' - Sources:') + for key, stats in self._stats.items(): + lines.append(f' - {key}:') + for k, v in stats.items(): + lines.append(f' - {k}: {v}') + return '\n'.join(lines) + + +class TextConditionedMixin: + def __init__(self, roots, **kwargs): + super().__init__(roots, **kwargs) + self.captions = {} + for instance in self.instances: + sha256 = instance[1] + self.captions[sha256] = json.loads(self.metadata.loc[sha256]['captions']) + + def filter_metadata(self, metadata): + metadata, stats = super().filter_metadata(metadata) + metadata = metadata[metadata['captions'].notna()] + stats['With captions'] = len(metadata) + return metadata, stats + + def get_instance(self, root, instance): + pack = super().get_instance(root, instance) + text = np.random.choice(self.captions[instance]) + pack['cond'] = text + return pack + + +class ImageConditionedMixin: + def __init__(self, roots, *, image_size=518, **kwargs): + self.image_size = image_size + super().__init__(roots, **kwargs) + + def filter_metadata(self, metadata): + metadata, stats = super().filter_metadata(metadata) + metadata = metadata[metadata[f'cond_rendered']] + stats['Cond rendered'] = len(metadata) + return metadata, stats + + def get_instance(self, root, instance): + pack = super().get_instance(root, instance) + + image_root = os.path.join(root, 'renders_cond', instance) + with open(os.path.join(image_root, 'transforms.json')) as f: + metadata = json.load(f) + n_views = len(metadata['frames']) + view = np.random.randint(n_views) + metadata = metadata['frames'][view] + + image_path = os.path.join(image_root, metadata['file_path']) + image = Image.open(image_path) + + alpha = np.array(image.getchannel(3)) + bbox = np.array(alpha).nonzero() + bbox = [bbox[1].min(), bbox[0].min(), bbox[1].max(), bbox[0].max()] + center = [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2] + hsize = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 2 + aug_size_ratio = 1.2 + aug_hsize = hsize * aug_size_ratio + aug_center_offset = [0, 0] + aug_center = [center[0] + aug_center_offset[0], center[1] + aug_center_offset[1]] + aug_bbox = [int(aug_center[0] - aug_hsize), int(aug_center[1] - aug_hsize), int(aug_center[0] + aug_hsize), int(aug_center[1] + aug_hsize)] + image = image.crop(aug_bbox) + + image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS) + alpha = image.getchannel(3) + image = image.convert('RGB') + image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0 + alpha = torch.tensor(np.array(alpha)).float() / 255.0 + image = image * alpha.unsqueeze(0) + pack['cond'] = image + + return pack + \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/datasets/sparse_feat2render.py b/third_party/TRELLIS/trellis/datasets/sparse_feat2render.py new file mode 100644 index 0000000000000000000000000000000000000000..65bf488aeaab5ace5d55612dafc9d4959dc154c6 --- /dev/null +++ b/third_party/TRELLIS/trellis/datasets/sparse_feat2render.py @@ -0,0 +1,134 @@ +import os +from PIL import Image +import json +import numpy as np +import pandas as pd +import torch +import utils3d.torch +from ..modules.sparse.basic import SparseTensor +from .components import StandardDatasetBase + + +class SparseFeat2Render(StandardDatasetBase): + """ + SparseFeat2Render dataset. + + Args: + roots (str): paths to the dataset + image_size (int): size of the image + model (str): model name + resolution (int): resolution of the data + min_aesthetic_score (float): minimum aesthetic score + max_num_voxels (int): maximum number of voxels + """ + def __init__( + self, + roots: str, + image_size: int, + model: str = 'dinov2_vitl14_reg', + resolution: int = 64, + min_aesthetic_score: float = 5.0, + max_num_voxels: int = 32768, + ): + self.image_size = image_size + self.model = model + self.resolution = resolution + self.min_aesthetic_score = min_aesthetic_score + self.max_num_voxels = max_num_voxels + self.value_range = (0, 1) + + super().__init__(roots) + + def filter_metadata(self, metadata): + stats = {} + metadata = metadata[metadata[f'feature_{self.model}']] + stats['With features'] = len(metadata) + metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score] + stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata) + metadata = metadata[metadata['num_voxels'] <= self.max_num_voxels] + stats[f'Num voxels <= {self.max_num_voxels}'] = len(metadata) + return metadata, stats + + def _get_image(self, root, instance): + with open(os.path.join(root, 'renders', instance, 'transforms.json')) as f: + metadata = json.load(f) + n_views = len(metadata['frames']) + view = np.random.randint(n_views) + metadata = metadata['frames'][view] + fov = metadata['camera_angle_x'] + intrinsics = utils3d.torch.intrinsics_from_fov_xy(torch.tensor(fov), torch.tensor(fov)) + c2w = torch.tensor(metadata['transform_matrix']) + c2w[:3, 1:3] *= -1 + extrinsics = torch.inverse(c2w) + + image_path = os.path.join(root, 'renders', instance, metadata['file_path']) + image = Image.open(image_path) + alpha = image.getchannel(3) + image = image.convert('RGB') + image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS) + alpha = alpha.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS) + image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0 + alpha = torch.tensor(np.array(alpha)).float() / 255.0 + + return { + 'image': image, + 'alpha': alpha, + 'extrinsics': extrinsics, + 'intrinsics': intrinsics, + } + + def _get_feat(self, root, instance): + DATA_RESOLUTION = 64 + feats_path = os.path.join(root, 'features', self.model, f'{instance}.npz') + feats = np.load(feats_path, allow_pickle=True) + coords = torch.tensor(feats['indices']).int() + feats = torch.tensor(feats['patchtokens']).float() + + if self.resolution != DATA_RESOLUTION: + factor = DATA_RESOLUTION // self.resolution + coords = coords // factor + coords, idx = coords.unique(return_inverse=True, dim=0) + feats = torch.scatter_reduce( + torch.zeros(coords.shape[0], feats.shape[1], device=feats.device), + dim=0, + index=idx.unsqueeze(-1).expand(-1, feats.shape[1]), + src=feats, + reduce='mean' + ) + + return { + 'coords': coords, + 'feats': feats, + } + + @torch.no_grad() + def visualize_sample(self, sample: dict): + return sample['image'] + + @staticmethod + def collate_fn(batch): + pack = {} + coords = [] + for i, b in enumerate(batch): + coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1)) + coords = torch.cat(coords) + feats = torch.cat([b['feats'] for b in batch]) + pack['feats'] = SparseTensor( + coords=coords, + feats=feats, + ) + + pack['image'] = torch.stack([b['image'] for b in batch]) + pack['alpha'] = torch.stack([b['alpha'] for b in batch]) + pack['extrinsics'] = torch.stack([b['extrinsics'] for b in batch]) + pack['intrinsics'] = torch.stack([b['intrinsics'] for b in batch]) + + return pack + + def get_instance(self, root, instance): + image = self._get_image(root, instance) + feat = self._get_feat(root, instance) + return { + **image, + **feat, + } diff --git a/third_party/TRELLIS/trellis/datasets/sparse_structure.py b/third_party/TRELLIS/trellis/datasets/sparse_structure.py new file mode 100644 index 0000000000000000000000000000000000000000..863fd7355621a077c2dabab11779416172824027 --- /dev/null +++ b/third_party/TRELLIS/trellis/datasets/sparse_structure.py @@ -0,0 +1,107 @@ +import os +import json +from typing import Union +import numpy as np +import pandas as pd +import torch +from torch.utils.data import Dataset +import utils3d +from .components import StandardDatasetBase +from ..representations.octree import DfsOctree as Octree +from ..renderers import OctreeRenderer + + +class SparseStructure(StandardDatasetBase): + """ + Sparse structure dataset + + Args: + roots (str): path to the dataset + resolution (int): resolution of the voxel grid + min_aesthetic_score (float): minimum aesthetic score of the instances to be included in the dataset + """ + + def __init__(self, + roots, + resolution: int = 64, + min_aesthetic_score: float = 5.0, + ): + self.resolution = resolution + self.min_aesthetic_score = min_aesthetic_score + self.value_range = (0, 1) + + super().__init__(roots) + + def filter_metadata(self, metadata): + stats = {} + metadata = metadata[metadata[f'voxelized']] + stats['Voxelized'] = len(metadata) + metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score] + stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata) + return metadata, stats + + def get_instance(self, root, instance): + position = utils3d.io.read_ply(os.path.join(root, 'voxels', f'{instance}.ply'))[0] + coords = ((torch.tensor(position) + 0.5) * self.resolution).int().contiguous() + ss = torch.zeros(1, self.resolution, self.resolution, self.resolution, dtype=torch.long) + ss[:, coords[:, 0], coords[:, 1], coords[:, 2]] = 1 + return {'ss': ss} + + @torch.no_grad() + def visualize_sample(self, ss: Union[torch.Tensor, dict]): + ss = ss if isinstance(ss, torch.Tensor) else ss['ss'] + + renderer = OctreeRenderer() + renderer.rendering_options.resolution = 512 + renderer.rendering_options.near = 0.8 + renderer.rendering_options.far = 1.6 + renderer.rendering_options.bg_color = (0, 0, 0) + renderer.rendering_options.ssaa = 4 + renderer.pipe.primitive = 'voxel' + + # Build camera + yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2] + yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4) + yaws = [y + yaws_offset for y in yaws] + pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)] + + exts = [] + ints = [] + for yaw, pitch in zip(yaws, pitch): + orig = torch.tensor([ + np.sin(yaw) * np.cos(pitch), + np.cos(yaw) * np.cos(pitch), + np.sin(pitch), + ]).float().cuda() * 2 + fov = torch.deg2rad(torch.tensor(30)).cuda() + extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda()) + intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov) + exts.append(extrinsics) + ints.append(intrinsics) + + images = [] + + # Build each representation + ss = ss.cuda() + for i in range(ss.shape[0]): + representation = Octree( + depth=10, + aabb=[-0.5, -0.5, -0.5, 1, 1, 1], + device='cuda', + primitive='voxel', + sh_degree=0, + primitive_config={'solid': True}, + ) + coords = torch.nonzero(ss[i, 0], as_tuple=False) + representation.position = coords.float() / self.resolution + representation.depth = torch.full((representation.position.shape[0], 1), int(np.log2(self.resolution)), dtype=torch.uint8, device='cuda') + + image = torch.zeros(3, 1024, 1024).cuda() + tile = [2, 2] + for j, (ext, intr) in enumerate(zip(exts, ints)): + res = renderer.render(representation, ext, intr, colors_overwrite=representation.position) + image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color'] + images.append(image) + + return torch.stack(images) + \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/datasets/sparse_structure_latent.py b/third_party/TRELLIS/trellis/datasets/sparse_structure_latent.py new file mode 100644 index 0000000000000000000000000000000000000000..af4c5d1a81e3e19b7fe629d524965ed3c74aa3a1 --- /dev/null +++ b/third_party/TRELLIS/trellis/datasets/sparse_structure_latent.py @@ -0,0 +1,188 @@ +import os +import json +from typing import * +import numpy as np +import torch +import utils3d +from ..representations.octree import DfsOctree as Octree +from ..renderers import OctreeRenderer +from .components import StandardDatasetBase, TextConditionedMixin, ImageConditionedMixin +from .. import models + + +class SparseStructureLatentVisMixin: + def __init__( + self, + *args, + pretrained_ss_dec: str = 'microsoft/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16', + ss_dec_path: Optional[str] = None, + ss_dec_ckpt: Optional[str] = None, + **kwargs + ): + super().__init__(*args, **kwargs) + self.ss_dec = None + self.pretrained_ss_dec = pretrained_ss_dec + self.ss_dec_path = ss_dec_path + self.ss_dec_ckpt = ss_dec_ckpt + + def _loading_ss_dec(self): + if self.ss_dec is not None: + return + if self.ss_dec_path is not None: + cfg = json.load(open(os.path.join(self.ss_dec_path, 'config.json'), 'r')) + decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args']) + ckpt_path = os.path.join(self.ss_dec_path, 'ckpts', f'decoder_{self.ss_dec_ckpt}.pt') + decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True)) + else: + decoder = models.from_pretrained(self.pretrained_ss_dec) + self.ss_dec = decoder.cuda().eval() + + def _delete_ss_dec(self): + del self.ss_dec + self.ss_dec = None + + @torch.no_grad() + def decode_latent(self, z, batch_size=4): + self._loading_ss_dec() + ss = [] + if self.normalization is not None: + z = z * self.std.to(z.device) + self.mean.to(z.device) + for i in range(0, z.shape[0], batch_size): + ss.append(self.ss_dec(z[i:i+batch_size])) + ss = torch.cat(ss, dim=0) + self._delete_ss_dec() + return ss + + @torch.no_grad() + def visualize_sample(self, x_0: Union[torch.Tensor, dict]): + x_0 = x_0 if isinstance(x_0, torch.Tensor) else x_0['x_0'] + x_0 = self.decode_latent(x_0.cuda()) + + renderer = OctreeRenderer() + renderer.rendering_options.resolution = 512 + renderer.rendering_options.near = 0.8 + renderer.rendering_options.far = 1.6 + renderer.rendering_options.bg_color = (0, 0, 0) + renderer.rendering_options.ssaa = 4 + renderer.pipe.primitive = 'voxel' + + # Build camera + yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2] + yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4) + yaws = [y + yaws_offset for y in yaws] + pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)] + + exts = [] + ints = [] + for yaw, pitch in zip(yaws, pitch): + orig = torch.tensor([ + np.sin(yaw) * np.cos(pitch), + np.cos(yaw) * np.cos(pitch), + np.sin(pitch), + ]).float().cuda() * 2 + fov = torch.deg2rad(torch.tensor(30)).cuda() + extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda()) + intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov) + exts.append(extrinsics) + ints.append(intrinsics) + + images = [] + + # Build each representation + x_0 = x_0.cuda() + for i in range(x_0.shape[0]): + representation = Octree( + depth=10, + aabb=[-0.5, -0.5, -0.5, 1, 1, 1], + device='cuda', + primitive='voxel', + sh_degree=0, + primitive_config={'solid': True}, + ) + coords = torch.nonzero(x_0[i, 0] > 0, as_tuple=False) + resolution = x_0.shape[-1] + representation.position = coords.float() / resolution + representation.depth = torch.full((representation.position.shape[0], 1), int(np.log2(resolution)), dtype=torch.uint8, device='cuda') + + image = torch.zeros(3, 1024, 1024).cuda() + tile = [2, 2] + for j, (ext, intr) in enumerate(zip(exts, ints)): + res = renderer.render(representation, ext, intr, colors_overwrite=representation.position) + image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color'] + images.append(image) + + return torch.stack(images) + + +class SparseStructureLatent(SparseStructureLatentVisMixin, StandardDatasetBase): + """ + Sparse structure latent dataset + + Args: + roots (str): path to the dataset + latent_model (str): name of the latent model + min_aesthetic_score (float): minimum aesthetic score + normalization (dict): normalization stats + pretrained_ss_dec (str): name of the pretrained sparse structure decoder + ss_dec_path (str): path to the sparse structure decoder, if given, will override the pretrained_ss_dec + ss_dec_ckpt (str): name of the sparse structure decoder checkpoint + """ + def __init__(self, + roots: str, + *, + latent_model: str, + min_aesthetic_score: float = 5.0, + normalization: Optional[dict] = None, + pretrained_ss_dec: str = 'microsoft/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16', + ss_dec_path: Optional[str] = None, + ss_dec_ckpt: Optional[str] = None, + ): + self.latent_model = latent_model + self.min_aesthetic_score = min_aesthetic_score + self.normalization = normalization + self.value_range = (0, 1) + + super().__init__( + roots, + pretrained_ss_dec=pretrained_ss_dec, + ss_dec_path=ss_dec_path, + ss_dec_ckpt=ss_dec_ckpt, + ) + + if self.normalization is not None: + self.mean = torch.tensor(self.normalization['mean']).reshape(-1, 1, 1, 1) + self.std = torch.tensor(self.normalization['std']).reshape(-1, 1, 1, 1) + + def filter_metadata(self, metadata): + stats = {} + metadata = metadata[metadata[f'ss_latent_{self.latent_model}']] + stats['With sparse structure latents'] = len(metadata) + metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score] + stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata) + return metadata, stats + + def get_instance(self, root, instance): + latent = np.load(os.path.join(root, 'ss_latents', self.latent_model, f'{instance}.npz')) + z = torch.tensor(latent['mean']).float() + if self.normalization is not None: + z = (z - self.mean) / self.std + + pack = { + 'x_0': z, + } + return pack + + +class TextConditionedSparseStructureLatent(TextConditionedMixin, SparseStructureLatent): + """ + Text-conditioned sparse structure dataset + """ + pass + + +class ImageConditionedSparseStructureLatent(ImageConditionedMixin, SparseStructureLatent): + """ + Image-conditioned sparse structure dataset + """ + pass + \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/datasets/structured_latent.py b/third_party/TRELLIS/trellis/datasets/structured_latent.py new file mode 100644 index 0000000000000000000000000000000000000000..f8b8f24963ec008264396148f2c92afd5fae6d52 --- /dev/null +++ b/third_party/TRELLIS/trellis/datasets/structured_latent.py @@ -0,0 +1,217 @@ +import json +import os +from typing import * +import numpy as np +import torch +import utils3d.torch +from .components import StandardDatasetBase, TextConditionedMixin, ImageConditionedMixin +from ..modules.sparse.basic import SparseTensor +from .. import models +from ..utils.render_utils import get_renderer +from ..utils.data_utils import load_balanced_group_indices + + +class SLatVisMixin: + def __init__( + self, + *args, + pretrained_slat_dec: str = 'microsoft/TRELLIS-image-large/ckpts/slat_dec_gs_swin8_B_64l8gs32_fp16', + slat_dec_path: Optional[str] = None, + slat_dec_ckpt: Optional[str] = None, + **kwargs + ): + super().__init__(*args, **kwargs) + self.slat_dec = None + self.pretrained_slat_dec = pretrained_slat_dec + self.slat_dec_path = slat_dec_path + self.slat_dec_ckpt = slat_dec_ckpt + + def _loading_slat_dec(self): + if self.slat_dec is not None: + return + if self.slat_dec_path is not None: + cfg = json.load(open(os.path.join(self.slat_dec_path, 'config.json'), 'r')) + decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args']) + ckpt_path = os.path.join(self.slat_dec_path, 'ckpts', f'decoder_{self.slat_dec_ckpt}.pt') + decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True)) + else: + decoder = models.from_pretrained(self.pretrained_slat_dec) + self.slat_dec = decoder.cuda().eval() + + def _delete_slat_dec(self): + del self.slat_dec + self.slat_dec = None + + @torch.no_grad() + def decode_latent(self, z, batch_size=4): + self._loading_slat_dec() + reps = [] + if self.normalization is not None: + z = z * self.std.to(z.device) + self.mean.to(z.device) + for i in range(0, z.shape[0], batch_size): + reps.append(self.slat_dec(z[i:i+batch_size])) + reps = sum(reps, []) + self._delete_slat_dec() + return reps + + @torch.no_grad() + def visualize_sample(self, x_0: Union[SparseTensor, dict]): + x_0 = x_0 if isinstance(x_0, SparseTensor) else x_0['x_0'] + reps = self.decode_latent(x_0.cuda()) + + # Build camera + yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2] + yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4) + yaws = [y + yaws_offset for y in yaws] + pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)] + + exts = [] + ints = [] + for yaw, pitch in zip(yaws, pitch): + orig = torch.tensor([ + np.sin(yaw) * np.cos(pitch), + np.cos(yaw) * np.cos(pitch), + np.sin(pitch), + ]).float().cuda() * 2 + fov = torch.deg2rad(torch.tensor(40)).cuda() + extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda()) + intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov) + exts.append(extrinsics) + ints.append(intrinsics) + + renderer = get_renderer(reps[0]) + images = [] + for representation in reps: + image = torch.zeros(3, 1024, 1024).cuda() + tile = [2, 2] + for j, (ext, intr) in enumerate(zip(exts, ints)): + res = renderer.render(representation, ext, intr) + image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color'] + images.append(image) + images = torch.stack(images) + + return images + + +class SLat(SLatVisMixin, StandardDatasetBase): + """ + structured latent dataset + + Args: + roots (str): path to the dataset + latent_model (str): name of the latent model + min_aesthetic_score (float): minimum aesthetic score + max_num_voxels (int): maximum number of voxels + normalization (dict): normalization stats + pretrained_slat_dec (str): name of the pretrained slat decoder + slat_dec_path (str): path to the slat decoder, if given, will override the pretrained_slat_dec + slat_dec_ckpt (str): name of the slat decoder checkpoint + """ + def __init__(self, + roots: str, + *, + latent_model: str, + min_aesthetic_score: float = 5.0, + max_num_voxels: int = 32768, + normalization: Optional[dict] = None, + pretrained_slat_dec: str = 'microsoft/TRELLIS-image-large/ckpts/slat_dec_gs_swin8_B_64l8gs32_fp16', + slat_dec_path: Optional[str] = None, + slat_dec_ckpt: Optional[str] = None, + ): + self.normalization = normalization + self.latent_model = latent_model + self.min_aesthetic_score = min_aesthetic_score + self.max_num_voxels = max_num_voxels + self.value_range = (0, 1) + + super().__init__( + roots, + pretrained_slat_dec=pretrained_slat_dec, + slat_dec_path=slat_dec_path, + slat_dec_ckpt=slat_dec_ckpt, + ) + + self.loads = [self.metadata.loc[sha256, 'num_voxels'] for _, sha256 in self.instances] + + if self.normalization is not None: + self.mean = torch.tensor(self.normalization['mean']).reshape(1, -1) + self.std = torch.tensor(self.normalization['std']).reshape(1, -1) + + def filter_metadata(self, metadata): + stats = {} + metadata = metadata[metadata[f'latent_{self.latent_model}']] + stats['With latent'] = len(metadata) + metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score] + stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata) + metadata = metadata[metadata['num_voxels'] <= self.max_num_voxels] + stats[f'Num voxels <= {self.max_num_voxels}'] = len(metadata) + return metadata, stats + + def get_instance(self, root, instance): + data = np.load(os.path.join(root, 'latents', self.latent_model, f'{instance}.npz')) + coords = torch.tensor(data['coords']).int() + feats = torch.tensor(data['feats']).float() + if self.normalization is not None: + feats = (feats - self.mean) / self.std + return { + 'coords': coords, + 'feats': feats, + } + + @staticmethod + def collate_fn(batch, split_size=None): + if split_size is None: + group_idx = [list(range(len(batch)))] + else: + group_idx = load_balanced_group_indices([b['coords'].shape[0] for b in batch], split_size) + packs = [] + for group in group_idx: + sub_batch = [batch[i] for i in group] + pack = {} + coords = [] + feats = [] + layout = [] + start = 0 + for i, b in enumerate(sub_batch): + coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1)) + feats.append(b['feats']) + layout.append(slice(start, start + b['coords'].shape[0])) + start += b['coords'].shape[0] + coords = torch.cat(coords) + feats = torch.cat(feats) + pack['x_0'] = SparseTensor( + coords=coords, + feats=feats, + ) + pack['x_0']._shape = torch.Size([len(group), *sub_batch[0]['feats'].shape[1:]]) + pack['x_0'].register_spatial_cache('layout', layout) + + # collate other data + keys = [k for k in sub_batch[0].keys() if k not in ['coords', 'feats']] + for k in keys: + if isinstance(sub_batch[0][k], torch.Tensor): + pack[k] = torch.stack([b[k] for b in sub_batch]) + elif isinstance(sub_batch[0][k], list): + pack[k] = sum([b[k] for b in sub_batch], []) + else: + pack[k] = [b[k] for b in sub_batch] + + packs.append(pack) + + if split_size is None: + return packs[0] + return packs + + +class TextConditionedSLat(TextConditionedMixin, SLat): + """ + Text conditioned structured latent dataset + """ + pass + + +class ImageConditionedSLat(ImageConditionedMixin, SLat): + """ + Image conditioned structured latent dataset + """ + pass diff --git a/third_party/TRELLIS/trellis/datasets/structured_latent2render.py b/third_party/TRELLIS/trellis/datasets/structured_latent2render.py new file mode 100644 index 0000000000000000000000000000000000000000..466737ab36c229aaeb368ad4b3af5cb97393eee1 --- /dev/null +++ b/third_party/TRELLIS/trellis/datasets/structured_latent2render.py @@ -0,0 +1,160 @@ +import os +from PIL import Image +import json +import numpy as np +import torch +import utils3d.torch +from ..modules.sparse.basic import SparseTensor +from .components import StandardDatasetBase + + +class SLat2Render(StandardDatasetBase): + """ + Dataset for Structured Latent and rendered images. + + Args: + roots (str): paths to the dataset + image_size (int): size of the image + latent_model (str): latent model name + min_aesthetic_score (float): minimum aesthetic score + max_num_voxels (int): maximum number of voxels + """ + def __init__( + self, + roots: str, + image_size: int, + latent_model: str, + min_aesthetic_score: float = 5.0, + max_num_voxels: int = 32768, + ): + self.image_size = image_size + self.latent_model = latent_model + self.min_aesthetic_score = min_aesthetic_score + self.max_num_voxels = max_num_voxels + self.value_range = (0, 1) + + super().__init__(roots) + + def filter_metadata(self, metadata): + stats = {} + metadata = metadata[metadata[f'latent_{self.latent_model}']] + stats['With latent'] = len(metadata) + metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score] + stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata) + metadata = metadata[metadata['num_voxels'] <= self.max_num_voxels] + stats[f'Num voxels <= {self.max_num_voxels}'] = len(metadata) + return metadata, stats + + def _get_image(self, root, instance): + with open(os.path.join(root, 'renders', instance, 'transforms.json')) as f: + metadata = json.load(f) + n_views = len(metadata['frames']) + view = np.random.randint(n_views) + metadata = metadata['frames'][view] + fov = metadata['camera_angle_x'] + intrinsics = utils3d.torch.intrinsics_from_fov_xy(torch.tensor(fov), torch.tensor(fov)) + c2w = torch.tensor(metadata['transform_matrix']) + c2w[:3, 1:3] *= -1 + extrinsics = torch.inverse(c2w) + + image_path = os.path.join(root, 'renders', instance, metadata['file_path']) + image = Image.open(image_path) + alpha = image.getchannel(3) + image = image.convert('RGB') + image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS) + alpha = alpha.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS) + image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0 + alpha = torch.tensor(np.array(alpha)).float() / 255.0 + + return { + 'image': image, + 'alpha': alpha, + 'extrinsics': extrinsics, + 'intrinsics': intrinsics, + } + + def _get_latent(self, root, instance): + data = np.load(os.path.join(root, 'latents', self.latent_model, f'{instance}.npz')) + coords = torch.tensor(data['coords']).int() + feats = torch.tensor(data['feats']).float() + return { + 'coords': coords, + 'feats': feats, + } + + @torch.no_grad() + def visualize_sample(self, sample: dict): + return sample['image'] + + @staticmethod + def collate_fn(batch): + pack = {} + coords = [] + for i, b in enumerate(batch): + coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1)) + coords = torch.cat(coords) + feats = torch.cat([b['feats'] for b in batch]) + pack['latents'] = SparseTensor( + coords=coords, + feats=feats, + ) + + # collate other data + keys = [k for k in batch[0].keys() if k not in ['coords', 'feats']] + for k in keys: + if isinstance(batch[0][k], torch.Tensor): + pack[k] = torch.stack([b[k] for b in batch]) + elif isinstance(batch[0][k], list): + pack[k] = sum([b[k] for b in batch], []) + else: + pack[k] = [b[k] for b in batch] + + return pack + + def get_instance(self, root, instance): + image = self._get_image(root, instance) + latent = self._get_latent(root, instance) + return { + **image, + **latent, + } + + +class Slat2RenderGeo(SLat2Render): + def __init__( + self, + roots: str, + image_size: int, + latent_model: str, + min_aesthetic_score: float = 5.0, + max_num_voxels: int = 32768, + ): + super().__init__( + roots, + image_size, + latent_model, + min_aesthetic_score, + max_num_voxels, + ) + + def _get_geo(self, root, instance): + verts, face = utils3d.io.read_ply(os.path.join(root, 'renders', instance, 'mesh.ply')) + mesh = { + "vertices" : torch.from_numpy(verts), + "faces" : torch.from_numpy(face), + } + return { + "mesh" : mesh, + } + + def get_instance(self, root, instance): + image = self._get_image(root, instance) + latent = self._get_latent(root, instance) + geo = self._get_geo(root, instance) + return { + **image, + **latent, + **geo, + } + + \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/models/__init__.py b/third_party/TRELLIS/trellis/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ae9d17cbaa645a61d7d173cb185cdb090f96b74f --- /dev/null +++ b/third_party/TRELLIS/trellis/models/__init__.py @@ -0,0 +1,96 @@ +import importlib + +__attributes = { + 'SparseStructureEncoder': 'sparse_structure_vae', + 'SparseStructureDecoder': 'sparse_structure_vae', + + 'SparseStructureFlowModel': 'sparse_structure_flow', + + 'SLatEncoder': 'structured_latent_vae', + 'SLatGaussianDecoder': 'structured_latent_vae', + 'SLatRadianceFieldDecoder': 'structured_latent_vae', + 'SLatMeshDecoder': 'structured_latent_vae', + 'ElasticSLatEncoder': 'structured_latent_vae', + 'ElasticSLatGaussianDecoder': 'structured_latent_vae', + 'ElasticSLatRadianceFieldDecoder': 'structured_latent_vae', + 'ElasticSLatMeshDecoder': 'structured_latent_vae', + + 'SLatFlowModel': 'structured_latent_flow', + 'ElasticSLatFlowModel': 'structured_latent_flow', +} + +__submodules = [] + +__all__ = list(__attributes.keys()) + __submodules + +def __getattr__(name): + if name not in globals(): + if name in __attributes: + module_name = __attributes[name] + module = importlib.import_module(f".{module_name}", __name__) + globals()[name] = getattr(module, name) + elif name in __submodules: + module = importlib.import_module(f".{name}", __name__) + globals()[name] = module + else: + raise AttributeError(f"module {__name__} has no attribute {name}") + return globals()[name] + + +def from_pretrained(path: str, **kwargs): + """ + Load a model from a pretrained checkpoint. + + Args: + path: The path to the checkpoint. Can be either local path or a Hugging Face model name. + NOTE: config file and model file should take the name f'{path}.json' and f'{path}.safetensors' respectively. + **kwargs: Additional arguments for the model constructor. + """ + import os + import json + from safetensors.torch import load_file + is_local = os.path.exists(f"{path}.json") and os.path.exists(f"{path}.safetensors") + + if is_local: + config_file = f"{path}.json" + model_file = f"{path}.safetensors" + else: + from huggingface_hub import hf_hub_download + path_parts = path.split('/') + repo_id = f'{path_parts[0]}/{path_parts[1]}' + model_name = '/'.join(path_parts[2:]) + config_file = hf_hub_download(repo_id, f"{model_name}.json") + model_file = hf_hub_download(repo_id, f"{model_name}.safetensors") + + with open(config_file, 'r') as f: + config = json.load(f) + model = __getattr__(config['name'])(**config['args'], **kwargs) + model.load_state_dict(load_file(model_file)) + + return model + + +# For Pylance +if __name__ == '__main__': + from .sparse_structure_vae import ( + SparseStructureEncoder, + SparseStructureDecoder, + ) + + from .sparse_structure_flow import SparseStructureFlowModel + + from .structured_latent_vae import ( + SLatEncoder, + SLatGaussianDecoder, + SLatRadianceFieldDecoder, + SLatMeshDecoder, + ElasticSLatEncoder, + ElasticSLatGaussianDecoder, + ElasticSLatRadianceFieldDecoder, + ElasticSLatMeshDecoder, + ) + + from .structured_latent_flow import ( + SLatFlowModel, + ElasticSLatFlowModel, + ) diff --git a/third_party/TRELLIS/trellis/models/sparse_elastic_mixin.py b/third_party/TRELLIS/trellis/models/sparse_elastic_mixin.py new file mode 100644 index 0000000000000000000000000000000000000000..66d204c89bedabc2afd1795cdfc6f5d58a6b1ac0 --- /dev/null +++ b/third_party/TRELLIS/trellis/models/sparse_elastic_mixin.py @@ -0,0 +1,24 @@ +from contextlib import contextmanager +from typing import * +import math +from ..modules import sparse as sp +from ..utils.elastic_utils import ElasticModuleMixin + + +class SparseTransformerElasticMixin(ElasticModuleMixin): + def _get_input_size(self, x: sp.SparseTensor, *args, **kwargs): + return x.feats.shape[0] + + @contextmanager + def with_mem_ratio(self, mem_ratio=1.0): + if mem_ratio == 1.0: + yield 1.0 + return + num_blocks = len(self.blocks) + num_checkpoint_blocks = min(math.ceil((1 - mem_ratio) * num_blocks) + 1, num_blocks) + exact_mem_ratio = 1 - (num_checkpoint_blocks - 1) / num_blocks + for i in range(num_blocks): + self.blocks[i].use_checkpoint = i < num_checkpoint_blocks + yield exact_mem_ratio + for i in range(num_blocks): + self.blocks[i].use_checkpoint = False diff --git a/third_party/TRELLIS/trellis/models/sparse_structure_flow.py b/third_party/TRELLIS/trellis/models/sparse_structure_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..aee71a9686fd3795960cf1df970e9b8db0ebd57a --- /dev/null +++ b/third_party/TRELLIS/trellis/models/sparse_structure_flow.py @@ -0,0 +1,200 @@ +from typing import * +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from ..modules.utils import convert_module_to_f16, convert_module_to_f32 +from ..modules.transformer import AbsolutePositionEmbedder, ModulatedTransformerCrossBlock +from ..modules.spatial import patchify, unpatchify + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + + Args: + t: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + dim: the dimension of the output. + max_period: controls the minimum frequency of the embeddings. + + Returns: + an (N, D) Tensor of positional embeddings. + """ + # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + half = dim // 2 + freqs = torch.exp( + -np.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half + ).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + t_emb = self.mlp(t_freq) + return t_emb + + +class SparseStructureFlowModel(nn.Module): + def __init__( + self, + resolution: int, + in_channels: int, + model_channels: int, + cond_channels: int, + out_channels: int, + num_blocks: int, + num_heads: Optional[int] = None, + num_head_channels: Optional[int] = 64, + mlp_ratio: float = 4, + patch_size: int = 2, + pe_mode: Literal["ape", "rope"] = "ape", + use_fp16: bool = False, + use_checkpoint: bool = False, + share_mod: bool = False, + qk_rms_norm: bool = False, + qk_rms_norm_cross: bool = False, + ): + super().__init__() + self.resolution = resolution + self.in_channels = in_channels + self.model_channels = model_channels + self.cond_channels = cond_channels + self.out_channels = out_channels + self.num_blocks = num_blocks + self.num_heads = num_heads or model_channels // num_head_channels + self.mlp_ratio = mlp_ratio + self.patch_size = patch_size + self.pe_mode = pe_mode + self.use_fp16 = use_fp16 + self.use_checkpoint = use_checkpoint + self.share_mod = share_mod + self.qk_rms_norm = qk_rms_norm + self.qk_rms_norm_cross = qk_rms_norm_cross + self.dtype = torch.float16 if use_fp16 else torch.float32 + + self.t_embedder = TimestepEmbedder(model_channels) + if share_mod: + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(model_channels, 6 * model_channels, bias=True) + ) + + if pe_mode == "ape": + pos_embedder = AbsolutePositionEmbedder(model_channels, 3) + coords = torch.meshgrid(*[torch.arange(res, device=self.device) for res in [resolution // patch_size] * 3], indexing='ij') + coords = torch.stack(coords, dim=-1).reshape(-1, 3) + pos_emb = pos_embedder(coords) + self.register_buffer("pos_emb", pos_emb) + + self.input_layer = nn.Linear(in_channels * patch_size**3, model_channels) + + self.blocks = nn.ModuleList([ + ModulatedTransformerCrossBlock( + model_channels, + cond_channels, + num_heads=self.num_heads, + mlp_ratio=self.mlp_ratio, + attn_mode='full', + use_checkpoint=self.use_checkpoint, + use_rope=(pe_mode == "rope"), + share_mod=share_mod, + qk_rms_norm=self.qk_rms_norm, + qk_rms_norm_cross=self.qk_rms_norm_cross, + ) + for _ in range(num_blocks) + ]) + + self.out_layer = nn.Linear(model_channels, out_channels * patch_size**3) + + self.initialize_weights() + if use_fp16: + self.convert_to_fp16() + + @property + def device(self) -> torch.device: + """ + Return the device of the model. + """ + return next(self.parameters()).device + + def convert_to_fp16(self) -> None: + """ + Convert the torso of the model to float16. + """ + self.blocks.apply(convert_module_to_f16) + + def convert_to_fp32(self) -> None: + """ + Convert the torso of the model to float32. + """ + self.blocks.apply(convert_module_to_f32) + + def initialize_weights(self) -> None: + # Initialize transformer layers: + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + self.apply(_basic_init) + + # Initialize timestep embedding MLP: + nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) + + # Zero-out adaLN modulation layers in DiT blocks: + if self.share_mod: + nn.init.constant_(self.adaLN_modulation[-1].weight, 0) + nn.init.constant_(self.adaLN_modulation[-1].bias, 0) + else: + for block in self.blocks: + nn.init.constant_(block.adaLN_modulation[-1].weight, 0) + nn.init.constant_(block.adaLN_modulation[-1].bias, 0) + + # Zero-out output layers: + nn.init.constant_(self.out_layer.weight, 0) + nn.init.constant_(self.out_layer.bias, 0) + + def forward(self, x: torch.Tensor, t: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: + assert [*x.shape] == [x.shape[0], self.in_channels, *[self.resolution] * 3], \ + f"Input shape mismatch, got {x.shape}, expected {[x.shape[0], self.in_channels, *[self.resolution] * 3]}" + + h = patchify(x, self.patch_size) + h = h.view(*h.shape[:2], -1).permute(0, 2, 1).contiguous() + + h = self.input_layer(h) + h = h + self.pos_emb[None] + t_emb = self.t_embedder(t) + if self.share_mod: + t_emb = self.adaLN_modulation(t_emb) + t_emb = t_emb.type(self.dtype) + h = h.type(self.dtype) + cond = cond.type(self.dtype) + for block in self.blocks: + h = block(h, t_emb, cond) + h = h.type(x.dtype) + h = F.layer_norm(h, h.shape[-1:]) + h = self.out_layer(h) + + h = h.permute(0, 2, 1).view(h.shape[0], h.shape[2], *[self.resolution // self.patch_size] * 3) + h = unpatchify(h, self.patch_size).contiguous() + + return h diff --git a/third_party/TRELLIS/trellis/models/sparse_structure_vae.py b/third_party/TRELLIS/trellis/models/sparse_structure_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..c3e09136cf294c4c1b47b0f09fa6ee57bad2166d --- /dev/null +++ b/third_party/TRELLIS/trellis/models/sparse_structure_vae.py @@ -0,0 +1,306 @@ +from typing import * +import torch +import torch.nn as nn +import torch.nn.functional as F +from ..modules.norm import GroupNorm32, ChannelLayerNorm32 +from ..modules.spatial import pixel_shuffle_3d +from ..modules.utils import zero_module, convert_module_to_f16, convert_module_to_f32 + + +def norm_layer(norm_type: str, *args, **kwargs) -> nn.Module: + """ + Return a normalization layer. + """ + if norm_type == "group": + return GroupNorm32(32, *args, **kwargs) + elif norm_type == "layer": + return ChannelLayerNorm32(*args, **kwargs) + else: + raise ValueError(f"Invalid norm type {norm_type}") + + +class ResBlock3d(nn.Module): + def __init__( + self, + channels: int, + out_channels: Optional[int] = None, + norm_type: Literal["group", "layer"] = "layer", + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + + self.norm1 = norm_layer(norm_type, channels) + self.norm2 = norm_layer(norm_type, self.out_channels) + self.conv1 = nn.Conv3d(channels, self.out_channels, 3, padding=1) + self.conv2 = zero_module(nn.Conv3d(self.out_channels, self.out_channels, 3, padding=1)) + self.skip_connection = nn.Conv3d(channels, self.out_channels, 1) if channels != self.out_channels else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.norm1(x) + h = F.silu(h) + h = self.conv1(h) + h = self.norm2(h) + h = F.silu(h) + h = self.conv2(h) + h = h + self.skip_connection(x) + return h + + +class DownsampleBlock3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + mode: Literal["conv", "avgpool"] = "conv", + ): + assert mode in ["conv", "avgpool"], f"Invalid mode {mode}" + + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + if mode == "conv": + self.conv = nn.Conv3d(in_channels, out_channels, 2, stride=2) + elif mode == "avgpool": + assert in_channels == out_channels, "Pooling mode requires in_channels to be equal to out_channels" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if hasattr(self, "conv"): + return self.conv(x) + else: + return F.avg_pool3d(x, 2) + + +class UpsampleBlock3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + mode: Literal["conv", "nearest"] = "conv", + ): + assert mode in ["conv", "nearest"], f"Invalid mode {mode}" + + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + if mode == "conv": + self.conv = nn.Conv3d(in_channels, out_channels*8, 3, padding=1) + elif mode == "nearest": + assert in_channels == out_channels, "Nearest mode requires in_channels to be equal to out_channels" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if hasattr(self, "conv"): + x = self.conv(x) + return pixel_shuffle_3d(x, 2) + else: + return F.interpolate(x, scale_factor=2, mode="nearest") + + +class SparseStructureEncoder(nn.Module): + """ + Encoder for Sparse Structure (\mathcal{E}_S in the paper Sec. 3.3). + + Args: + in_channels (int): Channels of the input. + latent_channels (int): Channels of the latent representation. + num_res_blocks (int): Number of residual blocks at each resolution. + channels (List[int]): Channels of the encoder blocks. + num_res_blocks_middle (int): Number of residual blocks in the middle. + norm_type (Literal["group", "layer"]): Type of normalization layer. + use_fp16 (bool): Whether to use FP16. + """ + def __init__( + self, + in_channels: int, + latent_channels: int, + num_res_blocks: int, + channels: List[int], + num_res_blocks_middle: int = 2, + norm_type: Literal["group", "layer"] = "layer", + use_fp16: bool = False, + ): + super().__init__() + self.in_channels = in_channels + self.latent_channels = latent_channels + self.num_res_blocks = num_res_blocks + self.channels = channels + self.num_res_blocks_middle = num_res_blocks_middle + self.norm_type = norm_type + self.use_fp16 = use_fp16 + self.dtype = torch.float16 if use_fp16 else torch.float32 + + self.input_layer = nn.Conv3d(in_channels, channels[0], 3, padding=1) + + self.blocks = nn.ModuleList([]) + for i, ch in enumerate(channels): + self.blocks.extend([ + ResBlock3d(ch, ch) + for _ in range(num_res_blocks) + ]) + if i < len(channels) - 1: + self.blocks.append( + DownsampleBlock3d(ch, channels[i+1]) + ) + + self.middle_block = nn.Sequential(*[ + ResBlock3d(channels[-1], channels[-1]) + for _ in range(num_res_blocks_middle) + ]) + + self.out_layer = nn.Sequential( + norm_layer(norm_type, channels[-1]), + nn.SiLU(), + nn.Conv3d(channels[-1], latent_channels*2, 3, padding=1) + ) + + if use_fp16: + self.convert_to_fp16() + + @property + def device(self) -> torch.device: + """ + Return the device of the model. + """ + return next(self.parameters()).device + + def convert_to_fp16(self) -> None: + """ + Convert the torso of the model to float16. + """ + self.use_fp16 = True + self.dtype = torch.float16 + self.blocks.apply(convert_module_to_f16) + self.middle_block.apply(convert_module_to_f16) + + def convert_to_fp32(self) -> None: + """ + Convert the torso of the model to float32. + """ + self.use_fp16 = False + self.dtype = torch.float32 + self.blocks.apply(convert_module_to_f32) + self.middle_block.apply(convert_module_to_f32) + + def forward(self, x: torch.Tensor, sample_posterior: bool = False, return_raw: bool = False) -> torch.Tensor: + h = self.input_layer(x) + h = h.type(self.dtype) + + for block in self.blocks: + h = block(h) + h = self.middle_block(h) + + h = h.type(x.dtype) + h = self.out_layer(h) + + mean, logvar = h.chunk(2, dim=1) + + if sample_posterior: + std = torch.exp(0.5 * logvar) + z = mean + std * torch.randn_like(std) + else: + z = mean + + if return_raw: + return z, mean, logvar + return z + + +class SparseStructureDecoder(nn.Module): + """ + Decoder for Sparse Structure (\mathcal{D}_S in the paper Sec. 3.3). + + Args: + out_channels (int): Channels of the output. + latent_channels (int): Channels of the latent representation. + num_res_blocks (int): Number of residual blocks at each resolution. + channels (List[int]): Channels of the decoder blocks. + num_res_blocks_middle (int): Number of residual blocks in the middle. + norm_type (Literal["group", "layer"]): Type of normalization layer. + use_fp16 (bool): Whether to use FP16. + """ + def __init__( + self, + out_channels: int, + latent_channels: int, + num_res_blocks: int, + channels: List[int], + num_res_blocks_middle: int = 2, + norm_type: Literal["group", "layer"] = "layer", + use_fp16: bool = False, + ): + super().__init__() + self.out_channels = out_channels + self.latent_channels = latent_channels + self.num_res_blocks = num_res_blocks + self.channels = channels + self.num_res_blocks_middle = num_res_blocks_middle + self.norm_type = norm_type + self.use_fp16 = use_fp16 + self.dtype = torch.float16 if use_fp16 else torch.float32 + + self.input_layer = nn.Conv3d(latent_channels, channels[0], 3, padding=1) + + self.middle_block = nn.Sequential(*[ + ResBlock3d(channels[0], channels[0]) + for _ in range(num_res_blocks_middle) + ]) + + self.blocks = nn.ModuleList([]) + for i, ch in enumerate(channels): + self.blocks.extend([ + ResBlock3d(ch, ch) + for _ in range(num_res_blocks) + ]) + if i < len(channels) - 1: + self.blocks.append( + UpsampleBlock3d(ch, channels[i+1]) + ) + + self.out_layer = nn.Sequential( + norm_layer(norm_type, channels[-1]), + nn.SiLU(), + nn.Conv3d(channels[-1], out_channels, 3, padding=1) + ) + + if use_fp16: + self.convert_to_fp16() + + @property + def device(self) -> torch.device: + """ + Return the device of the model. + """ + return next(self.parameters()).device + + def convert_to_fp16(self) -> None: + """ + Convert the torso of the model to float16. + """ + self.use_fp16 = True + self.dtype = torch.float16 + self.blocks.apply(convert_module_to_f16) + self.middle_block.apply(convert_module_to_f16) + + def convert_to_fp32(self) -> None: + """ + Convert the torso of the model to float32. + """ + self.use_fp16 = False + self.dtype = torch.float32 + self.blocks.apply(convert_module_to_f32) + self.middle_block.apply(convert_module_to_f32) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.input_layer(x) + + h = h.type(self.dtype) + + h = self.middle_block(h) + for block in self.blocks: + h = block(h) + + h = h.type(x.dtype) + h = self.out_layer(h) + return h diff --git a/third_party/TRELLIS/trellis/models/structured_latent_flow.py b/third_party/TRELLIS/trellis/models/structured_latent_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..d661f55d581e20daf7312f15813c7a2fe888b92e --- /dev/null +++ b/third_party/TRELLIS/trellis/models/structured_latent_flow.py @@ -0,0 +1,295 @@ +from typing import * +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from ..modules.utils import zero_module, convert_module_to_f16, convert_module_to_f32 +from ..modules.transformer import AbsolutePositionEmbedder +from ..modules.norm import LayerNorm32 +from ..modules import sparse as sp +from ..modules.sparse.transformer import ModulatedSparseTransformerCrossBlock +from .sparse_structure_flow import TimestepEmbedder +from .sparse_elastic_mixin import SparseTransformerElasticMixin + + +class SparseResBlock3d(nn.Module): + def __init__( + self, + channels: int, + emb_channels: int, + out_channels: Optional[int] = None, + downsample: bool = False, + upsample: bool = False, + ): + super().__init__() + self.channels = channels + self.emb_channels = emb_channels + self.out_channels = out_channels or channels + self.downsample = downsample + self.upsample = upsample + + assert not (downsample and upsample), "Cannot downsample and upsample at the same time" + + self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6) + self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6) + self.conv1 = sp.SparseConv3d(channels, self.out_channels, 3) + self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3)) + self.emb_layers = nn.Sequential( + nn.SiLU(), + nn.Linear(emb_channels, 2 * self.out_channels, bias=True), + ) + self.skip_connection = sp.SparseLinear(channels, self.out_channels) if channels != self.out_channels else nn.Identity() + self.updown = None + if self.downsample: + self.updown = sp.SparseDownsample(2) + elif self.upsample: + self.updown = sp.SparseUpsample(2) + + def _updown(self, x: sp.SparseTensor) -> sp.SparseTensor: + if self.updown is not None: + x = self.updown(x) + return x + + def forward(self, x: sp.SparseTensor, emb: torch.Tensor) -> sp.SparseTensor: + emb_out = self.emb_layers(emb).type(x.dtype) + scale, shift = torch.chunk(emb_out, 2, dim=1) + + x = self._updown(x) + h = x.replace(self.norm1(x.feats)) + h = h.replace(F.silu(h.feats)) + h = self.conv1(h) + h = h.replace(self.norm2(h.feats)) * (1 + scale) + shift + h = h.replace(F.silu(h.feats)) + h = self.conv2(h) + h = h + self.skip_connection(x) + + return h + + +class SLatFlowModel(nn.Module): + def __init__( + self, + resolution: int, + in_channels: int, + model_channels: int, + cond_channels: int, + out_channels: int, + num_blocks: int, + num_heads: Optional[int] = None, + num_head_channels: Optional[int] = 64, + mlp_ratio: float = 4, + patch_size: int = 2, + num_io_res_blocks: int = 2, + io_block_channels: List[int] = None, + pe_mode: Literal["ape", "rope"] = "ape", + use_fp16: bool = False, + use_checkpoint: bool = False, + use_skip_connection: bool = True, + share_mod: bool = False, + qk_rms_norm: bool = False, + qk_rms_norm_cross: bool = False, + ): + super().__init__() + self.resolution = resolution + self.in_channels = in_channels + self.model_channels = model_channels + self.cond_channels = cond_channels + self.out_channels = out_channels + self.num_blocks = num_blocks + self.num_heads = num_heads or model_channels // num_head_channels + self.mlp_ratio = mlp_ratio + self.patch_size = patch_size + self.num_io_res_blocks = num_io_res_blocks + self.io_block_channels = io_block_channels + self.pe_mode = pe_mode + self.use_fp16 = use_fp16 + self.use_checkpoint = use_checkpoint + self.use_skip_connection = use_skip_connection + self.share_mod = share_mod + self.qk_rms_norm = qk_rms_norm + self.qk_rms_norm_cross = qk_rms_norm_cross + self.dtype = torch.float16 if use_fp16 else torch.float32 + + if self.io_block_channels is not None: + assert int(np.log2(patch_size)) == np.log2(patch_size), "Patch size must be a power of 2" + assert np.log2(patch_size) == len(io_block_channels), "Number of IO ResBlocks must match the number of stages" + + self.t_embedder = TimestepEmbedder(model_channels) + if share_mod: + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(model_channels, 6 * model_channels, bias=True) + ) + + if pe_mode == "ape": + self.pos_embedder = AbsolutePositionEmbedder(model_channels) + + self.input_layer = sp.SparseLinear(in_channels, model_channels if io_block_channels is None else io_block_channels[0]) + + self.input_blocks = nn.ModuleList([]) + if io_block_channels is not None: + for chs, next_chs in zip(io_block_channels, io_block_channels[1:] + [model_channels]): + self.input_blocks.extend([ + SparseResBlock3d( + chs, + model_channels, + out_channels=chs, + ) + for _ in range(num_io_res_blocks-1) + ]) + self.input_blocks.append( + SparseResBlock3d( + chs, + model_channels, + out_channels=next_chs, + downsample=True, + ) + ) + + self.blocks = nn.ModuleList([ + ModulatedSparseTransformerCrossBlock( + model_channels, + cond_channels, + num_heads=self.num_heads, + mlp_ratio=self.mlp_ratio, + attn_mode='full', + use_checkpoint=self.use_checkpoint, + use_rope=(pe_mode == "rope"), + share_mod=self.share_mod, + qk_rms_norm=self.qk_rms_norm, + qk_rms_norm_cross=self.qk_rms_norm_cross, + ) + for _ in range(num_blocks) + ]) + + self.out_blocks = nn.ModuleList([]) + if io_block_channels is not None: + for chs, prev_chs in zip(reversed(io_block_channels), [model_channels] + list(reversed(io_block_channels[1:]))): + self.out_blocks.append( + SparseResBlock3d( + prev_chs * 2 if self.use_skip_connection else prev_chs, + model_channels, + out_channels=chs, + upsample=True, + ) + ) + self.out_blocks.extend([ + SparseResBlock3d( + chs * 2 if self.use_skip_connection else chs, + model_channels, + out_channels=chs, + ) + for _ in range(num_io_res_blocks-1) + ]) + + self.out_layer = sp.SparseLinear(model_channels if io_block_channels is None else io_block_channels[0], out_channels) + + self.initialize_weights() + if use_fp16: + self.convert_to_fp16() + + @property + def device(self) -> torch.device: + """ + Return the device of the model. + """ + return next(self.parameters()).device + + def convert_to_fp16(self) -> None: + """ + Convert the torso of the model to float16. + """ + self.input_blocks.apply(convert_module_to_f16) + self.blocks.apply(convert_module_to_f16) + self.out_blocks.apply(convert_module_to_f16) + + def convert_to_fp32(self) -> None: + """ + Convert the torso of the model to float32. + """ + self.input_blocks.apply(convert_module_to_f32) + self.blocks.apply(convert_module_to_f32) + self.out_blocks.apply(convert_module_to_f32) + + def initialize_weights(self) -> None: + # Initialize transformer layers: + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + self.apply(_basic_init) + + # Initialize timestep embedding MLP: + nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) + + # Zero-out adaLN modulation layers in DiT blocks: + if self.share_mod: + nn.init.constant_(self.adaLN_modulation[-1].weight, 0) + nn.init.constant_(self.adaLN_modulation[-1].bias, 0) + else: + for block in self.blocks: + nn.init.constant_(block.adaLN_modulation[-1].weight, 0) + nn.init.constant_(block.adaLN_modulation[-1].bias, 0) + + # Zero-out output layers: + nn.init.constant_(self.out_layer.weight, 0) + nn.init.constant_(self.out_layer.bias, 0) + + def forward( + self, + x: sp.SparseTensor, + t: torch.Tensor, + cond: torch.Tensor, + cond_list: Optional[List[torch.Tensor]] = None, + coords_dense_indices: Optional[torch.Tensor] = None, + self_attn_region: Optional[torch.Tensor] = None, + self_attn_region_boost: Optional[float] = None, + ) -> sp.SparseTensor: + h = self.input_layer(x).type(self.dtype) + t_emb = self.t_embedder(t) + if self.share_mod: + t_emb = self.adaLN_modulation(t_emb) + t_emb = t_emb.type(self.dtype) + cond = cond.type(self.dtype) + if cond_list is not None: + cond_list = [local_cond.type(self.dtype) for local_cond in cond_list] + + skips = [] + # pack with input blocks + for block in self.input_blocks: + h = block(h, t_emb) + skips.append(h.feats) + + if self.pe_mode == "ape": + h = h + self.pos_embedder(h.coords[:, 1:]).type(self.dtype) + for block in self.blocks: + h = block( + h, + t_emb, + cond, + context_list=cond_list, + coords_dense_indices=coords_dense_indices, + self_attn_region=self_attn_region, + self_attn_region_boost=self_attn_region_boost, + ) + + # unpack with output blocks + for block, skip in zip(self.out_blocks, reversed(skips)): + if self.use_skip_connection: + h = block(h.replace(torch.cat([h.feats, skip], dim=1)), t_emb) + else: + h = block(h, t_emb) + + h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:])) + h = self.out_layer(h.type(x.dtype)) + return h + + +class ElasticSLatFlowModel(SparseTransformerElasticMixin, SLatFlowModel): + """ + SLat Flow Model with elastic memory management. + Used for training with low VRAM. + """ + pass diff --git a/third_party/TRELLIS/trellis/models/structured_latent_vae/__init__.py b/third_party/TRELLIS/trellis/models/structured_latent_vae/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4e2ac3531d71b5f4fb168cae62e776a7d4d55576 --- /dev/null +++ b/third_party/TRELLIS/trellis/models/structured_latent_vae/__init__.py @@ -0,0 +1,4 @@ +from .encoder import SLatEncoder, ElasticSLatEncoder +from .decoder_gs import SLatGaussianDecoder, ElasticSLatGaussianDecoder +from .decoder_rf import SLatRadianceFieldDecoder, ElasticSLatRadianceFieldDecoder +from .decoder_mesh import SLatMeshDecoder, ElasticSLatMeshDecoder diff --git a/third_party/TRELLIS/trellis/models/structured_latent_vae/base.py b/third_party/TRELLIS/trellis/models/structured_latent_vae/base.py new file mode 100644 index 0000000000000000000000000000000000000000..ab0bf6a850b1c146e081c32ad92c7c44ead5ef6e --- /dev/null +++ b/third_party/TRELLIS/trellis/models/structured_latent_vae/base.py @@ -0,0 +1,117 @@ +from typing import * +import torch +import torch.nn as nn +from ...modules.utils import convert_module_to_f16, convert_module_to_f32 +from ...modules import sparse as sp +from ...modules.transformer import AbsolutePositionEmbedder +from ...modules.sparse.transformer import SparseTransformerBlock + + +def block_attn_config(self): + """ + Return the attention configuration of the model. + """ + for i in range(self.num_blocks): + if self.attn_mode == "shift_window": + yield "serialized", self.window_size, 0, (16 * (i % 2),) * 3, sp.SerializeMode.Z_ORDER + elif self.attn_mode == "shift_sequence": + yield "serialized", self.window_size, self.window_size // 2 * (i % 2), (0, 0, 0), sp.SerializeMode.Z_ORDER + elif self.attn_mode == "shift_order": + yield "serialized", self.window_size, 0, (0, 0, 0), sp.SerializeModes[i % 4] + elif self.attn_mode == "full": + yield "full", None, None, None, None + elif self.attn_mode == "swin": + yield "windowed", self.window_size, None, self.window_size // 2 * (i % 2), None + + +class SparseTransformerBase(nn.Module): + """ + Sparse Transformer without output layers. + Serve as the base class for encoder and decoder. + """ + def __init__( + self, + in_channels: int, + model_channels: int, + num_blocks: int, + num_heads: Optional[int] = None, + num_head_channels: Optional[int] = 64, + mlp_ratio: float = 4.0, + attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "full", + window_size: Optional[int] = None, + pe_mode: Literal["ape", "rope"] = "ape", + use_fp16: bool = False, + use_checkpoint: bool = False, + qk_rms_norm: bool = False, + ): + super().__init__() + self.in_channels = in_channels + self.model_channels = model_channels + self.num_blocks = num_blocks + self.window_size = window_size + self.num_heads = num_heads or model_channels // num_head_channels + self.mlp_ratio = mlp_ratio + self.attn_mode = attn_mode + self.pe_mode = pe_mode + self.use_fp16 = use_fp16 + self.use_checkpoint = use_checkpoint + self.qk_rms_norm = qk_rms_norm + self.dtype = torch.float16 if use_fp16 else torch.float32 + + if pe_mode == "ape": + self.pos_embedder = AbsolutePositionEmbedder(model_channels) + + self.input_layer = sp.SparseLinear(in_channels, model_channels) + self.blocks = nn.ModuleList([ + SparseTransformerBlock( + model_channels, + num_heads=self.num_heads, + mlp_ratio=self.mlp_ratio, + attn_mode=attn_mode, + window_size=window_size, + shift_sequence=shift_sequence, + shift_window=shift_window, + serialize_mode=serialize_mode, + use_checkpoint=self.use_checkpoint, + use_rope=(pe_mode == "rope"), + qk_rms_norm=self.qk_rms_norm, + ) + for attn_mode, window_size, shift_sequence, shift_window, serialize_mode in block_attn_config(self) + ]) + + @property + def device(self) -> torch.device: + """ + Return the device of the model. + """ + return next(self.parameters()).device + + def convert_to_fp16(self) -> None: + """ + Convert the torso of the model to float16. + """ + self.blocks.apply(convert_module_to_f16) + + def convert_to_fp32(self) -> None: + """ + Convert the torso of the model to float32. + """ + self.blocks.apply(convert_module_to_f32) + + def initialize_weights(self) -> None: + # Initialize transformer layers: + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + self.apply(_basic_init) + + def forward(self, x: sp.SparseTensor) -> sp.SparseTensor: + h = self.input_layer(x) + if self.pe_mode == "ape": + h = h + self.pos_embedder(x.coords[:, 1:]) + h = h.type(self.dtype) + for block in self.blocks: + h = block(h) + return h diff --git a/third_party/TRELLIS/trellis/models/structured_latent_vae/decoder_gs.py b/third_party/TRELLIS/trellis/models/structured_latent_vae/decoder_gs.py new file mode 100644 index 0000000000000000000000000000000000000000..47304f381833bf9529f6576000972ae14f0cc4ba --- /dev/null +++ b/third_party/TRELLIS/trellis/models/structured_latent_vae/decoder_gs.py @@ -0,0 +1,131 @@ +from typing import * +import torch +import torch.nn as nn +import torch.nn.functional as F +from ...modules import sparse as sp +from ...utils.random_utils import hammersley_sequence +from .base import SparseTransformerBase +from ...representations import Gaussian +from ..sparse_elastic_mixin import SparseTransformerElasticMixin + + +class SLatGaussianDecoder(SparseTransformerBase): + def __init__( + self, + resolution: int, + model_channels: int, + latent_channels: int, + num_blocks: int, + num_heads: Optional[int] = None, + num_head_channels: Optional[int] = 64, + mlp_ratio: float = 4, + attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "swin", + window_size: int = 8, + pe_mode: Literal["ape", "rope"] = "ape", + use_fp16: bool = False, + use_checkpoint: bool = False, + qk_rms_norm: bool = False, + representation_config: dict = None, + ): + super().__init__( + in_channels=latent_channels, + model_channels=model_channels, + num_blocks=num_blocks, + num_heads=num_heads, + num_head_channels=num_head_channels, + mlp_ratio=mlp_ratio, + attn_mode=attn_mode, + window_size=window_size, + pe_mode=pe_mode, + use_fp16=use_fp16, + use_checkpoint=use_checkpoint, + qk_rms_norm=qk_rms_norm, + ) + self.resolution = resolution + self.rep_config = representation_config + self._calc_layout() + self.out_layer = sp.SparseLinear(model_channels, self.out_channels) + self._build_perturbation() + + self.initialize_weights() + if use_fp16: + self.convert_to_fp16() + + def initialize_weights(self) -> None: + super().initialize_weights() + # Zero-out output layers: + nn.init.constant_(self.out_layer.weight, 0) + nn.init.constant_(self.out_layer.bias, 0) + + def _build_perturbation(self) -> None: + perturbation = [hammersley_sequence(3, i, self.rep_config['num_gaussians']) for i in range(self.rep_config['num_gaussians'])] + perturbation = torch.tensor(perturbation).float() * 2 - 1 + perturbation = perturbation / self.rep_config['voxel_size'] + perturbation = torch.atanh(perturbation).to(self.device) + self.register_buffer('offset_perturbation', perturbation) + + def _calc_layout(self) -> None: + self.layout = { + '_xyz' : {'shape': (self.rep_config['num_gaussians'], 3), 'size': self.rep_config['num_gaussians'] * 3}, + '_features_dc' : {'shape': (self.rep_config['num_gaussians'], 1, 3), 'size': self.rep_config['num_gaussians'] * 3}, + '_scaling' : {'shape': (self.rep_config['num_gaussians'], 3), 'size': self.rep_config['num_gaussians'] * 3}, + '_rotation' : {'shape': (self.rep_config['num_gaussians'], 4), 'size': self.rep_config['num_gaussians'] * 4}, + '_opacity' : {'shape': (self.rep_config['num_gaussians'], 1), 'size': self.rep_config['num_gaussians']}, + } + start = 0 + for k, v in self.layout.items(): + v['range'] = (start, start + v['size']) + start += v['size'] + self.out_channels = start + + def to_representation(self, x: sp.SparseTensor) -> List[Gaussian]: + """ + Convert a batch of network outputs to 3D representations. + + Args: + x: The [N x * x C] sparse tensor output by the network. + + Returns: + list of representations + """ + ret = [] + for i in range(x.shape[0]): + representation = Gaussian( + sh_degree=0, + aabb=[-0.5, -0.5, -0.5, 1.0, 1.0, 1.0], + mininum_kernel_size = self.rep_config['3d_filter_kernel_size'], + scaling_bias = self.rep_config['scaling_bias'], + opacity_bias = self.rep_config['opacity_bias'], + scaling_activation = self.rep_config['scaling_activation'] + ) + xyz = (x.coords[x.layout[i]][:, 1:].float() + 0.5) / self.resolution + for k, v in self.layout.items(): + if k == '_xyz': + offset = x.feats[x.layout[i]][:, v['range'][0]:v['range'][1]].reshape(-1, *v['shape']) + offset = offset * self.rep_config['lr'][k] + if self.rep_config['perturb_offset']: + offset = offset + self.offset_perturbation + offset = torch.tanh(offset) / self.resolution * 0.5 * self.rep_config['voxel_size'] + _xyz = xyz.unsqueeze(1) + offset + setattr(representation, k, _xyz.flatten(0, 1)) + else: + feats = x.feats[x.layout[i]][:, v['range'][0]:v['range'][1]].reshape(-1, *v['shape']).flatten(0, 1) + feats = feats * self.rep_config['lr'][k] + setattr(representation, k, feats) + ret.append(representation) + return ret + + def forward(self, x: sp.SparseTensor) -> List[Gaussian]: + h = super().forward(x) + h = h.type(x.dtype) + h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:])) + h = self.out_layer(h) + return self.to_representation(h) + + +class ElasticSLatGaussianDecoder(SparseTransformerElasticMixin, SLatGaussianDecoder): + """ + Slat VAE Gaussian decoder with elastic memory management. + Used for training with low VRAM. + """ + pass diff --git a/third_party/TRELLIS/trellis/models/structured_latent_vae/decoder_mesh.py b/third_party/TRELLIS/trellis/models/structured_latent_vae/decoder_mesh.py new file mode 100644 index 0000000000000000000000000000000000000000..f9f5659d1d0556bcf51b973ace584fdb347ac4d6 --- /dev/null +++ b/third_party/TRELLIS/trellis/models/structured_latent_vae/decoder_mesh.py @@ -0,0 +1,176 @@ +from typing import * +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from ...modules.utils import zero_module, convert_module_to_f16, convert_module_to_f32 +from ...modules import sparse as sp +from .base import SparseTransformerBase +from ...representations import MeshExtractResult +from ...representations.mesh import SparseFeatures2Mesh +from ..sparse_elastic_mixin import SparseTransformerElasticMixin + + +class SparseSubdivideBlock3d(nn.Module): + """ + A 3D subdivide block that can subdivide the sparse tensor. + + Args: + channels: channels in the inputs and outputs. + out_channels: if specified, the number of output channels. + num_groups: the number of groups for the group norm. + """ + def __init__( + self, + channels: int, + resolution: int, + out_channels: Optional[int] = None, + num_groups: int = 32 + ): + super().__init__() + self.channels = channels + self.resolution = resolution + self.out_resolution = resolution * 2 + self.out_channels = out_channels or channels + + self.act_layers = nn.Sequential( + sp.SparseGroupNorm32(num_groups, channels), + sp.SparseSiLU() + ) + + self.sub = sp.SparseSubdivide() + + self.out_layers = nn.Sequential( + sp.SparseConv3d(channels, self.out_channels, 3, indice_key=f"res_{self.out_resolution}"), + sp.SparseGroupNorm32(num_groups, self.out_channels), + sp.SparseSiLU(), + zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3, indice_key=f"res_{self.out_resolution}")), + ) + + if self.out_channels == channels: + self.skip_connection = nn.Identity() + else: + self.skip_connection = sp.SparseConv3d(channels, self.out_channels, 1, indice_key=f"res_{self.out_resolution}") + + def forward(self, x: sp.SparseTensor) -> sp.SparseTensor: + """ + Apply the block to a Tensor, conditioned on a timestep embedding. + + Args: + x: an [N x C x ...] Tensor of features. + Returns: + an [N x C x ...] Tensor of outputs. + """ + h = self.act_layers(x) + h = self.sub(h) + x = self.sub(x) + h = self.out_layers(h) + h = h + self.skip_connection(x) + return h + + +class SLatMeshDecoder(SparseTransformerBase): + def __init__( + self, + resolution: int, + model_channels: int, + latent_channels: int, + num_blocks: int, + num_heads: Optional[int] = None, + num_head_channels: Optional[int] = 64, + mlp_ratio: float = 4, + attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "swin", + window_size: int = 8, + pe_mode: Literal["ape", "rope"] = "ape", + use_fp16: bool = False, + use_checkpoint: bool = False, + qk_rms_norm: bool = False, + representation_config: dict = None, + ): + super().__init__( + in_channels=latent_channels, + model_channels=model_channels, + num_blocks=num_blocks, + num_heads=num_heads, + num_head_channels=num_head_channels, + mlp_ratio=mlp_ratio, + attn_mode=attn_mode, + window_size=window_size, + pe_mode=pe_mode, + use_fp16=use_fp16, + use_checkpoint=use_checkpoint, + qk_rms_norm=qk_rms_norm, + ) + self.resolution = resolution + self.rep_config = representation_config + self.mesh_extractor = SparseFeatures2Mesh(res=self.resolution*4, use_color=self.rep_config.get('use_color', False)) + self.out_channels = self.mesh_extractor.feats_channels + self.upsample = nn.ModuleList([ + SparseSubdivideBlock3d( + channels=model_channels, + resolution=resolution, + out_channels=model_channels // 4 + ), + SparseSubdivideBlock3d( + channels=model_channels // 4, + resolution=resolution * 2, + out_channels=model_channels // 8 + ) + ]) + self.out_layer = sp.SparseLinear(model_channels // 8, self.out_channels) + + self.initialize_weights() + if use_fp16: + self.convert_to_fp16() + + def initialize_weights(self) -> None: + super().initialize_weights() + # Zero-out output layers: + nn.init.constant_(self.out_layer.weight, 0) + nn.init.constant_(self.out_layer.bias, 0) + + def convert_to_fp16(self) -> None: + """ + Convert the torso of the model to float16. + """ + super().convert_to_fp16() + self.upsample.apply(convert_module_to_f16) + + def convert_to_fp32(self) -> None: + """ + Convert the torso of the model to float32. + """ + super().convert_to_fp32() + self.upsample.apply(convert_module_to_f32) + + def to_representation(self, x: sp.SparseTensor) -> List[MeshExtractResult]: + """ + Convert a batch of network outputs to 3D representations. + + Args: + x: The [N x * x C] sparse tensor output by the network. + + Returns: + list of representations + """ + ret = [] + for i in range(x.shape[0]): + mesh = self.mesh_extractor(x[i], training=self.training) + ret.append(mesh) + return ret + + def forward(self, x: sp.SparseTensor) -> List[MeshExtractResult]: + h = super().forward(x) + for block in self.upsample: + h = block(h) + h = h.type(x.dtype) + h = self.out_layer(h) + return self.to_representation(h) + + +class ElasticSLatMeshDecoder(SparseTransformerElasticMixin, SLatMeshDecoder): + """ + Slat VAE Mesh decoder with elastic memory management. + Used for training with low VRAM. + """ + pass diff --git a/third_party/TRELLIS/trellis/models/structured_latent_vae/decoder_rf.py b/third_party/TRELLIS/trellis/models/structured_latent_vae/decoder_rf.py new file mode 100644 index 0000000000000000000000000000000000000000..bb1809449f689dfa7bb03673782039f3015d8479 --- /dev/null +++ b/third_party/TRELLIS/trellis/models/structured_latent_vae/decoder_rf.py @@ -0,0 +1,113 @@ +from typing import * +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from ...modules import sparse as sp +from .base import SparseTransformerBase +from ...representations import Strivec +from ..sparse_elastic_mixin import SparseTransformerElasticMixin + + +class SLatRadianceFieldDecoder(SparseTransformerBase): + def __init__( + self, + resolution: int, + model_channels: int, + latent_channels: int, + num_blocks: int, + num_heads: Optional[int] = None, + num_head_channels: Optional[int] = 64, + mlp_ratio: float = 4, + attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "swin", + window_size: int = 8, + pe_mode: Literal["ape", "rope"] = "ape", + use_fp16: bool = False, + use_checkpoint: bool = False, + qk_rms_norm: bool = False, + representation_config: dict = None, + ): + super().__init__( + in_channels=latent_channels, + model_channels=model_channels, + num_blocks=num_blocks, + num_heads=num_heads, + num_head_channels=num_head_channels, + mlp_ratio=mlp_ratio, + attn_mode=attn_mode, + window_size=window_size, + pe_mode=pe_mode, + use_fp16=use_fp16, + use_checkpoint=use_checkpoint, + qk_rms_norm=qk_rms_norm, + ) + self.resolution = resolution + self.rep_config = representation_config + self._calc_layout() + self.out_layer = sp.SparseLinear(model_channels, self.out_channels) + + self.initialize_weights() + if use_fp16: + self.convert_to_fp16() + + def initialize_weights(self) -> None: + super().initialize_weights() + # Zero-out output layers: + nn.init.constant_(self.out_layer.weight, 0) + nn.init.constant_(self.out_layer.bias, 0) + + def _calc_layout(self) -> None: + self.layout = { + 'trivec': {'shape': (self.rep_config['rank'], 3, self.rep_config['dim']), 'size': self.rep_config['rank'] * 3 * self.rep_config['dim']}, + 'density': {'shape': (self.rep_config['rank'],), 'size': self.rep_config['rank']}, + 'features_dc': {'shape': (self.rep_config['rank'], 1, 3), 'size': self.rep_config['rank'] * 3}, + } + start = 0 + for k, v in self.layout.items(): + v['range'] = (start, start + v['size']) + start += v['size'] + self.out_channels = start + + def to_representation(self, x: sp.SparseTensor) -> List[Strivec]: + """ + Convert a batch of network outputs to 3D representations. + + Args: + x: The [N x * x C] sparse tensor output by the network. + + Returns: + list of representations + """ + ret = [] + for i in range(x.shape[0]): + representation = Strivec( + sh_degree=0, + resolution=self.resolution, + aabb=[-0.5, -0.5, -0.5, 1, 1, 1], + rank=self.rep_config['rank'], + dim=self.rep_config['dim'], + device='cuda', + ) + representation.density_shift = 0.0 + representation.position = (x.coords[x.layout[i]][:, 1:].float() + 0.5) / self.resolution + representation.depth = torch.full((representation.position.shape[0], 1), int(np.log2(self.resolution)), dtype=torch.uint8, device='cuda') + for k, v in self.layout.items(): + setattr(representation, k, x.feats[x.layout[i]][:, v['range'][0]:v['range'][1]].reshape(-1, *v['shape'])) + representation.trivec = representation.trivec + 1 + ret.append(representation) + return ret + + def forward(self, x: sp.SparseTensor) -> List[Strivec]: + h = super().forward(x) + h = h.type(x.dtype) + h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:])) + h = self.out_layer(h) + return self.to_representation(h) + + +class ElasticSLatRadianceFieldDecoder(SparseTransformerElasticMixin, SLatRadianceFieldDecoder): + """ + Slat VAE Radiance Field Decoder with elastic memory management. + Used for training with low VRAM. + """ + pass diff --git a/third_party/TRELLIS/trellis/models/structured_latent_vae/encoder.py b/third_party/TRELLIS/trellis/models/structured_latent_vae/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..77f6fb193ea2bee41fd0b77cd06e37bd14bf6793 --- /dev/null +++ b/third_party/TRELLIS/trellis/models/structured_latent_vae/encoder.py @@ -0,0 +1,80 @@ +from typing import * +import torch +import torch.nn as nn +import torch.nn.functional as F +from ...modules import sparse as sp +from .base import SparseTransformerBase +from ..sparse_elastic_mixin import SparseTransformerElasticMixin + + +class SLatEncoder(SparseTransformerBase): + def __init__( + self, + resolution: int, + in_channels: int, + model_channels: int, + latent_channels: int, + num_blocks: int, + num_heads: Optional[int] = None, + num_head_channels: Optional[int] = 64, + mlp_ratio: float = 4, + attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "swin", + window_size: int = 8, + pe_mode: Literal["ape", "rope"] = "ape", + use_fp16: bool = False, + use_checkpoint: bool = False, + qk_rms_norm: bool = False, + ): + super().__init__( + in_channels=in_channels, + model_channels=model_channels, + num_blocks=num_blocks, + num_heads=num_heads, + num_head_channels=num_head_channels, + mlp_ratio=mlp_ratio, + attn_mode=attn_mode, + window_size=window_size, + pe_mode=pe_mode, + use_fp16=use_fp16, + use_checkpoint=use_checkpoint, + qk_rms_norm=qk_rms_norm, + ) + self.resolution = resolution + self.out_layer = sp.SparseLinear(model_channels, 2 * latent_channels) + + self.initialize_weights() + if use_fp16: + self.convert_to_fp16() + + def initialize_weights(self) -> None: + super().initialize_weights() + # Zero-out output layers: + nn.init.constant_(self.out_layer.weight, 0) + nn.init.constant_(self.out_layer.bias, 0) + + def forward(self, x: sp.SparseTensor, sample_posterior=True, return_raw=False): + h = super().forward(x) + h = h.type(x.dtype) + h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:])) + h = self.out_layer(h) + + # Sample from the posterior distribution + mean, logvar = h.feats.chunk(2, dim=-1) + if sample_posterior: + std = torch.exp(0.5 * logvar) + z = mean + std * torch.randn_like(std) + else: + z = mean + z = h.replace(z) + + if return_raw: + return z, mean, logvar + else: + return z + + +class ElasticSLatEncoder(SparseTransformerElasticMixin, SLatEncoder): + """ + SLat VAE encoder with elastic memory management. + Used for training with low VRAM. + """ diff --git a/third_party/TRELLIS/trellis/modules/attention/__init__.py b/third_party/TRELLIS/trellis/modules/attention/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f452320d5dbc4c0aa1664e33f76c56ff4bbe2039 --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/attention/__init__.py @@ -0,0 +1,36 @@ +from typing import * + +BACKEND = 'flash_attn' +DEBUG = False + +def __from_env(): + import os + + global BACKEND + global DEBUG + + env_attn_backend = os.environ.get('ATTN_BACKEND') + env_sttn_debug = os.environ.get('ATTN_DEBUG') + + if env_attn_backend is not None and env_attn_backend in ['xformers', 'flash_attn', 'sdpa', 'naive']: + BACKEND = env_attn_backend + if env_sttn_debug is not None: + DEBUG = env_sttn_debug == '1' + + print(f"[ATTENTION] Using backend: {BACKEND}") + + +__from_env() + + +def set_backend(backend: Literal['xformers', 'flash_attn']): + global BACKEND + BACKEND = backend + +def set_debug(debug: bool): + global DEBUG + DEBUG = debug + + +from .full_attn import * +from .modules import * diff --git a/third_party/TRELLIS/trellis/modules/attention/full_attn.py b/third_party/TRELLIS/trellis/modules/attention/full_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..d9ebf6380a78906d4c6e969c63223fb7b398e5a7 --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/attention/full_attn.py @@ -0,0 +1,140 @@ +from typing import * +import torch +import math +from . import DEBUG, BACKEND + +if BACKEND == 'xformers': + import xformers.ops as xops +elif BACKEND == 'flash_attn': + import flash_attn +elif BACKEND == 'sdpa': + from torch.nn.functional import scaled_dot_product_attention as sdpa +elif BACKEND == 'naive': + pass +else: + raise ValueError(f"Unknown attention backend: {BACKEND}") + + +__all__ = [ + 'scaled_dot_product_attention', +] + + +def _naive_sdpa(q, k, v): + """ + Naive implementation of scaled dot product attention. + """ + q = q.permute(0, 2, 1, 3) # [N, H, L, C] + k = k.permute(0, 2, 1, 3) # [N, H, L, C] + v = v.permute(0, 2, 1, 3) # [N, H, L, C] + scale_factor = 1 / math.sqrt(q.size(-1)) + attn_weight = q @ k.transpose(-2, -1) * scale_factor + attn_weight = torch.softmax(attn_weight, dim=-1) + out = attn_weight @ v + out = out.permute(0, 2, 1, 3) # [N, L, H, C] + return out + + +@overload +def scaled_dot_product_attention(qkv: torch.Tensor) -> torch.Tensor: + """ + Apply scaled dot product attention. + + Args: + qkv (torch.Tensor): A [N, L, 3, H, C] tensor containing Qs, Ks, and Vs. + """ + ... + +@overload +def scaled_dot_product_attention(q: torch.Tensor, kv: torch.Tensor) -> torch.Tensor: + """ + Apply scaled dot product attention. + + Args: + q (torch.Tensor): A [N, L, H, C] tensor containing Qs. + kv (torch.Tensor): A [N, L, 2, H, C] tensor containing Ks and Vs. + """ + ... + +@overload +def scaled_dot_product_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + """ + Apply scaled dot product attention. + + Args: + q (torch.Tensor): A [N, L, H, Ci] tensor containing Qs. + k (torch.Tensor): A [N, L, H, Ci] tensor containing Ks. + v (torch.Tensor): A [N, L, H, Co] tensor containing Vs. + + Note: + k and v are assumed to have the same coordinate map. + """ + ... + +def scaled_dot_product_attention(*args, **kwargs): + arg_names_dict = { + 1: ['qkv'], + 2: ['q', 'kv'], + 3: ['q', 'k', 'v'] + } + num_all_args = len(args) + len(kwargs) + assert num_all_args in arg_names_dict, f"Invalid number of arguments, got {num_all_args}, expected 1, 2, or 3" + for key in arg_names_dict[num_all_args][len(args):]: + assert key in kwargs, f"Missing argument {key}" + + if num_all_args == 1: + qkv = args[0] if len(args) > 0 else kwargs['qkv'] + assert len(qkv.shape) == 5 and qkv.shape[2] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, L, 3, H, C]" + device = qkv.device + + elif num_all_args == 2: + q = args[0] if len(args) > 0 else kwargs['q'] + kv = args[1] if len(args) > 1 else kwargs['kv'] + assert q.shape[0] == kv.shape[0], f"Batch size mismatch, got {q.shape[0]} and {kv.shape[0]}" + assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, C]" + assert len(kv.shape) == 5, f"Invalid shape for kv, got {kv.shape}, expected [N, L, 2, H, C]" + device = q.device + + elif num_all_args == 3: + q = args[0] if len(args) > 0 else kwargs['q'] + k = args[1] if len(args) > 1 else kwargs['k'] + v = args[2] if len(args) > 2 else kwargs['v'] + assert q.shape[0] == k.shape[0] == v.shape[0], f"Batch size mismatch, got {q.shape[0]}, {k.shape[0]}, and {v.shape[0]}" + assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, Ci]" + assert len(k.shape) == 4, f"Invalid shape for k, got {k.shape}, expected [N, L, H, Ci]" + assert len(v.shape) == 4, f"Invalid shape for v, got {v.shape}, expected [N, L, H, Co]" + device = q.device + + if BACKEND == 'xformers': + if num_all_args == 1: + q, k, v = qkv.unbind(dim=2) + elif num_all_args == 2: + k, v = kv.unbind(dim=2) + out = xops.memory_efficient_attention(q, k, v) + elif BACKEND == 'flash_attn': + if num_all_args == 1: + out = flash_attn.flash_attn_qkvpacked_func(qkv) + elif num_all_args == 2: + out = flash_attn.flash_attn_kvpacked_func(q, kv) + elif num_all_args == 3: + out = flash_attn.flash_attn_func(q, k, v) + elif BACKEND == 'sdpa': + if num_all_args == 1: + q, k, v = qkv.unbind(dim=2) + elif num_all_args == 2: + k, v = kv.unbind(dim=2) + q = q.permute(0, 2, 1, 3) # [N, H, L, C] + k = k.permute(0, 2, 1, 3) # [N, H, L, C] + v = v.permute(0, 2, 1, 3) # [N, H, L, C] + out = sdpa(q, k, v) # [N, H, L, C] + out = out.permute(0, 2, 1, 3) # [N, L, H, C] + elif BACKEND == 'naive': + if num_all_args == 1: + q, k, v = qkv.unbind(dim=2) + elif num_all_args == 2: + k, v = kv.unbind(dim=2) + out = _naive_sdpa(q, k, v) + else: + raise ValueError(f"Unknown attention module: {BACKEND}") + + return out diff --git a/third_party/TRELLIS/trellis/modules/attention/modules.py b/third_party/TRELLIS/trellis/modules/attention/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..dbe6235c27134f0477e48d3e12de3068c6a500ef --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/attention/modules.py @@ -0,0 +1,146 @@ +from typing import * +import torch +import torch.nn as nn +import torch.nn.functional as F +from .full_attn import scaled_dot_product_attention + + +class MultiHeadRMSNorm(nn.Module): + def __init__(self, dim: int, heads: int): + super().__init__() + self.scale = dim ** 0.5 + self.gamma = nn.Parameter(torch.ones(heads, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (F.normalize(x.float(), dim = -1) * self.gamma * self.scale).to(x.dtype) + + +class RotaryPositionEmbedder(nn.Module): + def __init__(self, hidden_size: int, in_channels: int = 3): + super().__init__() + assert hidden_size % 2 == 0, "Hidden size must be divisible by 2" + self.hidden_size = hidden_size + self.in_channels = in_channels + self.freq_dim = hidden_size // in_channels // 2 + self.freqs = torch.arange(self.freq_dim, dtype=torch.float32) / self.freq_dim + self.freqs = 1.0 / (10000 ** self.freqs) + + def _get_phases(self, indices: torch.Tensor) -> torch.Tensor: + self.freqs = self.freqs.to(indices.device) + phases = torch.outer(indices, self.freqs) + phases = torch.polar(torch.ones_like(phases), phases) + return phases + + def _rotary_embedding(self, x: torch.Tensor, phases: torch.Tensor) -> torch.Tensor: + x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + x_rotated = x_complex * phases + x_embed = torch.view_as_real(x_rotated).reshape(*x_rotated.shape[:-1], -1).to(x.dtype) + return x_embed + + def forward(self, q: torch.Tensor, k: torch.Tensor, indices: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Args: + q (sp.SparseTensor): [..., N, D] tensor of queries + k (sp.SparseTensor): [..., N, D] tensor of keys + indices (torch.Tensor): [..., N, C] tensor of spatial positions + """ + if indices is None: + indices = torch.arange(q.shape[-2], device=q.device) + if len(q.shape) > 2: + indices = indices.unsqueeze(0).expand(q.shape[:-2] + (-1,)) + + phases = self._get_phases(indices.reshape(-1)).reshape(*indices.shape[:-1], -1) + if phases.shape[1] < self.hidden_size // 2: + phases = torch.cat([phases, torch.polar( + torch.ones(*phases.shape[:-1], self.hidden_size // 2 - phases.shape[1], device=phases.device), + torch.zeros(*phases.shape[:-1], self.hidden_size // 2 - phases.shape[1], device=phases.device) + )], dim=-1) + q_embed = self._rotary_embedding(q, phases) + k_embed = self._rotary_embedding(k, phases) + return q_embed, k_embed + + +class MultiHeadAttention(nn.Module): + def __init__( + self, + channels: int, + num_heads: int, + ctx_channels: Optional[int]=None, + type: Literal["self", "cross"] = "self", + attn_mode: Literal["full", "windowed"] = "full", + window_size: Optional[int] = None, + shift_window: Optional[Tuple[int, int, int]] = None, + qkv_bias: bool = True, + use_rope: bool = False, + qk_rms_norm: bool = False, + ): + super().__init__() + assert channels % num_heads == 0 + assert type in ["self", "cross"], f"Invalid attention type: {type}" + assert attn_mode in ["full", "windowed"], f"Invalid attention mode: {attn_mode}" + assert type == "self" or attn_mode == "full", "Cross-attention only supports full attention" + + if attn_mode == "windowed": + raise NotImplementedError("Windowed attention is not yet implemented") + + self.channels = channels + self.head_dim = channels // num_heads + self.ctx_channels = ctx_channels if ctx_channels is not None else channels + self.num_heads = num_heads + self._type = type + self.attn_mode = attn_mode + self.window_size = window_size + self.shift_window = shift_window + self.use_rope = use_rope + self.qk_rms_norm = qk_rms_norm + + if self._type == "self": + self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias) + else: + self.to_q = nn.Linear(channels, channels, bias=qkv_bias) + self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias) + + if self.qk_rms_norm: + self.q_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads) + self.k_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads) + + self.to_out = nn.Linear(channels, channels) + + if use_rope: + self.rope = RotaryPositionEmbedder(channels) + + def forward(self, x: torch.Tensor, context: Optional[torch.Tensor] = None, indices: Optional[torch.Tensor] = None) -> torch.Tensor: + B, L, C = x.shape + if self._type == "self": + qkv = self.to_qkv(x) + qkv = qkv.reshape(B, L, 3, self.num_heads, -1) + if self.use_rope: + q, k, v = qkv.unbind(dim=2) + q, k = self.rope(q, k, indices) + qkv = torch.stack([q, k, v], dim=2) + if self.attn_mode == "full": + if self.qk_rms_norm: + q, k, v = qkv.unbind(dim=2) + q = self.q_rms_norm(q) + k = self.k_rms_norm(k) + h = scaled_dot_product_attention(q, k, v) + else: + h = scaled_dot_product_attention(qkv) + elif self.attn_mode == "windowed": + raise NotImplementedError("Windowed attention is not yet implemented") + else: + Lkv = context.shape[1] + q = self.to_q(x) + kv = self.to_kv(context) + q = q.reshape(B, L, self.num_heads, -1) + kv = kv.reshape(B, Lkv, 2, self.num_heads, -1) + if self.qk_rms_norm: + q = self.q_rms_norm(q) + k, v = kv.unbind(dim=2) + k = self.k_rms_norm(k) + h = scaled_dot_product_attention(q, k, v) + else: + h = scaled_dot_product_attention(q, kv) + h = h.reshape(B, L, -1) + h = self.to_out(h) + return h diff --git a/third_party/TRELLIS/trellis/modules/norm.py b/third_party/TRELLIS/trellis/modules/norm.py new file mode 100644 index 0000000000000000000000000000000000000000..09035726081fb7afda2c62504d5474cfa483c58f --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/norm.py @@ -0,0 +1,25 @@ +import torch +import torch.nn as nn + + +class LayerNorm32(nn.LayerNorm): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return super().forward(x.float()).type(x.dtype) + + +class GroupNorm32(nn.GroupNorm): + """ + A GroupNorm layer that converts to float32 before the forward pass. + """ + def forward(self, x: torch.Tensor) -> torch.Tensor: + return super().forward(x.float()).type(x.dtype) + + +class ChannelLayerNorm32(LayerNorm32): + def forward(self, x: torch.Tensor) -> torch.Tensor: + DIM = x.dim() + x = x.permute(0, *range(2, DIM), 1).contiguous() + x = super().forward(x) + x = x.permute(0, DIM-1, *range(1, DIM-1)).contiguous() + return x + \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/modules/sparse/__init__.py b/third_party/TRELLIS/trellis/modules/sparse/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..726756c16dcfe0f04de0d2ea5bdce499fa220160 --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/__init__.py @@ -0,0 +1,102 @@ +from typing import * + +BACKEND = 'spconv' +DEBUG = False +ATTN = 'flash_attn' + +def __from_env(): + import os + + global BACKEND + global DEBUG + global ATTN + + env_sparse_backend = os.environ.get('SPARSE_BACKEND') + env_sparse_debug = os.environ.get('SPARSE_DEBUG') + env_sparse_attn = os.environ.get('SPARSE_ATTN_BACKEND') + if env_sparse_attn is None: + env_sparse_attn = os.environ.get('ATTN_BACKEND') + + if env_sparse_backend is not None and env_sparse_backend in ['spconv', 'torchsparse']: + BACKEND = env_sparse_backend + if env_sparse_debug is not None: + DEBUG = env_sparse_debug == '1' + if env_sparse_attn is not None and env_sparse_attn in ['xformers', 'flash_attn']: + ATTN = env_sparse_attn + + print(f"[SPARSE] Backend: {BACKEND}, Attention: {ATTN}") + + +__from_env() + + +def set_backend(backend: Literal['spconv', 'torchsparse']): + global BACKEND + BACKEND = backend + +def set_debug(debug: bool): + global DEBUG + DEBUG = debug + +def set_attn(attn: Literal['xformers', 'flash_attn']): + global ATTN + ATTN = attn + + +import importlib + +__attributes = { + 'SparseTensor': 'basic', + 'sparse_batch_broadcast': 'basic', + 'sparse_batch_op': 'basic', + 'sparse_cat': 'basic', + 'sparse_unbind': 'basic', + 'SparseGroupNorm': 'norm', + 'SparseLayerNorm': 'norm', + 'SparseGroupNorm32': 'norm', + 'SparseLayerNorm32': 'norm', + 'SparseReLU': 'nonlinearity', + 'SparseSiLU': 'nonlinearity', + 'SparseGELU': 'nonlinearity', + 'SparseActivation': 'nonlinearity', + 'SparseLinear': 'linear', + 'sparse_scaled_dot_product_attention': 'attention', + 'SerializeMode': 'attention', + 'sparse_serialized_scaled_dot_product_self_attention': 'attention', + 'sparse_windowed_scaled_dot_product_self_attention': 'attention', + 'SparseMultiHeadAttention': 'attention', + 'SparseConv3d': 'conv', + 'SparseInverseConv3d': 'conv', + 'SparseDownsample': 'spatial', + 'SparseUpsample': 'spatial', + 'SparseSubdivide' : 'spatial' +} + +__submodules = ['transformer'] + +__all__ = list(__attributes.keys()) + __submodules + +def __getattr__(name): + if name not in globals(): + if name in __attributes: + module_name = __attributes[name] + module = importlib.import_module(f".{module_name}", __name__) + globals()[name] = getattr(module, name) + elif name in __submodules: + module = importlib.import_module(f".{name}", __name__) + globals()[name] = module + else: + raise AttributeError(f"module {__name__} has no attribute {name}") + return globals()[name] + + +# For Pylance +if __name__ == '__main__': + from .basic import * + from .norm import * + from .nonlinearity import * + from .linear import * + from .attention import * + from .conv import * + from .spatial import * + import transformer diff --git a/third_party/TRELLIS/trellis/modules/sparse/attention/__init__.py b/third_party/TRELLIS/trellis/modules/sparse/attention/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..32b3c2c837c613e41755ac4c85f9ed057a6f5bfb --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/attention/__init__.py @@ -0,0 +1,4 @@ +from .full_attn import * +from .serialized_attn import * +from .windowed_attn import * +from .modules import * diff --git a/third_party/TRELLIS/trellis/modules/sparse/attention/full_attn.py b/third_party/TRELLIS/trellis/modules/sparse/attention/full_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..d794684c4c70ad4ae6fca95465c471214e9e8e61 --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/attention/full_attn.py @@ -0,0 +1,347 @@ +from typing import * +import torch +from .. import SparseTensor +from .. import DEBUG, ATTN + +if ATTN == 'xformers': + import xformers.ops as xops +elif ATTN == 'flash_attn': + import flash_attn +else: + raise ValueError(f"Unknown attention module: {ATTN}") + + +__all__ = [ + 'sparse_scaled_dot_product_attention', + 'sparse_block_diagonal_self_attention', + 'sparse_region_boost_self_attention', +] + + +def sparse_block_diagonal_self_attention(qkv: SparseTensor, group_ids: torch.Tensor) -> SparseTensor: + """Full self-attention restricted so each voxel only attends within its own (batch, group). + + Used for mixed text/image conditioning: voxels of one modality's region are made blind to the + other region's voxels, so a strongly-conditioned region cannot bleed into a cfg=0 region through + global self-attention. Implemented as block-diagonal attention over the packed sequence (via + variable-length sequences), which is memory-safe (no N x N mask) and cheaper than full attention. + + Args: + qkv (SparseTensor): [N, 3, H, C] packed q/k/v (RoPE/rms already applied by the caller). + group_ids (torch.Tensor): [T] long per-voxel group id, aligned to qkv.feats rows. + """ + feats = qkv.feats # [T, 3, H, C] + T = feats.shape[0] + device = feats.device + if T == 0: + return qkv.replace(feats[:, 0]) + + # Block key = (batch index, group id); different batch items or groups never attend together. + batch_idx = qkv.coords[:, 0].long() + num_groups = int(group_ids.max().item()) + 1 + block_key = batch_idx * num_groups + group_ids.long() + + # Reorder so each block is contiguous (order within a block is irrelevant given baked RoPE). + perm = torch.argsort(block_key) + inv = torch.empty_like(perm) + inv[perm] = torch.arange(T, device=device) + feats_sorted = feats[perm] + keys_sorted = block_key[perm] + _, counts = torch.unique_consecutive(keys_sorted, return_counts=True) + + if ATTN == 'flash_attn': + cu_seqlens = torch.cat([counts.new_zeros(1), counts.cumsum(0)]).to(torch.int32) + out_sorted = flash_attn.flash_attn_varlen_qkvpacked_func( + feats_sorted, cu_seqlens, int(counts.max().item())) + elif ATTN == 'xformers': + q, k, v = feats_sorted.unbind(dim=1) # each [T, H, C] + mask = xops.fmha.BlockDiagonalMask.from_seqlens(counts.tolist()) + out_sorted = xops.memory_efficient_attention( + q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0), mask)[0] + else: + raise ValueError(f"Unknown attention module: {ATTN}") + + out = out_sorted[inv] + return qkv.replace(out) + + +def _lse_to_token_major(lse: torch.Tensor, tokens: int, heads: int) -> torch.Tensor: + """Normalize flash-attn's varlen log-sum-exp to [T, H, 1] float. + + It comes back as [heads, tokens], and some builds pad the token axis out to a multiple + of the kernel's block size, so trim to the real token count after transposing. + """ + if lse.shape[0] == heads and lse.shape[0] != tokens: + lse = lse.transpose(0, 1) + return lse[:tokens].unsqueeze(-1).float() + + +def sparse_region_boost_self_attention(qkv: SparseTensor, group_ids: torch.Tensor, boost: float) -> SparseTensor: + """Full self-attention with in-region attention weights multiplied by `boost`. + + The soft counterpart of :func:`sparse_block_diagonal_self_attention`: a voxel still sees the + whole object, so parts stay globally coherent, but it weighs its own region `boost` times + more heavily. That is what keeps one condition's identity from spilling into a region that is + unguided (cfg=0) on the current step of alternating text/image conditioning, without the + independently-generated look that hard blocking produces. + + Multiplying the in-region weights by `b` is exactly adding `log b` to the in-region logits + before the softmax. Materializing that bias would need an O(T^2) score matrix in every block + of every step; it is not necessary, because the result is recoverable from two flash calls + and their log-sum-exps: + + out = (b*Z_in*out_in + (Z_all*out_all - Z_in*out_in)) / (b*Z_in + (Z_all - Z_in)) + + where `Z = exp(lse)` is each softmax's unnormalised mass and `Z_in <= Z_all` always. + + Args: + qkv (SparseTensor): [N, 3, H, C] packed q/k/v (RoPE/rms already applied by the caller). + group_ids (torch.Tensor): [T] long per-voxel region id, aligned to qkv.feats rows. + boost (float): in-region weight multiplier; 1.0 is plain full attention. + """ + feats = qkv.feats # [T, 3, H, C] + T = feats.shape[0] + device = feats.device + if T == 0: + return qkv.replace(feats[:, 0]) + if boost == 1.0: + return sparse_scaled_dot_product_attention(qkv) + if ATTN != 'flash_attn': + raise NotImplementedError( + f"region-boosted self-attention needs the flash_attn backend for its log-sum-exp, " + f"but the sparse attention backend is {ATTN!r}. Set ATTN=flash_attn, or disable the " + f"boost (sim_guidance.mixed_self_attn_boost: 1.0) to fall back to stock attention." + ) + + # Block key = (batch index, region id), so sorting makes each region contiguous AND, because + # batch is the major key, keeps each batch item contiguous too - both segmentations below + # read straight off the same permutation. + batch_idx = qkv.coords[:, 0].long() + num_groups = int(group_ids.max().item()) + 1 + block_key = batch_idx * num_groups + group_ids.long() + + perm = torch.argsort(block_key) + inv = torch.empty_like(perm) + inv[perm] = torch.arange(T, device=device) + feats_sorted = feats[perm].contiguous() + heads = feats_sorted.shape[2] + + # Restricted softmax: each region attends only within itself. + _, region_counts = torch.unique_consecutive(block_key[perm], return_counts=True) + cu_in = torch.cat([region_counts.new_zeros(1), region_counts.cumsum(0)]).to(torch.int32) + out_in, lse_in, _ = flash_attn.flash_attn_varlen_qkvpacked_func( + feats_sorted, cu_in, int(region_counts.max().item()), return_attn_probs=True) + + # Unrestricted softmax over the same permuted layout: blocks are whole batch items. + _, batch_counts = torch.unique_consecutive(batch_idx[perm], return_counts=True) + cu_all = torch.cat([batch_counts.new_zeros(1), batch_counts.cumsum(0)]).to(torch.int32) + out_all, lse_all, _ = flash_attn.flash_attn_varlen_qkvpacked_func( + feats_sorted, cu_all, int(batch_counts.max().item()), return_attn_probs=True) + + # Work relative to lse_all for numerical stability; Z_in / Z_all lands in (0, 1]. + li = _lse_to_token_major(lse_in, T, heads) + la = _lse_to_token_major(lse_all, T, heads) + w_in = torch.exp(li - la) + numer = (boost - 1.0) * w_in * out_in.float() + out_all.float() + denom = (boost - 1.0) * w_in + 1.0 + out_sorted = (numer / denom).to(feats.dtype) + return qkv.replace(out_sorted[inv]) + + +@overload +def sparse_scaled_dot_product_attention(qkv: SparseTensor) -> SparseTensor: + """ + Apply scaled dot product attention to a sparse tensor. + + Args: + qkv (SparseTensor): A [N, *, 3, H, C] sparse tensor containing Qs, Ks, and Vs. + """ + ... + +@overload +def sparse_scaled_dot_product_attention(q: SparseTensor, kv: Union[SparseTensor, torch.Tensor]) -> SparseTensor: + """ + Apply scaled dot product attention to a sparse tensor. + + Args: + q (SparseTensor): A [N, *, H, C] sparse tensor containing Qs. + kv (SparseTensor or torch.Tensor): A [N, *, 2, H, C] sparse tensor or a [N, L, 2, H, C] dense tensor containing Ks and Vs. + """ + ... + +@overload +def sparse_scaled_dot_product_attention(q: torch.Tensor, kv: SparseTensor) -> torch.Tensor: + """ + Apply scaled dot product attention to a sparse tensor. + + Args: + q (SparseTensor): A [N, L, H, C] dense tensor containing Qs. + kv (SparseTensor or torch.Tensor): A [N, *, 2, H, C] sparse tensor containing Ks and Vs. + """ + ... + +@overload +def sparse_scaled_dot_product_attention(q: SparseTensor, k: SparseTensor, v: SparseTensor) -> SparseTensor: + """ + Apply scaled dot product attention to a sparse tensor. + + Args: + q (SparseTensor): A [N, *, H, Ci] sparse tensor containing Qs. + k (SparseTensor): A [N, *, H, Ci] sparse tensor containing Ks. + v (SparseTensor): A [N, *, H, Co] sparse tensor containing Vs. + + Note: + k and v are assumed to have the same coordinate map. + """ + ... + +@overload +def sparse_scaled_dot_product_attention(q: SparseTensor, k: torch.Tensor, v: torch.Tensor) -> SparseTensor: + """ + Apply scaled dot product attention to a sparse tensor. + + Args: + q (SparseTensor): A [N, *, H, Ci] sparse tensor containing Qs. + k (torch.Tensor): A [N, L, H, Ci] dense tensor containing Ks. + v (torch.Tensor): A [N, L, H, Co] dense tensor containing Vs. + """ + ... + +@overload +def sparse_scaled_dot_product_attention(q: torch.Tensor, k: SparseTensor, v: SparseTensor) -> torch.Tensor: + """ + Apply scaled dot product attention to a sparse tensor. + + Args: + q (torch.Tensor): A [N, L, H, Ci] dense tensor containing Qs. + k (SparseTensor): A [N, *, H, Ci] sparse tensor containing Ks. + v (SparseTensor): A [N, *, H, Co] sparse tensor containing Vs. + """ + ... + +def sparse_scaled_dot_product_attention(*args, **kwargs): + arg_names_dict = { + 1: ['qkv'], + 2: ['q', 'kv'], + 3: ['q', 'k', 'v'] + } + num_all_args = len(args) + len(kwargs) + assert num_all_args in arg_names_dict, f"Invalid number of arguments, got {num_all_args}, expected 1, 2, or 3" + for key in arg_names_dict[num_all_args][len(args):]: + assert key in kwargs, f"Missing argument {key}" + + if num_all_args == 1: + qkv = args[0] if len(args) > 0 else kwargs['qkv'] + assert isinstance(qkv, SparseTensor), f"qkv must be a SparseTensor, got {type(qkv)}" + assert len(qkv.shape) == 4 and qkv.shape[1] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, *, 3, H, C]" + device = qkv.device + + s = qkv + q_seqlen = [qkv.layout[i].stop - qkv.layout[i].start for i in range(qkv.shape[0])] + kv_seqlen = q_seqlen + qkv = qkv.feats # [T, 3, H, C] + + elif num_all_args == 2: + q = args[0] if len(args) > 0 else kwargs['q'] + kv = args[1] if len(args) > 1 else kwargs['kv'] + assert isinstance(q, SparseTensor) and isinstance(kv, (SparseTensor, torch.Tensor)) or \ + isinstance(q, torch.Tensor) and isinstance(kv, SparseTensor), \ + f"Invalid types, got {type(q)} and {type(kv)}" + assert q.shape[0] == kv.shape[0], f"Batch size mismatch, got {q.shape[0]} and {kv.shape[0]}" + device = q.device + + if isinstance(q, SparseTensor): + assert len(q.shape) == 3, f"Invalid shape for q, got {q.shape}, expected [N, *, H, C]" + s = q + q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])] + q = q.feats # [T_Q, H, C] + else: + assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, C]" + s = None + N, L, H, C = q.shape + q_seqlen = [L] * N + q = q.reshape(N * L, H, C) # [T_Q, H, C] + + if isinstance(kv, SparseTensor): + assert len(kv.shape) == 4 and kv.shape[1] == 2, f"Invalid shape for kv, got {kv.shape}, expected [N, *, 2, H, C]" + kv_seqlen = [kv.layout[i].stop - kv.layout[i].start for i in range(kv.shape[0])] + kv = kv.feats # [T_KV, 2, H, C] + else: + assert len(kv.shape) == 5, f"Invalid shape for kv, got {kv.shape}, expected [N, L, 2, H, C]" + N, L, _, H, C = kv.shape + kv_seqlen = [L] * N + kv = kv.reshape(N * L, 2, H, C) # [T_KV, 2, H, C] + + elif num_all_args == 3: + q = args[0] if len(args) > 0 else kwargs['q'] + k = args[1] if len(args) > 1 else kwargs['k'] + v = args[2] if len(args) > 2 else kwargs['v'] + assert isinstance(q, SparseTensor) and isinstance(k, (SparseTensor, torch.Tensor)) and type(k) == type(v) or \ + isinstance(q, torch.Tensor) and isinstance(k, SparseTensor) and isinstance(v, SparseTensor), \ + f"Invalid types, got {type(q)}, {type(k)}, and {type(v)}" + assert q.shape[0] == k.shape[0] == v.shape[0], f"Batch size mismatch, got {q.shape[0]}, {k.shape[0]}, and {v.shape[0]}" + device = q.device + + if isinstance(q, SparseTensor): + assert len(q.shape) == 3, f"Invalid shape for q, got {q.shape}, expected [N, *, H, Ci]" + s = q + q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])] + q = q.feats # [T_Q, H, Ci] + else: + assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, Ci]" + s = None + N, L, H, CI = q.shape + q_seqlen = [L] * N + q = q.reshape(N * L, H, CI) # [T_Q, H, Ci] + + if isinstance(k, SparseTensor): + assert len(k.shape) == 3, f"Invalid shape for k, got {k.shape}, expected [N, *, H, Ci]" + assert len(v.shape) == 3, f"Invalid shape for v, got {v.shape}, expected [N, *, H, Co]" + kv_seqlen = [k.layout[i].stop - k.layout[i].start for i in range(k.shape[0])] + k = k.feats # [T_KV, H, Ci] + v = v.feats # [T_KV, H, Co] + else: + assert len(k.shape) == 4, f"Invalid shape for k, got {k.shape}, expected [N, L, H, Ci]" + assert len(v.shape) == 4, f"Invalid shape for v, got {v.shape}, expected [N, L, H, Co]" + N, L, H, CI, CO = *k.shape, v.shape[-1] + kv_seqlen = [L] * N + k = k.reshape(N * L, H, CI) # [T_KV, H, Ci] + v = v.reshape(N * L, H, CO) # [T_KV, H, Co] + + if DEBUG: + if s is not None: + for i in range(s.shape[0]): + assert (s.coords[s.layout[i]] == i).all(), f"SparseScaledDotProductSelfAttention: batch index mismatch" + if num_all_args in [2, 3]: + assert q.shape[:2] == [1, sum(q_seqlen)], f"SparseScaledDotProductSelfAttention: q shape mismatch" + if num_all_args == 3: + assert k.shape[:2] == [1, sum(kv_seqlen)], f"SparseScaledDotProductSelfAttention: k shape mismatch" + assert v.shape[:2] == [1, sum(kv_seqlen)], f"SparseScaledDotProductSelfAttention: v shape mismatch" + + if ATTN == 'xformers': + if num_all_args == 1: + q, k, v = qkv.unbind(dim=1) + elif num_all_args == 2: + k, v = kv.unbind(dim=1) + q = q.unsqueeze(0) + k = k.unsqueeze(0) + v = v.unsqueeze(0) + mask = xops.fmha.BlockDiagonalMask.from_seqlens(q_seqlen, kv_seqlen) + out = xops.memory_efficient_attention(q, k, v, mask)[0] + elif ATTN == 'flash_attn': + cu_seqlens_q = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(q_seqlen), dim=0)]).int().to(device) + if num_all_args in [2, 3]: + cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device) + if num_all_args == 1: + out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv, cu_seqlens_q, max(q_seqlen)) + elif num_all_args == 2: + out = flash_attn.flash_attn_varlen_kvpacked_func(q, kv, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen)) + elif num_all_args == 3: + out = flash_attn.flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen)) + else: + raise ValueError(f"Unknown attention module: {ATTN}") + + if s is not None: + return s.replace(out) + else: + return out.reshape(N, L, H, -1) diff --git a/third_party/TRELLIS/trellis/modules/sparse/attention/modules.py b/third_party/TRELLIS/trellis/modules/sparse/attention/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..124dbcd8c4988df357e8dc230a97c41e17838543 --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/attention/modules.py @@ -0,0 +1,183 @@ +from typing import * +import torch +import torch.nn as nn +import torch.nn.functional as F +from .. import SparseTensor +from .full_attn import sparse_scaled_dot_product_attention, sparse_block_diagonal_self_attention, sparse_region_boost_self_attention +from .serialized_attn import SerializeMode, sparse_serialized_scaled_dot_product_self_attention +from .windowed_attn import sparse_windowed_scaled_dot_product_self_attention +from ...attention import RotaryPositionEmbedder + + +class SparseMultiHeadRMSNorm(nn.Module): + def __init__(self, dim: int, heads: int): + super().__init__() + self.scale = dim ** 0.5 + self.gamma = nn.Parameter(torch.ones(heads, dim)) + + def forward(self, x: Union[SparseTensor, torch.Tensor]) -> Union[SparseTensor, torch.Tensor]: + x_type = x.dtype + x = x.float() + if isinstance(x, SparseTensor): + x = x.replace(F.normalize(x.feats, dim=-1)) + else: + x = F.normalize(x, dim=-1) + return (x * self.gamma * self.scale).to(x_type) + + +class SparseMultiHeadAttention(nn.Module): + def __init__( + self, + channels: int, + num_heads: int, + ctx_channels: Optional[int] = None, + type: Literal["self", "cross"] = "self", + attn_mode: Literal["full", "serialized", "windowed"] = "full", + window_size: Optional[int] = None, + shift_sequence: Optional[int] = None, + shift_window: Optional[Tuple[int, int, int]] = None, + serialize_mode: Optional[SerializeMode] = None, + qkv_bias: bool = True, + use_rope: bool = False, + qk_rms_norm: bool = False, + ): + super().__init__() + assert channels % num_heads == 0 + assert type in ["self", "cross"], f"Invalid attention type: {type}" + assert attn_mode in ["full", "serialized", "windowed"], f"Invalid attention mode: {attn_mode}" + assert type == "self" or attn_mode == "full", "Cross-attention only supports full attention" + assert type == "self" or use_rope is False, "Rotary position embeddings only supported for self-attention" + self.channels = channels + self.ctx_channels = ctx_channels if ctx_channels is not None else channels + self.num_heads = num_heads + self._type = type + self.attn_mode = attn_mode + self.window_size = window_size + self.shift_sequence = shift_sequence + self.shift_window = shift_window + self.serialize_mode = serialize_mode + self.use_rope = use_rope + self.qk_rms_norm = qk_rms_norm + + if self._type == "self": + self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias) + else: + self.to_q = nn.Linear(channels, channels, bias=qkv_bias) + self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias) + + if self.qk_rms_norm: + self.q_rms_norm = SparseMultiHeadRMSNorm(channels // num_heads, num_heads) + self.k_rms_norm = SparseMultiHeadRMSNorm(channels // num_heads, num_heads) + + self.to_out = nn.Linear(channels, channels) + + if use_rope: + self.rope = RotaryPositionEmbedder(channels) + + @staticmethod + def _linear(module: nn.Linear, x: Union[SparseTensor, torch.Tensor]) -> Union[SparseTensor, torch.Tensor]: + if isinstance(x, SparseTensor): + return x.replace(module(x.feats)) + else: + return module(x) + + @staticmethod + def _reshape_chs(x: Union[SparseTensor, torch.Tensor], shape: Tuple[int, ...]) -> Union[SparseTensor, torch.Tensor]: + if isinstance(x, SparseTensor): + return x.reshape(*shape) + else: + return x.reshape(*x.shape[:2], *shape) + + def _fused_pre(self, x: Union[SparseTensor, torch.Tensor], num_fused: int) -> Union[SparseTensor, torch.Tensor]: + if isinstance(x, SparseTensor): + x_feats = x.feats.unsqueeze(0) + else: + x_feats = x + x_feats = x_feats.reshape(*x_feats.shape[:2], num_fused, self.num_heads, -1) + return x.replace(x_feats.squeeze(0)) if isinstance(x, SparseTensor) else x_feats + + def _rope(self, qkv: SparseTensor) -> SparseTensor: + q, k, v = qkv.feats.unbind(dim=1) # [T, H, C] + q, k = self.rope(q, k, qkv.coords[:, 1:]) + qkv = qkv.replace(torch.stack([q, k, v], dim=1)) + return qkv + + def forward(self, x: Union[SparseTensor, torch.Tensor], context: Optional[Union[SparseTensor, torch.Tensor]] = None, context_list=None, coords_dense_indices=None, self_attn_region=None, self_attn_region_boost=None) -> Union[SparseTensor, torch.Tensor]: + if self._type == "self": + qkv = self._linear(self.to_qkv, x) + qkv = self._fused_pre(qkv, num_fused=3) + if self.use_rope: + qkv = self._rope(qkv) + if self.qk_rms_norm: + q, k, v = qkv.unbind(dim=1) + q = self.q_rms_norm(q) + k = self.k_rms_norm(k) + qkv = qkv.replace(torch.stack([q.feats, k.feats, v.feats], dim=1)) + if self.attn_mode == "full": + if self_attn_region is not None and isinstance(qkv, SparseTensor): + # Mixed conditioning: keep cross-region self-attention from letting a strongly + # conditioned region bleed into a cfg=0 region. self_attn_region is a + # (1,1,R,R,R) grid giving each voxel's region id, looked up by its coords + # (same mechanism as coords_dense_indices in cross-attention). + group_ids = self_attn_region[ + 0, 0, qkv.coords[:, 1], qkv.coords[:, 2], qkv.coords[:, 3]] + if self_attn_region_boost is None: + # Hard: a voxel cannot see other regions at all. + h = sparse_block_diagonal_self_attention(qkv, group_ids) + else: + # Soft: still full attention, but in-region weights are upweighted + # before the softmax, so parts stay globally coherent. + h = sparse_region_boost_self_attention(qkv, group_ids, self_attn_region_boost) + else: + h = sparse_scaled_dot_product_attention(qkv) + elif self.attn_mode == "serialized": + h = sparse_serialized_scaled_dot_product_self_attention( + qkv, self.window_size, serialize_mode=self.serialize_mode, shift_sequence=self.shift_sequence, shift_window=self.shift_window + ) + elif self.attn_mode == "windowed": + h = sparse_windowed_scaled_dot_product_self_attention( + qkv, self.window_size, shift_window=self.shift_window + ) + else: + q = self._linear(self.to_q, x) + q = self._reshape_chs(q, (self.num_heads, -1)) + + if context_list is not None: + # Per-superquadric local routing: compute attention for every embedding in + # context_list, then assign each voxel to its corresponding SQ output. + # context_list[0] = global/fallback embedding + # context_list[i+1] = local embedding for SQ i + # coords_dense_indices: (1,1,32,32,32) int tensor, 0=unassigned, 1..N=SQ index + hs = [] + for ctx in context_list: + kv = self._linear(self.to_kv, ctx.to(context.dtype)) + kv = self._fused_pre(kv, num_fused=2) + if self.qk_rms_norm: + q_n = self.q_rms_norm(q) + k, v = kv.unbind(dim=1) + k = self.k_rms_norm(k) + kv = kv.replace(torch.stack([k.feats, v.feats], dim=1)) + hs.append(sparse_scaled_dot_product_attention(q_n, kv)) + else: + hs.append(sparse_scaled_dot_product_attention(q, kv)) + indices = coords_dense_indices[ + 0, 0, q.coords[:, 1], q.coords[:, 2], q.coords[:, 3]] + # Vectorised routing: stack all per-context outputs and select + # the correct one for each voxel in a single gather operation. + # hs_feats: (num_ctx, N, ...) indices: (N,) + hs_feats = torch.stack([hi.feats for hi in hs], dim=0) + vox_idx = torch.arange(hs_feats.shape[1], device=hs_feats.device) + selected = hs_feats[indices.long(), vox_idx] + h = hs[0].replace(selected) + else: + kv = self._linear(self.to_kv, context) + kv = self._fused_pre(kv, num_fused=2) + if self.qk_rms_norm: + q = self.q_rms_norm(q) + k, v = kv.unbind(dim=1) + k = self.k_rms_norm(k) + kv = kv.replace(torch.stack([k.feats, v.feats], dim=1)) + h = sparse_scaled_dot_product_attention(q, kv) + h = self._reshape_chs(h, (-1,)) + h = self._linear(self.to_out, h) + return h diff --git a/third_party/TRELLIS/trellis/modules/sparse/attention/serialized_attn.py b/third_party/TRELLIS/trellis/modules/sparse/attention/serialized_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..5950b75b2f5a6d6e79ab6d472b8501aaa5ec4a26 --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/attention/serialized_attn.py @@ -0,0 +1,193 @@ +from typing import * +from enum import Enum +import torch +import math +from .. import SparseTensor +from .. import DEBUG, ATTN + +if ATTN == 'xformers': + import xformers.ops as xops +elif ATTN == 'flash_attn': + import flash_attn +else: + raise ValueError(f"Unknown attention module: {ATTN}") + + +__all__ = [ + 'sparse_serialized_scaled_dot_product_self_attention', +] + + +class SerializeMode(Enum): + Z_ORDER = 0 + Z_ORDER_TRANSPOSED = 1 + HILBERT = 2 + HILBERT_TRANSPOSED = 3 + + +SerializeModes = [ + SerializeMode.Z_ORDER, + SerializeMode.Z_ORDER_TRANSPOSED, + SerializeMode.HILBERT, + SerializeMode.HILBERT_TRANSPOSED +] + + +def calc_serialization( + tensor: SparseTensor, + window_size: int, + serialize_mode: SerializeMode = SerializeMode.Z_ORDER, + shift_sequence: int = 0, + shift_window: Tuple[int, int, int] = (0, 0, 0) +) -> Tuple[torch.Tensor, torch.Tensor, List[int]]: + """ + Calculate serialization and partitioning for a set of coordinates. + + Args: + tensor (SparseTensor): The input tensor. + window_size (int): The window size to use. + serialize_mode (SerializeMode): The serialization mode to use. + shift_sequence (int): The shift of serialized sequence. + shift_window (Tuple[int, int, int]): The shift of serialized coordinates. + + Returns: + (torch.Tensor, torch.Tensor): Forwards and backwards indices. + """ + fwd_indices = [] + bwd_indices = [] + seq_lens = [] + seq_batch_indices = [] + offsets = [0] + + if 'vox2seq' not in globals(): + import vox2seq + + # Serialize the input + serialize_coords = tensor.coords[:, 1:].clone() + serialize_coords += torch.tensor(shift_window, dtype=torch.int32, device=tensor.device).reshape(1, 3) + if serialize_mode == SerializeMode.Z_ORDER: + code = vox2seq.encode(serialize_coords, mode='z_order', permute=[0, 1, 2]) + elif serialize_mode == SerializeMode.Z_ORDER_TRANSPOSED: + code = vox2seq.encode(serialize_coords, mode='z_order', permute=[1, 0, 2]) + elif serialize_mode == SerializeMode.HILBERT: + code = vox2seq.encode(serialize_coords, mode='hilbert', permute=[0, 1, 2]) + elif serialize_mode == SerializeMode.HILBERT_TRANSPOSED: + code = vox2seq.encode(serialize_coords, mode='hilbert', permute=[1, 0, 2]) + else: + raise ValueError(f"Unknown serialize mode: {serialize_mode}") + + for bi, s in enumerate(tensor.layout): + num_points = s.stop - s.start + num_windows = (num_points + window_size - 1) // window_size + valid_window_size = num_points / num_windows + to_ordered = torch.argsort(code[s.start:s.stop]) + if num_windows == 1: + fwd_indices.append(to_ordered) + bwd_indices.append(torch.zeros_like(to_ordered).scatter_(0, to_ordered, torch.arange(num_points, device=tensor.device))) + fwd_indices[-1] += s.start + bwd_indices[-1] += offsets[-1] + seq_lens.append(num_points) + seq_batch_indices.append(bi) + offsets.append(offsets[-1] + seq_lens[-1]) + else: + # Partition the input + offset = 0 + mids = [(i + 0.5) * valid_window_size + shift_sequence for i in range(num_windows)] + split = [math.floor(i * valid_window_size + shift_sequence) for i in range(num_windows + 1)] + bwd_index = torch.zeros((num_points,), dtype=torch.int64, device=tensor.device) + for i in range(num_windows): + mid = mids[i] + valid_start = split[i] + valid_end = split[i + 1] + padded_start = math.floor(mid - 0.5 * window_size) + padded_end = padded_start + window_size + fwd_indices.append(to_ordered[torch.arange(padded_start, padded_end, device=tensor.device) % num_points]) + offset += valid_start - padded_start + bwd_index.scatter_(0, fwd_indices[-1][valid_start-padded_start:valid_end-padded_start], torch.arange(offset, offset + valid_end - valid_start, device=tensor.device)) + offset += padded_end - valid_start + fwd_indices[-1] += s.start + seq_lens.extend([window_size] * num_windows) + seq_batch_indices.extend([bi] * num_windows) + bwd_indices.append(bwd_index + offsets[-1]) + offsets.append(offsets[-1] + num_windows * window_size) + + fwd_indices = torch.cat(fwd_indices) + bwd_indices = torch.cat(bwd_indices) + + return fwd_indices, bwd_indices, seq_lens, seq_batch_indices + + +def sparse_serialized_scaled_dot_product_self_attention( + qkv: SparseTensor, + window_size: int, + serialize_mode: SerializeMode = SerializeMode.Z_ORDER, + shift_sequence: int = 0, + shift_window: Tuple[int, int, int] = (0, 0, 0) +) -> SparseTensor: + """ + Apply serialized scaled dot product self attention to a sparse tensor. + + Args: + qkv (SparseTensor): [N, *, 3, H, C] sparse tensor containing Qs, Ks, and Vs. + window_size (int): The window size to use. + serialize_mode (SerializeMode): The serialization mode to use. + shift_sequence (int): The shift of serialized sequence. + shift_window (Tuple[int, int, int]): The shift of serialized coordinates. + shift (int): The shift to use. + """ + assert len(qkv.shape) == 4 and qkv.shape[1] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, *, 3, H, C]" + + serialization_spatial_cache_name = f'serialization_{serialize_mode}_{window_size}_{shift_sequence}_{shift_window}' + serialization_spatial_cache = qkv.get_spatial_cache(serialization_spatial_cache_name) + if serialization_spatial_cache is None: + fwd_indices, bwd_indices, seq_lens, seq_batch_indices = calc_serialization(qkv, window_size, serialize_mode, shift_sequence, shift_window) + qkv.register_spatial_cache(serialization_spatial_cache_name, (fwd_indices, bwd_indices, seq_lens, seq_batch_indices)) + else: + fwd_indices, bwd_indices, seq_lens, seq_batch_indices = serialization_spatial_cache + + M = fwd_indices.shape[0] + T = qkv.feats.shape[0] + H = qkv.feats.shape[2] + C = qkv.feats.shape[3] + + qkv_feats = qkv.feats[fwd_indices] # [M, 3, H, C] + + if DEBUG: + start = 0 + qkv_coords = qkv.coords[fwd_indices] + for i in range(len(seq_lens)): + assert (qkv_coords[start:start+seq_lens[i], 0] == seq_batch_indices[i]).all(), f"SparseWindowedScaledDotProductSelfAttention: batch index mismatch" + start += seq_lens[i] + + if all([seq_len == window_size for seq_len in seq_lens]): + B = len(seq_lens) + N = window_size + qkv_feats = qkv_feats.reshape(B, N, 3, H, C) + if ATTN == 'xformers': + q, k, v = qkv_feats.unbind(dim=2) # [B, N, H, C] + out = xops.memory_efficient_attention(q, k, v) # [B, N, H, C] + elif ATTN == 'flash_attn': + out = flash_attn.flash_attn_qkvpacked_func(qkv_feats) # [B, N, H, C] + else: + raise ValueError(f"Unknown attention module: {ATTN}") + out = out.reshape(B * N, H, C) # [M, H, C] + else: + if ATTN == 'xformers': + q, k, v = qkv_feats.unbind(dim=1) # [M, H, C] + q = q.unsqueeze(0) # [1, M, H, C] + k = k.unsqueeze(0) # [1, M, H, C] + v = v.unsqueeze(0) # [1, M, H, C] + mask = xops.fmha.BlockDiagonalMask.from_seqlens(seq_lens) + out = xops.memory_efficient_attention(q, k, v, mask)[0] # [M, H, C] + elif ATTN == 'flash_attn': + cu_seqlens = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(seq_lens), dim=0)], dim=0) \ + .to(qkv.device).int() + out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, cu_seqlens, max(seq_lens)) # [M, H, C] + + out = out[bwd_indices] # [T, H, C] + + if DEBUG: + qkv_coords = qkv_coords[bwd_indices] + assert torch.equal(qkv_coords, qkv.coords), "SparseWindowedScaledDotProductSelfAttention: coordinate mismatch" + + return qkv.replace(out) diff --git a/third_party/TRELLIS/trellis/modules/sparse/attention/windowed_attn.py b/third_party/TRELLIS/trellis/modules/sparse/attention/windowed_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..cd642c5252e29a3a5e59fad7ed3880b7b00bcf9a --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/attention/windowed_attn.py @@ -0,0 +1,135 @@ +from typing import * +import torch +import math +from .. import SparseTensor +from .. import DEBUG, ATTN + +if ATTN == 'xformers': + import xformers.ops as xops +elif ATTN == 'flash_attn': + import flash_attn +else: + raise ValueError(f"Unknown attention module: {ATTN}") + + +__all__ = [ + 'sparse_windowed_scaled_dot_product_self_attention', +] + + +def calc_window_partition( + tensor: SparseTensor, + window_size: Union[int, Tuple[int, ...]], + shift_window: Union[int, Tuple[int, ...]] = 0 +) -> Tuple[torch.Tensor, torch.Tensor, List[int], List[int]]: + """ + Calculate serialization and partitioning for a set of coordinates. + + Args: + tensor (SparseTensor): The input tensor. + window_size (int): The window size to use. + shift_window (Tuple[int, ...]): The shift of serialized coordinates. + + Returns: + (torch.Tensor): Forwards indices. + (torch.Tensor): Backwards indices. + (List[int]): Sequence lengths. + (List[int]): Sequence batch indices. + """ + DIM = tensor.coords.shape[1] - 1 + shift_window = (shift_window,) * DIM if isinstance(shift_window, int) else shift_window + window_size = (window_size,) * DIM if isinstance(window_size, int) else window_size + shifted_coords = tensor.coords.clone().detach() + shifted_coords[:, 1:] += torch.tensor(shift_window, device=tensor.device, dtype=torch.int32).unsqueeze(0) + + MAX_COORDS = shifted_coords[:, 1:].max(dim=0).values.tolist() + NUM_WINDOWS = [math.ceil((mc + 1) / ws) for mc, ws in zip(MAX_COORDS, window_size)] + OFFSET = torch.cumprod(torch.tensor([1] + NUM_WINDOWS[::-1]), dim=0).tolist()[::-1] + + shifted_coords[:, 1:] //= torch.tensor(window_size, device=tensor.device, dtype=torch.int32).unsqueeze(0) + shifted_indices = (shifted_coords * torch.tensor(OFFSET, device=tensor.device, dtype=torch.int32).unsqueeze(0)).sum(dim=1) + fwd_indices = torch.argsort(shifted_indices) + bwd_indices = torch.empty_like(fwd_indices) + bwd_indices[fwd_indices] = torch.arange(fwd_indices.shape[0], device=tensor.device) + seq_lens = torch.bincount(shifted_indices) + seq_batch_indices = torch.arange(seq_lens.shape[0], device=tensor.device, dtype=torch.int32) // OFFSET[0] + mask = seq_lens != 0 + seq_lens = seq_lens[mask].tolist() + seq_batch_indices = seq_batch_indices[mask].tolist() + + return fwd_indices, bwd_indices, seq_lens, seq_batch_indices + + +def sparse_windowed_scaled_dot_product_self_attention( + qkv: SparseTensor, + window_size: int, + shift_window: Tuple[int, int, int] = (0, 0, 0) +) -> SparseTensor: + """ + Apply windowed scaled dot product self attention to a sparse tensor. + + Args: + qkv (SparseTensor): [N, *, 3, H, C] sparse tensor containing Qs, Ks, and Vs. + window_size (int): The window size to use. + shift_window (Tuple[int, int, int]): The shift of serialized coordinates. + shift (int): The shift to use. + """ + assert len(qkv.shape) == 4 and qkv.shape[1] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, *, 3, H, C]" + + serialization_spatial_cache_name = f'window_partition_{window_size}_{shift_window}' + serialization_spatial_cache = qkv.get_spatial_cache(serialization_spatial_cache_name) + if serialization_spatial_cache is None: + fwd_indices, bwd_indices, seq_lens, seq_batch_indices = calc_window_partition(qkv, window_size, shift_window) + qkv.register_spatial_cache(serialization_spatial_cache_name, (fwd_indices, bwd_indices, seq_lens, seq_batch_indices)) + else: + fwd_indices, bwd_indices, seq_lens, seq_batch_indices = serialization_spatial_cache + + M = fwd_indices.shape[0] + T = qkv.feats.shape[0] + H = qkv.feats.shape[2] + C = qkv.feats.shape[3] + + qkv_feats = qkv.feats[fwd_indices] # [M, 3, H, C] + + if DEBUG: + start = 0 + qkv_coords = qkv.coords[fwd_indices] + for i in range(len(seq_lens)): + seq_coords = qkv_coords[start:start+seq_lens[i]] + assert (seq_coords[:, 0] == seq_batch_indices[i]).all(), f"SparseWindowedScaledDotProductSelfAttention: batch index mismatch" + assert (seq_coords[:, 1:].max(dim=0).values - seq_coords[:, 1:].min(dim=0).values < window_size).all(), \ + f"SparseWindowedScaledDotProductSelfAttention: window size exceeded" + start += seq_lens[i] + + if all([seq_len == window_size for seq_len in seq_lens]): + B = len(seq_lens) + N = window_size + qkv_feats = qkv_feats.reshape(B, N, 3, H, C) + if ATTN == 'xformers': + q, k, v = qkv_feats.unbind(dim=2) # [B, N, H, C] + out = xops.memory_efficient_attention(q, k, v) # [B, N, H, C] + elif ATTN == 'flash_attn': + out = flash_attn.flash_attn_qkvpacked_func(qkv_feats) # [B, N, H, C] + else: + raise ValueError(f"Unknown attention module: {ATTN}") + out = out.reshape(B * N, H, C) # [M, H, C] + else: + if ATTN == 'xformers': + q, k, v = qkv_feats.unbind(dim=1) # [M, H, C] + q = q.unsqueeze(0) # [1, M, H, C] + k = k.unsqueeze(0) # [1, M, H, C] + v = v.unsqueeze(0) # [1, M, H, C] + mask = xops.fmha.BlockDiagonalMask.from_seqlens(seq_lens) + out = xops.memory_efficient_attention(q, k, v, mask)[0] # [M, H, C] + elif ATTN == 'flash_attn': + cu_seqlens = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(seq_lens), dim=0)], dim=0) \ + .to(qkv.device).int() + out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, cu_seqlens, max(seq_lens)) # [M, H, C] + + out = out[bwd_indices] # [T, H, C] + + if DEBUG: + qkv_coords = qkv_coords[bwd_indices] + assert torch.equal(qkv_coords, qkv.coords), "SparseWindowedScaledDotProductSelfAttention: coordinate mismatch" + + return qkv.replace(out) diff --git a/third_party/TRELLIS/trellis/modules/sparse/basic.py b/third_party/TRELLIS/trellis/modules/sparse/basic.py new file mode 100644 index 0000000000000000000000000000000000000000..8837f44052f6d573d09e3bfb897e659e10516bb5 --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/basic.py @@ -0,0 +1,459 @@ +from typing import * +import torch +import torch.nn as nn +from . import BACKEND, DEBUG +SparseTensorData = None # Lazy import + + +__all__ = [ + 'SparseTensor', + 'sparse_batch_broadcast', + 'sparse_batch_op', + 'sparse_cat', + 'sparse_unbind', +] + + +class SparseTensor: + """ + Sparse tensor with support for both torchsparse and spconv backends. + + Parameters: + - feats (torch.Tensor): Features of the sparse tensor. + - coords (torch.Tensor): Coordinates of the sparse tensor. + - shape (torch.Size): Shape of the sparse tensor. + - layout (List[slice]): Layout of the sparse tensor for each batch + - data (SparseTensorData): Sparse tensor data used for convolusion + + NOTE: + - Data corresponding to a same batch should be contiguous. + - Coords should be in [0, 1023] + """ + @overload + def __init__(self, feats: torch.Tensor, coords: torch.Tensor, shape: Optional[torch.Size] = None, layout: Optional[List[slice]] = None, **kwargs): ... + + @overload + def __init__(self, data, shape: Optional[torch.Size] = None, layout: Optional[List[slice]] = None, **kwargs): ... + + def __init__(self, *args, **kwargs): + # Lazy import of sparse tensor backend + global SparseTensorData + if SparseTensorData is None: + import importlib + if BACKEND == 'torchsparse': + SparseTensorData = importlib.import_module('torchsparse').SparseTensor + elif BACKEND == 'spconv': + SparseTensorData = importlib.import_module('spconv.pytorch').SparseConvTensor + + method_id = 0 + if len(args) != 0: + method_id = 0 if isinstance(args[0], torch.Tensor) else 1 + else: + method_id = 1 if 'data' in kwargs else 0 + + if method_id == 0: + feats, coords, shape, layout = args + (None,) * (4 - len(args)) + if 'feats' in kwargs: + feats = kwargs['feats'] + del kwargs['feats'] + if 'coords' in kwargs: + coords = kwargs['coords'] + del kwargs['coords'] + if 'shape' in kwargs: + shape = kwargs['shape'] + del kwargs['shape'] + if 'layout' in kwargs: + layout = kwargs['layout'] + del kwargs['layout'] + + if shape is None: + shape = self.__cal_shape(feats, coords) + if layout is None: + layout = self.__cal_layout(coords, shape[0]) + if BACKEND == 'torchsparse': + self.data = SparseTensorData(feats, coords, **kwargs) + elif BACKEND == 'spconv': + spatial_shape = list(coords.max(0)[0] + 1)[1:] + self.data = SparseTensorData(feats.reshape(feats.shape[0], -1), coords, spatial_shape, shape[0], **kwargs) + self.data._features = feats + elif method_id == 1: + data, shape, layout = args + (None,) * (3 - len(args)) + if 'data' in kwargs: + data = kwargs['data'] + del kwargs['data'] + if 'shape' in kwargs: + shape = kwargs['shape'] + del kwargs['shape'] + if 'layout' in kwargs: + layout = kwargs['layout'] + del kwargs['layout'] + + self.data = data + if shape is None: + shape = self.__cal_shape(self.feats, self.coords) + if layout is None: + layout = self.__cal_layout(self.coords, shape[0]) + + self._shape = shape + self._layout = layout + self._scale = kwargs.get('scale', (1, 1, 1)) + self._spatial_cache = kwargs.get('spatial_cache', {}) + + if DEBUG: + try: + assert self.feats.shape[0] == self.coords.shape[0], f"Invalid feats shape: {self.feats.shape}, coords shape: {self.coords.shape}" + assert self.shape == self.__cal_shape(self.feats, self.coords), f"Invalid shape: {self.shape}" + assert self.layout == self.__cal_layout(self.coords, self.shape[0]), f"Invalid layout: {self.layout}" + for i in range(self.shape[0]): + assert torch.all(self.coords[self.layout[i], 0] == i), f"The data of batch {i} is not contiguous" + except Exception as e: + print('Debugging information:') + print(f"- Shape: {self.shape}") + print(f"- Layout: {self.layout}") + print(f"- Scale: {self._scale}") + print(f"- Coords: {self.coords}") + raise e + + def __cal_shape(self, feats, coords): + shape = [] + shape.append(coords[:, 0].max().item() + 1) + shape.extend([*feats.shape[1:]]) + return torch.Size(shape) + + def __cal_layout(self, coords, batch_size): + seq_len = torch.bincount(coords[:, 0], minlength=batch_size) + offset = torch.cumsum(seq_len, dim=0) + layout = [slice((offset[i] - seq_len[i]).item(), offset[i].item()) for i in range(batch_size)] + return layout + + @property + def shape(self) -> torch.Size: + return self._shape + + def dim(self) -> int: + return len(self.shape) + + @property + def layout(self) -> List[slice]: + return self._layout + + @property + def feats(self) -> torch.Tensor: + if BACKEND == 'torchsparse': + return self.data.F + elif BACKEND == 'spconv': + return self.data.features + + @feats.setter + def feats(self, value: torch.Tensor): + if BACKEND == 'torchsparse': + self.data.F = value + elif BACKEND == 'spconv': + self.data.features = value + + @property + def coords(self) -> torch.Tensor: + if BACKEND == 'torchsparse': + return self.data.C + elif BACKEND == 'spconv': + return self.data.indices + + @coords.setter + def coords(self, value: torch.Tensor): + if BACKEND == 'torchsparse': + self.data.C = value + elif BACKEND == 'spconv': + self.data.indices = value + + @property + def dtype(self): + return self.feats.dtype + + @property + def device(self): + return self.feats.device + + @overload + def to(self, dtype: torch.dtype) -> 'SparseTensor': ... + + @overload + def to(self, device: Optional[Union[str, torch.device]] = None, dtype: Optional[torch.dtype] = None) -> 'SparseTensor': ... + + def to(self, *args, **kwargs) -> 'SparseTensor': + device = None + dtype = None + if len(args) == 2: + device, dtype = args + elif len(args) == 1: + if isinstance(args[0], torch.dtype): + dtype = args[0] + else: + device = args[0] + if 'dtype' in kwargs: + assert dtype is None, "to() received multiple values for argument 'dtype'" + dtype = kwargs['dtype'] + if 'device' in kwargs: + assert device is None, "to() received multiple values for argument 'device'" + device = kwargs['device'] + + new_feats = self.feats.to(device=device, dtype=dtype) + new_coords = self.coords.to(device=device) + return self.replace(new_feats, new_coords) + + def type(self, dtype): + new_feats = self.feats.type(dtype) + return self.replace(new_feats) + + def cpu(self) -> 'SparseTensor': + new_feats = self.feats.cpu() + new_coords = self.coords.cpu() + return self.replace(new_feats, new_coords) + + def cuda(self) -> 'SparseTensor': + new_feats = self.feats.cuda() + new_coords = self.coords.cuda() + return self.replace(new_feats, new_coords) + + def half(self) -> 'SparseTensor': + new_feats = self.feats.half() + return self.replace(new_feats) + + def float(self) -> 'SparseTensor': + new_feats = self.feats.float() + return self.replace(new_feats) + + def detach(self) -> 'SparseTensor': + new_coords = self.coords.detach() + new_feats = self.feats.detach() + return self.replace(new_feats, new_coords) + + def dense(self) -> torch.Tensor: + if BACKEND == 'torchsparse': + return self.data.dense() + elif BACKEND == 'spconv': + return self.data.dense() + + def reshape(self, *shape) -> 'SparseTensor': + new_feats = self.feats.reshape(self.feats.shape[0], *shape) + return self.replace(new_feats) + + def unbind(self, dim: int) -> List['SparseTensor']: + return sparse_unbind(self, dim) + + def replace(self, feats: torch.Tensor, coords: Optional[torch.Tensor] = None) -> 'SparseTensor': + new_shape = [self.shape[0]] + new_shape.extend(feats.shape[1:]) + if BACKEND == 'torchsparse': + new_data = SparseTensorData( + feats=feats, + coords=self.data.coords if coords is None else coords, + stride=self.data.stride, + spatial_range=self.data.spatial_range, + ) + new_data._caches = self.data._caches + elif BACKEND == 'spconv': + new_data = SparseTensorData( + self.data.features.reshape(self.data.features.shape[0], -1), + self.data.indices, + self.data.spatial_shape, + self.data.batch_size, + self.data.grid, + self.data.voxel_num, + self.data.indice_dict + ) + new_data._features = feats + new_data.benchmark = self.data.benchmark + new_data.benchmark_record = self.data.benchmark_record + new_data.thrust_allocator = self.data.thrust_allocator + new_data._timer = self.data._timer + new_data.force_algo = self.data.force_algo + new_data.int8_scale = self.data.int8_scale + if coords is not None: + new_data.indices = coords + new_tensor = SparseTensor(new_data, shape=torch.Size(new_shape), layout=self.layout, scale=self._scale, spatial_cache=self._spatial_cache) + return new_tensor + + @staticmethod + def full(aabb, dim, value, dtype=torch.float32, device=None) -> 'SparseTensor': + N, C = dim + x = torch.arange(aabb[0], aabb[3] + 1) + y = torch.arange(aabb[1], aabb[4] + 1) + z = torch.arange(aabb[2], aabb[5] + 1) + coords = torch.stack(torch.meshgrid(x, y, z, indexing='ij'), dim=-1).reshape(-1, 3) + coords = torch.cat([ + torch.arange(N).view(-1, 1).repeat(1, coords.shape[0]).view(-1, 1), + coords.repeat(N, 1), + ], dim=1).to(dtype=torch.int32, device=device) + feats = torch.full((coords.shape[0], C), value, dtype=dtype, device=device) + return SparseTensor(feats=feats, coords=coords) + + def __merge_sparse_cache(self, other: 'SparseTensor') -> dict: + new_cache = {} + for k in set(list(self._spatial_cache.keys()) + list(other._spatial_cache.keys())): + if k in self._spatial_cache: + new_cache[k] = self._spatial_cache[k] + if k in other._spatial_cache: + if k not in new_cache: + new_cache[k] = other._spatial_cache[k] + else: + new_cache[k].update(other._spatial_cache[k]) + return new_cache + + def __neg__(self) -> 'SparseTensor': + return self.replace(-self.feats) + + def __elemwise__(self, other: Union[torch.Tensor, 'SparseTensor'], op: callable) -> 'SparseTensor': + if isinstance(other, torch.Tensor): + try: + other = torch.broadcast_to(other, self.shape) + other = sparse_batch_broadcast(self, other) + except: + pass + if isinstance(other, SparseTensor): + other = other.feats + new_feats = op(self.feats, other) + new_tensor = self.replace(new_feats) + if isinstance(other, SparseTensor): + new_tensor._spatial_cache = self.__merge_sparse_cache(other) + return new_tensor + + def __add__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor': + return self.__elemwise__(other, torch.add) + + def __radd__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor': + return self.__elemwise__(other, torch.add) + + def __sub__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor': + return self.__elemwise__(other, torch.sub) + + def __rsub__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor': + return self.__elemwise__(other, lambda x, y: torch.sub(y, x)) + + def __mul__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor': + return self.__elemwise__(other, torch.mul) + + def __rmul__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor': + return self.__elemwise__(other, torch.mul) + + def __truediv__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor': + return self.__elemwise__(other, torch.div) + + def __rtruediv__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor': + return self.__elemwise__(other, lambda x, y: torch.div(y, x)) + + def __getitem__(self, idx): + if isinstance(idx, int): + idx = [idx] + elif isinstance(idx, slice): + idx = range(*idx.indices(self.shape[0])) + elif isinstance(idx, torch.Tensor): + if idx.dtype == torch.bool: + assert idx.shape == (self.shape[0],), f"Invalid index shape: {idx.shape}" + idx = idx.nonzero().squeeze(1) + elif idx.dtype in [torch.int32, torch.int64]: + assert len(idx.shape) == 1, f"Invalid index shape: {idx.shape}" + else: + raise ValueError(f"Unknown index type: {idx.dtype}") + else: + raise ValueError(f"Unknown index type: {type(idx)}") + + coords = [] + feats = [] + for new_idx, old_idx in enumerate(idx): + coords.append(self.coords[self.layout[old_idx]].clone()) + coords[-1][:, 0] = new_idx + feats.append(self.feats[self.layout[old_idx]]) + coords = torch.cat(coords, dim=0).contiguous() + feats = torch.cat(feats, dim=0).contiguous() + return SparseTensor(feats=feats, coords=coords) + + def register_spatial_cache(self, key, value) -> None: + """ + Register a spatial cache. + The spatial cache can be any thing you want to cache. + The registery and retrieval of the cache is based on current scale. + """ + scale_key = str(self._scale) + if scale_key not in self._spatial_cache: + self._spatial_cache[scale_key] = {} + self._spatial_cache[scale_key][key] = value + + def get_spatial_cache(self, key=None): + """ + Get a spatial cache. + """ + scale_key = str(self._scale) + cur_scale_cache = self._spatial_cache.get(scale_key, {}) + if key is None: + return cur_scale_cache + return cur_scale_cache.get(key, None) + + +def sparse_batch_broadcast(input: SparseTensor, other: torch.Tensor) -> torch.Tensor: + """ + Broadcast a 1D tensor to a sparse tensor along the batch dimension then perform an operation. + + Args: + input (torch.Tensor): 1D tensor to broadcast. + target (SparseTensor): Sparse tensor to broadcast to. + op (callable): Operation to perform after broadcasting. Defaults to torch.add. + """ + coords, feats = input.coords, input.feats + broadcasted = torch.zeros_like(feats) + for k in range(input.shape[0]): + broadcasted[input.layout[k]] = other[k] + return broadcasted + + +def sparse_batch_op(input: SparseTensor, other: torch.Tensor, op: callable = torch.add) -> SparseTensor: + """ + Broadcast a 1D tensor to a sparse tensor along the batch dimension then perform an operation. + + Args: + input (torch.Tensor): 1D tensor to broadcast. + target (SparseTensor): Sparse tensor to broadcast to. + op (callable): Operation to perform after broadcasting. Defaults to torch.add. + """ + return input.replace(op(input.feats, sparse_batch_broadcast(input, other))) + + +def sparse_cat(inputs: List[SparseTensor], dim: int = 0) -> SparseTensor: + """ + Concatenate a list of sparse tensors. + + Args: + inputs (List[SparseTensor]): List of sparse tensors to concatenate. + """ + if dim == 0: + start = 0 + coords = [] + for input in inputs: + coords.append(input.coords.clone()) + coords[-1][:, 0] += start + start += input.shape[0] + coords = torch.cat(coords, dim=0) + feats = torch.cat([input.feats for input in inputs], dim=0) + output = SparseTensor( + coords=coords, + feats=feats, + ) + else: + feats = torch.cat([input.feats for input in inputs], dim=dim) + output = inputs[0].replace(feats) + + return output + + +def sparse_unbind(input: SparseTensor, dim: int) -> List[SparseTensor]: + """ + Unbind a sparse tensor along a dimension. + + Args: + input (SparseTensor): Sparse tensor to unbind. + dim (int): Dimension to unbind. + """ + if dim == 0: + return [input[i] for i in range(input.shape[0])] + else: + feats = input.feats.unbind(dim) + return [input.replace(f) for f in feats] diff --git a/third_party/TRELLIS/trellis/modules/sparse/conv/__init__.py b/third_party/TRELLIS/trellis/modules/sparse/conv/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..340a87126a8de574ee0276feb96b49824a2ce234 --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/conv/__init__.py @@ -0,0 +1,21 @@ +from .. import BACKEND + + +SPCONV_ALGO = 'auto' # 'auto', 'implicit_gemm', 'native' + +def __from_env(): + import os + + global SPCONV_ALGO + env_spconv_algo = os.environ.get('SPCONV_ALGO') + if env_spconv_algo is not None and env_spconv_algo in ['auto', 'implicit_gemm', 'native']: + SPCONV_ALGO = env_spconv_algo + print(f"[SPARSE][CONV] spconv algo: {SPCONV_ALGO}") + + +__from_env() + +if BACKEND == 'torchsparse': + from .conv_torchsparse import * +elif BACKEND == 'spconv': + from .conv_spconv import * diff --git a/third_party/TRELLIS/trellis/modules/sparse/conv/conv_spconv.py b/third_party/TRELLIS/trellis/modules/sparse/conv/conv_spconv.py new file mode 100644 index 0000000000000000000000000000000000000000..524bcd4a845b2d6bd090a5f74bc8859978727528 --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/conv/conv_spconv.py @@ -0,0 +1,80 @@ +import torch +import torch.nn as nn +from .. import SparseTensor +from .. import DEBUG +from . import SPCONV_ALGO + +class SparseConv3d(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, padding=None, bias=True, indice_key=None): + super(SparseConv3d, self).__init__() + if 'spconv' not in globals(): + import spconv.pytorch as spconv + algo = None + if SPCONV_ALGO == 'native': + algo = spconv.ConvAlgo.Native + elif SPCONV_ALGO == 'implicit_gemm': + algo = spconv.ConvAlgo.MaskImplicitGemm + if stride == 1 and (padding is None): + self.conv = spconv.SubMConv3d(in_channels, out_channels, kernel_size, dilation=dilation, bias=bias, indice_key=indice_key, algo=algo) + else: + self.conv = spconv.SparseConv3d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, padding=padding, bias=bias, indice_key=indice_key, algo=algo) + self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride, stride, stride) + self.padding = padding + + def forward(self, x: SparseTensor) -> SparseTensor: + spatial_changed = any(s != 1 for s in self.stride) or (self.padding is not None) + new_data = self.conv(x.data) + new_shape = [x.shape[0], self.conv.out_channels] + new_layout = None if spatial_changed else x.layout + + if spatial_changed and (x.shape[0] != 1): + # spconv was non-1 stride will break the contiguous of the output tensor, sort by the coords + fwd = new_data.indices[:, 0].argsort() + bwd = torch.zeros_like(fwd).scatter_(0, fwd, torch.arange(fwd.shape[0], device=fwd.device)) + sorted_feats = new_data.features[fwd] + sorted_coords = new_data.indices[fwd] + unsorted_data = new_data + new_data = spconv.SparseConvTensor(sorted_feats, sorted_coords, unsorted_data.spatial_shape, unsorted_data.batch_size) # type: ignore + + out = SparseTensor( + new_data, shape=torch.Size(new_shape), layout=new_layout, + scale=tuple([s * stride for s, stride in zip(x._scale, self.stride)]), + spatial_cache=x._spatial_cache, + ) + + if spatial_changed and (x.shape[0] != 1): + out.register_spatial_cache(f'conv_{self.stride}_unsorted_data', unsorted_data) + out.register_spatial_cache(f'conv_{self.stride}_sort_bwd', bwd) + + return out + + +class SparseInverseConv3d(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True, indice_key=None): + super(SparseInverseConv3d, self).__init__() + if 'spconv' not in globals(): + import spconv.pytorch as spconv + self.conv = spconv.SparseInverseConv3d(in_channels, out_channels, kernel_size, bias=bias, indice_key=indice_key) + self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride, stride, stride) + + def forward(self, x: SparseTensor) -> SparseTensor: + spatial_changed = any(s != 1 for s in self.stride) + if spatial_changed: + # recover the original spconv order + data = x.get_spatial_cache(f'conv_{self.stride}_unsorted_data') + bwd = x.get_spatial_cache(f'conv_{self.stride}_sort_bwd') + data = data.replace_feature(x.feats[bwd]) + if DEBUG: + assert torch.equal(data.indices, x.coords[bwd]), 'Recover the original order failed' + else: + data = x.data + + new_data = self.conv(data) + new_shape = [x.shape[0], self.conv.out_channels] + new_layout = None if spatial_changed else x.layout + out = SparseTensor( + new_data, shape=torch.Size(new_shape), layout=new_layout, + scale=tuple([s // stride for s, stride in zip(x._scale, self.stride)]), + spatial_cache=x._spatial_cache, + ) + return out diff --git a/third_party/TRELLIS/trellis/modules/sparse/conv/conv_torchsparse.py b/third_party/TRELLIS/trellis/modules/sparse/conv/conv_torchsparse.py new file mode 100644 index 0000000000000000000000000000000000000000..1d612582d4b31f90aca3c00b693bbbc2550dc62c --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/conv/conv_torchsparse.py @@ -0,0 +1,38 @@ +import torch +import torch.nn as nn +from .. import SparseTensor + + +class SparseConv3d(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True, indice_key=None): + super(SparseConv3d, self).__init__() + if 'torchsparse' not in globals(): + import torchsparse + self.conv = torchsparse.nn.Conv3d(in_channels, out_channels, kernel_size, stride, 0, dilation, bias) + + def forward(self, x: SparseTensor) -> SparseTensor: + out = self.conv(x.data) + new_shape = [x.shape[0], self.conv.out_channels] + out = SparseTensor(out, shape=torch.Size(new_shape), layout=x.layout if all(s == 1 for s in self.conv.stride) else None) + out._spatial_cache = x._spatial_cache + out._scale = tuple([s * stride for s, stride in zip(x._scale, self.conv.stride)]) + return out + + +class SparseInverseConv3d(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True, indice_key=None): + super(SparseInverseConv3d, self).__init__() + if 'torchsparse' not in globals(): + import torchsparse + self.conv = torchsparse.nn.Conv3d(in_channels, out_channels, kernel_size, stride, 0, dilation, bias, transposed=True) + + def forward(self, x: SparseTensor) -> SparseTensor: + out = self.conv(x.data) + new_shape = [x.shape[0], self.conv.out_channels] + out = SparseTensor(out, shape=torch.Size(new_shape), layout=x.layout if all(s == 1 for s in self.conv.stride) else None) + out._spatial_cache = x._spatial_cache + out._scale = tuple([s // stride for s, stride in zip(x._scale, self.conv.stride)]) + return out + + + diff --git a/third_party/TRELLIS/trellis/modules/sparse/linear.py b/third_party/TRELLIS/trellis/modules/sparse/linear.py new file mode 100644 index 0000000000000000000000000000000000000000..a854e77ce87d1a190b9730d91f363a821ff250bd --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/linear.py @@ -0,0 +1,15 @@ +import torch +import torch.nn as nn +from . import SparseTensor + +__all__ = [ + 'SparseLinear' +] + + +class SparseLinear(nn.Linear): + def __init__(self, in_features, out_features, bias=True): + super(SparseLinear, self).__init__(in_features, out_features, bias) + + def forward(self, input: SparseTensor) -> SparseTensor: + return input.replace(super().forward(input.feats)) diff --git a/third_party/TRELLIS/trellis/modules/sparse/nonlinearity.py b/third_party/TRELLIS/trellis/modules/sparse/nonlinearity.py new file mode 100644 index 0000000000000000000000000000000000000000..f200098dd82011a3aeee1688b9eb17018fa78295 --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/nonlinearity.py @@ -0,0 +1,35 @@ +import torch +import torch.nn as nn +from . import SparseTensor + +__all__ = [ + 'SparseReLU', + 'SparseSiLU', + 'SparseGELU', + 'SparseActivation' +] + + +class SparseReLU(nn.ReLU): + def forward(self, input: SparseTensor) -> SparseTensor: + return input.replace(super().forward(input.feats)) + + +class SparseSiLU(nn.SiLU): + def forward(self, input: SparseTensor) -> SparseTensor: + return input.replace(super().forward(input.feats)) + + +class SparseGELU(nn.GELU): + def forward(self, input: SparseTensor) -> SparseTensor: + return input.replace(super().forward(input.feats)) + + +class SparseActivation(nn.Module): + def __init__(self, activation: nn.Module): + super().__init__() + self.activation = activation + + def forward(self, input: SparseTensor) -> SparseTensor: + return input.replace(self.activation(input.feats)) + diff --git a/third_party/TRELLIS/trellis/modules/sparse/norm.py b/third_party/TRELLIS/trellis/modules/sparse/norm.py new file mode 100644 index 0000000000000000000000000000000000000000..6b38a36682c098210000dc31d68ddc31ccd2929d --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/norm.py @@ -0,0 +1,58 @@ +import torch +import torch.nn as nn +from . import SparseTensor +from . import DEBUG + +__all__ = [ + 'SparseGroupNorm', + 'SparseLayerNorm', + 'SparseGroupNorm32', + 'SparseLayerNorm32', +] + + +class SparseGroupNorm(nn.GroupNorm): + def __init__(self, num_groups, num_channels, eps=1e-5, affine=True): + super(SparseGroupNorm, self).__init__(num_groups, num_channels, eps, affine) + + def forward(self, input: SparseTensor) -> SparseTensor: + nfeats = torch.zeros_like(input.feats) + for k in range(input.shape[0]): + if DEBUG: + assert (input.coords[input.layout[k], 0] == k).all(), f"SparseGroupNorm: batch index mismatch" + bfeats = input.feats[input.layout[k]] + bfeats = bfeats.permute(1, 0).reshape(1, input.shape[1], -1) + bfeats = super().forward(bfeats) + bfeats = bfeats.reshape(input.shape[1], -1).permute(1, 0) + nfeats[input.layout[k]] = bfeats + return input.replace(nfeats) + + +class SparseLayerNorm(nn.LayerNorm): + def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True): + super(SparseLayerNorm, self).__init__(normalized_shape, eps, elementwise_affine) + + def forward(self, input: SparseTensor) -> SparseTensor: + nfeats = torch.zeros_like(input.feats) + for k in range(input.shape[0]): + bfeats = input.feats[input.layout[k]] + bfeats = bfeats.permute(1, 0).reshape(1, input.shape[1], -1) + bfeats = super().forward(bfeats) + bfeats = bfeats.reshape(input.shape[1], -1).permute(1, 0) + nfeats[input.layout[k]] = bfeats + return input.replace(nfeats) + + +class SparseGroupNorm32(SparseGroupNorm): + """ + A GroupNorm layer that converts to float32 before the forward pass. + """ + def forward(self, x: SparseTensor) -> SparseTensor: + return super().forward(x.float()).type(x.dtype) + +class SparseLayerNorm32(SparseLayerNorm): + """ + A LayerNorm layer that converts to float32 before the forward pass. + """ + def forward(self, x: SparseTensor) -> SparseTensor: + return super().forward(x.float()).type(x.dtype) diff --git a/third_party/TRELLIS/trellis/modules/sparse/spatial.py b/third_party/TRELLIS/trellis/modules/sparse/spatial.py new file mode 100644 index 0000000000000000000000000000000000000000..ad7121473f335b307e2f7ea5f05c964d3aec0440 --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/spatial.py @@ -0,0 +1,110 @@ +from typing import * +import torch +import torch.nn as nn +from . import SparseTensor + +__all__ = [ + 'SparseDownsample', + 'SparseUpsample', + 'SparseSubdivide' +] + + +class SparseDownsample(nn.Module): + """ + Downsample a sparse tensor by a factor of `factor`. + Implemented as average pooling. + """ + def __init__(self, factor: Union[int, Tuple[int, ...], List[int]]): + super(SparseDownsample, self).__init__() + self.factor = tuple(factor) if isinstance(factor, (list, tuple)) else factor + + def forward(self, input: SparseTensor) -> SparseTensor: + DIM = input.coords.shape[-1] - 1 + factor = self.factor if isinstance(self.factor, tuple) else (self.factor,) * DIM + assert DIM == len(factor), 'Input coordinates must have the same dimension as the downsample factor.' + + coord = list(input.coords.unbind(dim=-1)) + for i, f in enumerate(factor): + coord[i+1] = coord[i+1] // f + + MAX = [coord[i+1].max().item() + 1 for i in range(DIM)] + OFFSET = torch.cumprod(torch.tensor(MAX[::-1]), 0).tolist()[::-1] + [1] + code = sum([c * o for c, o in zip(coord, OFFSET)]) + code, idx = code.unique(return_inverse=True) + + new_feats = torch.scatter_reduce( + torch.zeros(code.shape[0], input.feats.shape[1], device=input.feats.device, dtype=input.feats.dtype), + dim=0, + index=idx.unsqueeze(1).expand(-1, input.feats.shape[1]), + src=input.feats, + reduce='mean' + ) + new_coords = torch.stack( + [code // OFFSET[0]] + + [(code // OFFSET[i+1]) % MAX[i] for i in range(DIM)], + dim=-1 + ) + out = SparseTensor(new_feats, new_coords, input.shape,) + out._scale = tuple([s // f for s, f in zip(input._scale, factor)]) + out._spatial_cache = input._spatial_cache + + out.register_spatial_cache(f'upsample_{factor}_coords', input.coords) + out.register_spatial_cache(f'upsample_{factor}_layout', input.layout) + out.register_spatial_cache(f'upsample_{factor}_idx', idx) + + return out + + +class SparseUpsample(nn.Module): + """ + Upsample a sparse tensor by a factor of `factor`. + Implemented as nearest neighbor interpolation. + """ + def __init__(self, factor: Union[int, Tuple[int, int, int], List[int]]): + super(SparseUpsample, self).__init__() + self.factor = tuple(factor) if isinstance(factor, (list, tuple)) else factor + + def forward(self, input: SparseTensor) -> SparseTensor: + DIM = input.coords.shape[-1] - 1 + factor = self.factor if isinstance(self.factor, tuple) else (self.factor,) * DIM + assert DIM == len(factor), 'Input coordinates must have the same dimension as the upsample factor.' + + new_coords = input.get_spatial_cache(f'upsample_{factor}_coords') + new_layout = input.get_spatial_cache(f'upsample_{factor}_layout') + idx = input.get_spatial_cache(f'upsample_{factor}_idx') + if any([x is None for x in [new_coords, new_layout, idx]]): + raise ValueError('Upsample cache not found. SparseUpsample must be paired with SparseDownsample.') + new_feats = input.feats[idx] + out = SparseTensor(new_feats, new_coords, input.shape, new_layout) + out._scale = tuple([s * f for s, f in zip(input._scale, factor)]) + out._spatial_cache = input._spatial_cache + return out + +class SparseSubdivide(nn.Module): + """ + Upsample a sparse tensor by a factor of `factor`. + Implemented as nearest neighbor interpolation. + """ + def __init__(self): + super(SparseSubdivide, self).__init__() + + def forward(self, input: SparseTensor) -> SparseTensor: + DIM = input.coords.shape[-1] - 1 + # upsample scale=2^DIM + n_cube = torch.ones([2] * DIM, device=input.device, dtype=torch.int) + n_coords = torch.nonzero(n_cube) + n_coords = torch.cat([torch.zeros_like(n_coords[:, :1]), n_coords], dim=-1) + factor = n_coords.shape[0] + assert factor == 2 ** DIM + # print(n_coords.shape) + new_coords = input.coords.clone() + new_coords[:, 1:] *= 2 + new_coords = new_coords.unsqueeze(1) + n_coords.unsqueeze(0).to(new_coords.dtype) + + new_feats = input.feats.unsqueeze(1).expand(input.feats.shape[0], factor, *input.feats.shape[1:]) + out = SparseTensor(new_feats.flatten(0, 1), new_coords.flatten(0, 1), input.shape) + out._scale = input._scale * 2 + out._spatial_cache = input._spatial_cache + return out + diff --git a/third_party/TRELLIS/trellis/modules/sparse/transformer/__init__.py b/third_party/TRELLIS/trellis/modules/sparse/transformer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b08b0d4e5bc24060a2cdc8df75d06dce122972bd --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/transformer/__init__.py @@ -0,0 +1,2 @@ +from .blocks import * +from .modulated import * \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/modules/sparse/transformer/blocks.py b/third_party/TRELLIS/trellis/modules/sparse/transformer/blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..9d037a49bf83e1c2dfb2f8c4b23d2e9d6c51e9f0 --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/transformer/blocks.py @@ -0,0 +1,151 @@ +from typing import * +import torch +import torch.nn as nn +from ..basic import SparseTensor +from ..linear import SparseLinear +from ..nonlinearity import SparseGELU +from ..attention import SparseMultiHeadAttention, SerializeMode +from ...norm import LayerNorm32 + + +class SparseFeedForwardNet(nn.Module): + def __init__(self, channels: int, mlp_ratio: float = 4.0): + super().__init__() + self.mlp = nn.Sequential( + SparseLinear(channels, int(channels * mlp_ratio)), + SparseGELU(approximate="tanh"), + SparseLinear(int(channels * mlp_ratio), channels), + ) + + def forward(self, x: SparseTensor) -> SparseTensor: + return self.mlp(x) + + +class SparseTransformerBlock(nn.Module): + """ + Sparse Transformer block (MSA + FFN). + """ + def __init__( + self, + channels: int, + num_heads: int, + mlp_ratio: float = 4.0, + attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "full", + window_size: Optional[int] = None, + shift_sequence: Optional[int] = None, + shift_window: Optional[Tuple[int, int, int]] = None, + serialize_mode: Optional[SerializeMode] = None, + use_checkpoint: bool = False, + use_rope: bool = False, + qk_rms_norm: bool = False, + qkv_bias: bool = True, + ln_affine: bool = False, + ): + super().__init__() + self.use_checkpoint = use_checkpoint + self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) + self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) + self.attn = SparseMultiHeadAttention( + channels, + num_heads=num_heads, + attn_mode=attn_mode, + window_size=window_size, + shift_sequence=shift_sequence, + shift_window=shift_window, + serialize_mode=serialize_mode, + qkv_bias=qkv_bias, + use_rope=use_rope, + qk_rms_norm=qk_rms_norm, + ) + self.mlp = SparseFeedForwardNet( + channels, + mlp_ratio=mlp_ratio, + ) + + def _forward(self, x: SparseTensor) -> SparseTensor: + h = x.replace(self.norm1(x.feats)) + h = self.attn(h) + x = x + h + h = x.replace(self.norm2(x.feats)) + h = self.mlp(h) + x = x + h + return x + + def forward(self, x: SparseTensor) -> SparseTensor: + if self.use_checkpoint: + return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False) + else: + return self._forward(x) + + +class SparseTransformerCrossBlock(nn.Module): + """ + Sparse Transformer cross-attention block (MSA + MCA + FFN). + """ + def __init__( + self, + channels: int, + ctx_channels: int, + num_heads: int, + mlp_ratio: float = 4.0, + attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "full", + window_size: Optional[int] = None, + shift_sequence: Optional[int] = None, + shift_window: Optional[Tuple[int, int, int]] = None, + serialize_mode: Optional[SerializeMode] = None, + use_checkpoint: bool = False, + use_rope: bool = False, + qk_rms_norm: bool = False, + qk_rms_norm_cross: bool = False, + qkv_bias: bool = True, + ln_affine: bool = False, + ): + super().__init__() + self.use_checkpoint = use_checkpoint + self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) + self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) + self.norm3 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) + self.self_attn = SparseMultiHeadAttention( + channels, + num_heads=num_heads, + type="self", + attn_mode=attn_mode, + window_size=window_size, + shift_sequence=shift_sequence, + shift_window=shift_window, + serialize_mode=serialize_mode, + qkv_bias=qkv_bias, + use_rope=use_rope, + qk_rms_norm=qk_rms_norm, + ) + self.cross_attn = SparseMultiHeadAttention( + channels, + ctx_channels=ctx_channels, + num_heads=num_heads, + type="cross", + attn_mode="full", + qkv_bias=qkv_bias, + qk_rms_norm=qk_rms_norm_cross, + ) + self.mlp = SparseFeedForwardNet( + channels, + mlp_ratio=mlp_ratio, + ) + + def _forward(self, x: SparseTensor, mod: torch.Tensor, context: torch.Tensor): + h = x.replace(self.norm1(x.feats)) + h = self.self_attn(h) + x = x + h + h = x.replace(self.norm2(x.feats)) + h = self.cross_attn(h, context) + x = x + h + h = x.replace(self.norm3(x.feats)) + h = self.mlp(h) + x = x + h + return x + + def forward(self, x: SparseTensor, context: torch.Tensor): + if self.use_checkpoint: + return torch.utils.checkpoint.checkpoint(self._forward, x, context, use_reentrant=False) + else: + return self._forward(x, context) diff --git a/third_party/TRELLIS/trellis/modules/sparse/transformer/modulated.py b/third_party/TRELLIS/trellis/modules/sparse/transformer/modulated.py new file mode 100644 index 0000000000000000000000000000000000000000..39cea5acc7713530d6230923e6af84b94fb75df2 --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/sparse/transformer/modulated.py @@ -0,0 +1,166 @@ +from typing import * +import torch +import torch.nn as nn +from ..basic import SparseTensor +from ..attention import SparseMultiHeadAttention, SerializeMode +from ...norm import LayerNorm32 +from .blocks import SparseFeedForwardNet + + +class ModulatedSparseTransformerBlock(nn.Module): + """ + Sparse Transformer block (MSA + FFN) with adaptive layer norm conditioning. + """ + def __init__( + self, + channels: int, + num_heads: int, + mlp_ratio: float = 4.0, + attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "full", + window_size: Optional[int] = None, + shift_sequence: Optional[int] = None, + shift_window: Optional[Tuple[int, int, int]] = None, + serialize_mode: Optional[SerializeMode] = None, + use_checkpoint: bool = False, + use_rope: bool = False, + qk_rms_norm: bool = False, + qkv_bias: bool = True, + share_mod: bool = False, + ): + super().__init__() + self.use_checkpoint = use_checkpoint + self.share_mod = share_mod + self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6) + self.norm2 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6) + self.attn = SparseMultiHeadAttention( + channels, + num_heads=num_heads, + attn_mode=attn_mode, + window_size=window_size, + shift_sequence=shift_sequence, + shift_window=shift_window, + serialize_mode=serialize_mode, + qkv_bias=qkv_bias, + use_rope=use_rope, + qk_rms_norm=qk_rms_norm, + ) + self.mlp = SparseFeedForwardNet( + channels, + mlp_ratio=mlp_ratio, + ) + if not share_mod: + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(channels, 6 * channels, bias=True) + ) + + def _forward(self, x: SparseTensor, mod: torch.Tensor) -> SparseTensor: + if self.share_mod: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1) + else: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1) + h = x.replace(self.norm1(x.feats)) + h = h * (1 + scale_msa) + shift_msa + h = self.attn(h) + h = h * gate_msa + x = x + h + h = x.replace(self.norm2(x.feats)) + h = h * (1 + scale_mlp) + shift_mlp + h = self.mlp(h) + h = h * gate_mlp + x = x + h + return x + + def forward(self, x: SparseTensor, mod: torch.Tensor) -> SparseTensor: + if self.use_checkpoint: + return torch.utils.checkpoint.checkpoint(self._forward, x, mod, use_reentrant=False) + else: + return self._forward(x, mod) + + +class ModulatedSparseTransformerCrossBlock(nn.Module): + """ + Sparse Transformer cross-attention block (MSA + MCA + FFN) with adaptive layer norm conditioning. + """ + def __init__( + self, + channels: int, + ctx_channels: int, + num_heads: int, + mlp_ratio: float = 4.0, + attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "full", + window_size: Optional[int] = None, + shift_sequence: Optional[int] = None, + shift_window: Optional[Tuple[int, int, int]] = None, + serialize_mode: Optional[SerializeMode] = None, + use_checkpoint: bool = False, + use_rope: bool = False, + qk_rms_norm: bool = False, + qk_rms_norm_cross: bool = False, + qkv_bias: bool = True, + share_mod: bool = False, + + ): + super().__init__() + self.use_checkpoint = use_checkpoint + self.share_mod = share_mod + self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6) + self.norm2 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6) + self.norm3 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6) + self.self_attn = SparseMultiHeadAttention( + channels, + num_heads=num_heads, + type="self", + attn_mode=attn_mode, + window_size=window_size, + shift_sequence=shift_sequence, + shift_window=shift_window, + serialize_mode=serialize_mode, + qkv_bias=qkv_bias, + use_rope=use_rope, + qk_rms_norm=qk_rms_norm, + ) + self.cross_attn = SparseMultiHeadAttention( + channels, + ctx_channels=ctx_channels, + num_heads=num_heads, + type="cross", + attn_mode="full", + qkv_bias=qkv_bias, + qk_rms_norm=qk_rms_norm_cross, + ) + self.mlp = SparseFeedForwardNet( + channels, + mlp_ratio=mlp_ratio, + ) + if not share_mod: + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(channels, 6 * channels, bias=True) + ) + + def _forward(self, x: SparseTensor, mod: torch.Tensor, context: torch.Tensor, context_list=None, coords_dense_indices=None, self_attn_region=None, self_attn_region_boost=None) -> SparseTensor: + if self.share_mod: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1) + else: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1) + h = x.replace(self.norm1(x.feats)) + h = h * (1 + scale_msa) + shift_msa + h = self.self_attn(h, self_attn_region=self_attn_region, self_attn_region_boost=self_attn_region_boost) + h = h * gate_msa + x = x + h + h = x.replace(self.norm2(x.feats)) + h = self.cross_attn(h, context, context_list, coords_dense_indices) + x = x + h + h = x.replace(self.norm3(x.feats)) + h = h * (1 + scale_mlp) + shift_mlp + h = self.mlp(h) + h = h * gate_mlp + x = x + h + return x + + def forward(self, x: SparseTensor, mod: torch.Tensor, context: torch.Tensor, context_list=None, coords_dense_indices=None, self_attn_region=None, self_attn_region_boost=None) -> SparseTensor: + if self.use_checkpoint: + return torch.utils.checkpoint.checkpoint(self._forward, x, mod, context, context_list, coords_dense_indices, self_attn_region, self_attn_region_boost, use_reentrant=False) + else: + return self._forward(x, mod, context, context_list, coords_dense_indices, self_attn_region, self_attn_region_boost) diff --git a/third_party/TRELLIS/trellis/modules/spatial.py b/third_party/TRELLIS/trellis/modules/spatial.py new file mode 100644 index 0000000000000000000000000000000000000000..79e268d36c2ba49b0275744022a1a1e19983dae3 --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/spatial.py @@ -0,0 +1,48 @@ +import torch + + +def pixel_shuffle_3d(x: torch.Tensor, scale_factor: int) -> torch.Tensor: + """ + 3D pixel shuffle. + """ + B, C, H, W, D = x.shape + C_ = C // scale_factor**3 + x = x.reshape(B, C_, scale_factor, scale_factor, scale_factor, H, W, D) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4) + x = x.reshape(B, C_, H*scale_factor, W*scale_factor, D*scale_factor) + return x + + +def patchify(x: torch.Tensor, patch_size: int): + """ + Patchify a tensor. + + Args: + x (torch.Tensor): (N, C, *spatial) tensor + patch_size (int): Patch size + """ + DIM = x.dim() - 2 + for d in range(2, DIM + 2): + assert x.shape[d] % patch_size == 0, f"Dimension {d} of input tensor must be divisible by patch size, got {x.shape[d]} and {patch_size}" + + x = x.reshape(*x.shape[:2], *sum([[x.shape[d] // patch_size, patch_size] for d in range(2, DIM + 2)], [])) + x = x.permute(0, 1, *([2 * i + 3 for i in range(DIM)] + [2 * i + 2 for i in range(DIM)])) + x = x.reshape(x.shape[0], x.shape[1] * (patch_size ** DIM), *(x.shape[-DIM:])) + return x + + +def unpatchify(x: torch.Tensor, patch_size: int): + """ + Unpatchify a tensor. + + Args: + x (torch.Tensor): (N, C, *spatial) tensor + patch_size (int): Patch size + """ + DIM = x.dim() - 2 + assert x.shape[1] % (patch_size ** DIM) == 0, f"Second dimension of input tensor must be divisible by patch size to unpatchify, got {x.shape[1]} and {patch_size ** DIM}" + + x = x.reshape(x.shape[0], x.shape[1] // (patch_size ** DIM), *([patch_size] * DIM), *(x.shape[-DIM:])) + x = x.permute(0, 1, *(sum([[2 + DIM + i, 2 + i] for i in range(DIM)], []))) + x = x.reshape(x.shape[0], x.shape[1], *[x.shape[2 + 2 * i] * patch_size for i in range(DIM)]) + return x diff --git a/third_party/TRELLIS/trellis/modules/transformer/__init__.py b/third_party/TRELLIS/trellis/modules/transformer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b08b0d4e5bc24060a2cdc8df75d06dce122972bd --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/transformer/__init__.py @@ -0,0 +1,2 @@ +from .blocks import * +from .modulated import * \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/modules/transformer/blocks.py b/third_party/TRELLIS/trellis/modules/transformer/blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..c37eb7ed92f4aacfc9e974a63b247589d95977da --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/transformer/blocks.py @@ -0,0 +1,182 @@ +from typing import * +import torch +import torch.nn as nn +from ..attention import MultiHeadAttention +from ..norm import LayerNorm32 + + +class AbsolutePositionEmbedder(nn.Module): + """ + Embeds spatial positions into vector representations. + """ + def __init__(self, channels: int, in_channels: int = 3): + super().__init__() + self.channels = channels + self.in_channels = in_channels + self.freq_dim = channels // in_channels // 2 + self.freqs = torch.arange(self.freq_dim, dtype=torch.float32) / self.freq_dim + self.freqs = 1.0 / (10000 ** self.freqs) + + def _sin_cos_embedding(self, x: torch.Tensor) -> torch.Tensor: + """ + Create sinusoidal position embeddings. + + Args: + x: a 1-D Tensor of N indices + + Returns: + an (N, D) Tensor of positional embeddings. + """ + self.freqs = self.freqs.to(x.device) + out = torch.outer(x, self.freqs) + out = torch.cat([torch.sin(out), torch.cos(out)], dim=-1) + return out + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x (torch.Tensor): (N, D) tensor of spatial positions + """ + N, D = x.shape + assert D == self.in_channels, "Input dimension must match number of input channels" + embed = self._sin_cos_embedding(x.reshape(-1)) + embed = embed.reshape(N, -1) + if embed.shape[1] < self.channels: + embed = torch.cat([embed, torch.zeros(N, self.channels - embed.shape[1], device=embed.device)], dim=-1) + return embed + + +class FeedForwardNet(nn.Module): + def __init__(self, channels: int, mlp_ratio: float = 4.0): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(channels, int(channels * mlp_ratio)), + nn.GELU(approximate="tanh"), + nn.Linear(int(channels * mlp_ratio), channels), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.mlp(x) + + +class TransformerBlock(nn.Module): + """ + Transformer block (MSA + FFN). + """ + def __init__( + self, + channels: int, + num_heads: int, + mlp_ratio: float = 4.0, + attn_mode: Literal["full", "windowed"] = "full", + window_size: Optional[int] = None, + shift_window: Optional[int] = None, + use_checkpoint: bool = False, + use_rope: bool = False, + qk_rms_norm: bool = False, + qkv_bias: bool = True, + ln_affine: bool = False, + ): + super().__init__() + self.use_checkpoint = use_checkpoint + self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) + self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) + self.attn = MultiHeadAttention( + channels, + num_heads=num_heads, + attn_mode=attn_mode, + window_size=window_size, + shift_window=shift_window, + qkv_bias=qkv_bias, + use_rope=use_rope, + qk_rms_norm=qk_rms_norm, + ) + self.mlp = FeedForwardNet( + channels, + mlp_ratio=mlp_ratio, + ) + + def _forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.norm1(x) + h = self.attn(h) + x = x + h + h = self.norm2(x) + h = self.mlp(h) + x = x + h + return x + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.use_checkpoint: + return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False) + else: + return self._forward(x) + + +class TransformerCrossBlock(nn.Module): + """ + Transformer cross-attention block (MSA + MCA + FFN). + """ + def __init__( + self, + channels: int, + ctx_channels: int, + num_heads: int, + mlp_ratio: float = 4.0, + attn_mode: Literal["full", "windowed"] = "full", + window_size: Optional[int] = None, + shift_window: Optional[Tuple[int, int, int]] = None, + use_checkpoint: bool = False, + use_rope: bool = False, + qk_rms_norm: bool = False, + qk_rms_norm_cross: bool = False, + qkv_bias: bool = True, + ln_affine: bool = False, + ): + super().__init__() + self.use_checkpoint = use_checkpoint + self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) + self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) + self.norm3 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) + self.self_attn = MultiHeadAttention( + channels, + num_heads=num_heads, + type="self", + attn_mode=attn_mode, + window_size=window_size, + shift_window=shift_window, + qkv_bias=qkv_bias, + use_rope=use_rope, + qk_rms_norm=qk_rms_norm, + ) + self.cross_attn = MultiHeadAttention( + channels, + ctx_channels=ctx_channels, + num_heads=num_heads, + type="cross", + attn_mode="full", + qkv_bias=qkv_bias, + qk_rms_norm=qk_rms_norm_cross, + ) + self.mlp = FeedForwardNet( + channels, + mlp_ratio=mlp_ratio, + ) + + def _forward(self, x: torch.Tensor, context: torch.Tensor): + h = self.norm1(x) + h = self.self_attn(h) + x = x + h + h = self.norm2(x) + h = self.cross_attn(h, context) + x = x + h + h = self.norm3(x) + h = self.mlp(h) + x = x + h + return x + + def forward(self, x: torch.Tensor, context: torch.Tensor): + if self.use_checkpoint: + return torch.utils.checkpoint.checkpoint(self._forward, x, context, use_reentrant=False) + else: + return self._forward(x, context) + \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/modules/transformer/modulated.py b/third_party/TRELLIS/trellis/modules/transformer/modulated.py new file mode 100644 index 0000000000000000000000000000000000000000..d4aeca0689e68f656b08f7aa822b7be839aa727d --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/transformer/modulated.py @@ -0,0 +1,157 @@ +from typing import * +import torch +import torch.nn as nn +from ..attention import MultiHeadAttention +from ..norm import LayerNorm32 +from .blocks import FeedForwardNet + + +class ModulatedTransformerBlock(nn.Module): + """ + Transformer block (MSA + FFN) with adaptive layer norm conditioning. + """ + def __init__( + self, + channels: int, + num_heads: int, + mlp_ratio: float = 4.0, + attn_mode: Literal["full", "windowed"] = "full", + window_size: Optional[int] = None, + shift_window: Optional[Tuple[int, int, int]] = None, + use_checkpoint: bool = False, + use_rope: bool = False, + qk_rms_norm: bool = False, + qkv_bias: bool = True, + share_mod: bool = False, + ): + super().__init__() + self.use_checkpoint = use_checkpoint + self.share_mod = share_mod + self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6) + self.norm2 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6) + self.attn = MultiHeadAttention( + channels, + num_heads=num_heads, + attn_mode=attn_mode, + window_size=window_size, + shift_window=shift_window, + qkv_bias=qkv_bias, + use_rope=use_rope, + qk_rms_norm=qk_rms_norm, + ) + self.mlp = FeedForwardNet( + channels, + mlp_ratio=mlp_ratio, + ) + if not share_mod: + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(channels, 6 * channels, bias=True) + ) + + def _forward(self, x: torch.Tensor, mod: torch.Tensor) -> torch.Tensor: + if self.share_mod: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1) + else: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1) + h = self.norm1(x) + h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1) + h = self.attn(h) + h = h * gate_msa.unsqueeze(1) + x = x + h + h = self.norm2(x) + h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1) + h = self.mlp(h) + h = h * gate_mlp.unsqueeze(1) + x = x + h + return x + + def forward(self, x: torch.Tensor, mod: torch.Tensor) -> torch.Tensor: + if self.use_checkpoint: + return torch.utils.checkpoint.checkpoint(self._forward, x, mod, use_reentrant=False) + else: + return self._forward(x, mod) + + +class ModulatedTransformerCrossBlock(nn.Module): + """ + Transformer cross-attention block (MSA + MCA + FFN) with adaptive layer norm conditioning. + """ + def __init__( + self, + channels: int, + ctx_channels: int, + num_heads: int, + mlp_ratio: float = 4.0, + attn_mode: Literal["full", "windowed"] = "full", + window_size: Optional[int] = None, + shift_window: Optional[Tuple[int, int, int]] = None, + use_checkpoint: bool = False, + use_rope: bool = False, + qk_rms_norm: bool = False, + qk_rms_norm_cross: bool = False, + qkv_bias: bool = True, + share_mod: bool = False, + ): + super().__init__() + self.use_checkpoint = use_checkpoint + self.share_mod = share_mod + self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6) + self.norm2 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6) + self.norm3 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6) + self.self_attn = MultiHeadAttention( + channels, + num_heads=num_heads, + type="self", + attn_mode=attn_mode, + window_size=window_size, + shift_window=shift_window, + qkv_bias=qkv_bias, + use_rope=use_rope, + qk_rms_norm=qk_rms_norm, + ) + self.cross_attn = MultiHeadAttention( + channels, + ctx_channels=ctx_channels, + num_heads=num_heads, + type="cross", + attn_mode="full", + qkv_bias=qkv_bias, + qk_rms_norm=qk_rms_norm_cross, + ) + self.mlp = FeedForwardNet( + channels, + mlp_ratio=mlp_ratio, + ) + if not share_mod: + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(channels, 6 * channels, bias=True) + ) + + def _forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor): + if self.share_mod: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1) + else: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1) + h = self.norm1(x) + h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1) + h = self.self_attn(h) + h = h * gate_msa.unsqueeze(1) + x = x + h + h = self.norm2(x) + h = self.cross_attn(h, context) + x = x + h + h = self.norm3(x) + h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1) + h = self.mlp(h) + h = h * gate_mlp.unsqueeze(1) + x = x + h + return x + + def forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor): + if self.use_checkpoint: + return torch.utils.checkpoint.checkpoint(self._forward, x, mod, context, use_reentrant=False) + else: + return self._forward(x, mod, context) + \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/modules/utils.py b/third_party/TRELLIS/trellis/modules/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f0afb1b6c767aa2ad00bad96649fb30315e696ea --- /dev/null +++ b/third_party/TRELLIS/trellis/modules/utils.py @@ -0,0 +1,54 @@ +import torch.nn as nn +from ..modules import sparse as sp + +FP16_MODULES = ( + nn.Conv1d, + nn.Conv2d, + nn.Conv3d, + nn.ConvTranspose1d, + nn.ConvTranspose2d, + nn.ConvTranspose3d, + nn.Linear, + sp.SparseConv3d, + sp.SparseInverseConv3d, + sp.SparseLinear, +) + +def convert_module_to_f16(l): + """ + Convert primitive modules to float16. + """ + if isinstance(l, FP16_MODULES): + for p in l.parameters(): + p.data = p.data.half() + + +def convert_module_to_f32(l): + """ + Convert primitive modules to float32, undoing convert_module_to_f16(). + """ + if isinstance(l, FP16_MODULES): + for p in l.parameters(): + p.data = p.data.float() + + +def zero_module(module): + """ + Zero out the parameters of a module and return it. + """ + for p in module.parameters(): + p.detach().zero_() + return module + + +def scale_module(module, scale): + """ + Scale the parameters of a module and return it. + """ + for p in module.parameters(): + p.detach().mul_(scale) + return module + + +def modulate(x, shift, scale): + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) diff --git a/third_party/TRELLIS/trellis/pipelines/__init__.py b/third_party/TRELLIS/trellis/pipelines/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..92b063773c5da1b1b34b68958fd650844b5ad55e --- /dev/null +++ b/third_party/TRELLIS/trellis/pipelines/__init__.py @@ -0,0 +1,47 @@ +from . import samplers +from .trellis_image_to_3d import TrellisImageTo3DPipeline +from .trellis_text_to_3d import TrellisTextTo3DPipeline + + +def from_pretrained(path: str): + """ + Load a pipeline from a model folder or a Hugging Face model hub. + + Args: + path: The path to the model. Can be either local path or a Hugging Face model name. + """ + import os + import json + is_local = os.path.exists(f"{path}/pipeline.json") + + if is_local: + config_file = f"{path}/pipeline.json" + else: + from huggingface_hub import hf_hub_download + config_file = hf_hub_download(path, "pipeline.json") + + with open(config_file, 'r') as f: + config = json.load(f) + return globals()[config['name']].from_pretrained(path) + + +def from_pretrained(path: str): + """ + Load a pipeline from a model folder or a Hugging Face model hub. + + Args: + path: The path to the model. Can be either local path or a Hugging Face model name. + """ + import os + import json + is_local = os.path.exists(f"{path}/pipeline.json") + + if is_local: + config_file = f"{path}/pipeline.json" + else: + from huggingface_hub import hf_hub_download + config_file = hf_hub_download(path, "pipeline.json") + + with open(config_file, 'r') as f: + config = json.load(f) + return globals()[config['name']].from_pretrained(path) diff --git a/third_party/TRELLIS/trellis/pipelines/base.py b/third_party/TRELLIS/trellis/pipelines/base.py new file mode 100644 index 0000000000000000000000000000000000000000..59edafb98b9649b3bc863ba1879e34de2ab7a857 --- /dev/null +++ b/third_party/TRELLIS/trellis/pipelines/base.py @@ -0,0 +1,70 @@ +from typing import * +import torch +import torch.nn as nn +from .. import models + + +class Pipeline: + """ + A base class for pipelines. + """ + def __init__( + self, + models: dict[str, nn.Module] = None, + ): + if models is None: + return + self.models = models + for model in self.models.values(): + model.eval() + + @staticmethod + def from_pretrained(path: str) -> "Pipeline": + """ + Load a pretrained model. + """ + import os + import json + is_local = os.path.exists(f"{path}/pipeline.json") + + if is_local: + config_file = f"{path}/pipeline.json" + else: + from huggingface_hub import hf_hub_download + config_file = hf_hub_download(path, "pipeline.json") + + with open(config_file, 'r') as f: + args = json.load(f)['args'] + + _models = {} + for k, v in args['models'].items(): + model_path = v + if is_local: + local_model_path = os.path.join(path, v) + if os.path.exists(f"{local_model_path}.json") and os.path.exists(f"{local_model_path}.safetensors"): + model_path = local_model_path + _models[k] = models.from_pretrained(model_path) + + new_pipeline = Pipeline(_models) + new_pipeline._pretrained_args = args + return new_pipeline + + @property + def device(self) -> torch.device: + for model in self.models.values(): + if hasattr(model, 'device'): + return model.device + for model in self.models.values(): + if hasattr(model, 'parameters'): + return next(model.parameters()).device + raise RuntimeError("No device found.") + + def to(self, device: torch.device) -> None: + for model in self.models.values(): + model.to(device) + + def cuda(self) -> None: + self.to(torch.device("cuda")) + + def cpu(self) -> None: + self.to(torch.device("cpu")) diff --git a/third_party/TRELLIS/trellis/pipelines/samplers/__init__.py b/third_party/TRELLIS/trellis/pipelines/samplers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..54d412fc5d8eb662081a92a56ad078243988c2f9 --- /dev/null +++ b/third_party/TRELLIS/trellis/pipelines/samplers/__init__.py @@ -0,0 +1,2 @@ +from .base import Sampler +from .flow_euler import FlowEulerSampler, FlowEulerCfgSampler, FlowEulerGuidanceIntervalSampler \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/pipelines/samplers/base.py b/third_party/TRELLIS/trellis/pipelines/samplers/base.py new file mode 100644 index 0000000000000000000000000000000000000000..1966ce787009a5ee0c1ed06dce491525ff1dbcbf --- /dev/null +++ b/third_party/TRELLIS/trellis/pipelines/samplers/base.py @@ -0,0 +1,20 @@ +from typing import * +from abc import ABC, abstractmethod + + +class Sampler(ABC): + """ + A base class for samplers. + """ + + @abstractmethod + def sample( + self, + model, + **kwargs + ): + """ + Sample from a model. + """ + pass + \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/pipelines/samplers/classifier_free_guidance_mixin.py b/third_party/TRELLIS/trellis/pipelines/samplers/classifier_free_guidance_mixin.py new file mode 100644 index 0000000000000000000000000000000000000000..463c91412f54323f2a613294aaa70fd366e74555 --- /dev/null +++ b/third_party/TRELLIS/trellis/pipelines/samplers/classifier_free_guidance_mixin.py @@ -0,0 +1,20 @@ +from typing import * + + +def _without_local_conditioning(kwargs: dict) -> dict: + kwargs = dict(kwargs) + kwargs.pop('cond_list', None) + kwargs.pop('coords_dense_indices', None) + return kwargs + + +class ClassifierFreeGuidanceSamplerMixin: + """ + A mixin class for samplers that apply classifier-free guidance. + """ + + def _inference_model(self, model, x_t, t, cond, neg_cond, cfg_strength, **kwargs): + pred = super()._inference_model(model, x_t, t, cond, **kwargs) + neg_pred = super()._inference_model( + model, x_t, t, neg_cond, **_without_local_conditioning(kwargs)) + return (1 + cfg_strength) * pred - cfg_strength * neg_pred diff --git a/third_party/TRELLIS/trellis/pipelines/samplers/flow_euler.py b/third_party/TRELLIS/trellis/pipelines/samplers/flow_euler.py new file mode 100644 index 0000000000000000000000000000000000000000..add95b82b44f8c72792141a3aca2f3d8e490d675 --- /dev/null +++ b/third_party/TRELLIS/trellis/pipelines/samplers/flow_euler.py @@ -0,0 +1,324 @@ +from typing import * +import torch +import numpy as np +from tqdm import tqdm +from easydict import EasyDict as edict +from .base import Sampler +from .classifier_free_guidance_mixin import ClassifierFreeGuidanceSamplerMixin +from .guidance_interval_mixin import GuidanceIntervalSamplerMixin +from torch.nn.functional import max_pool3d +import torch.nn.functional as F + + +class FlowEulerSampler(Sampler): + """ + Generate samples from a flow-matching model using Euler sampling. + + Args: + sigma_min: The minimum scale of noise in flow. + """ + def __init__( + self, + sigma_min: float, + ): + self.sigma_min = sigma_min + + def _eps_to_xstart(self, x_t, t, eps): + assert x_t.shape == eps.shape + return (x_t - (self.sigma_min + (1 - self.sigma_min) * t) * eps) / (1 - t) + + def _xstart_to_eps(self, x_t, t, x_0): + assert x_t.shape == x_0.shape + return (x_t - (1 - t) * x_0) / (self.sigma_min + (1 - self.sigma_min) * t) + + def _v_to_xstart_eps(self, x_t, t, v): + assert x_t.shape == v.shape + eps = (1 - t) * v + x_t + x_0 = (1 - self.sigma_min) * x_t - (self.sigma_min + (1 - self.sigma_min) * t) * v + return x_0, eps + + def _inference_model(self, model, x_t, t, cond=None, **kwargs): + t = torch.tensor([1000 * t] * x_t.shape[0], device=x_t.device, dtype=torch.float32) + if cond is not None and cond.shape[0] == 1 and x_t.shape[0] > 1: + cond = cond.repeat(x_t.shape[0], *([1] * (len(cond.shape) - 1))) + return model(x_t, t, cond, **kwargs) + + def _get_model_prediction(self, model, x_t, t, cond=None, **kwargs): + pred_v = self._inference_model(model, x_t, t, cond, **kwargs) + pred_x_0, pred_eps = self._v_to_xstart_eps(x_t=x_t, t=t, v=pred_v) + return pred_x_0, pred_eps, pred_v + + @torch.no_grad() + def sample_once( + self, + model, + x_t, + t: float, + t_prev: float, + cond: Optional[Any] = None, + **kwargs + ): + """ + Sample x_{t-1} from the model using Euler method. + + Args: + model: The model to sample from. + x_t: The [N x C x ...] tensor of noisy inputs at time t. + t: The current timestep. + t_prev: The previous timestep. + cond: conditional information. + **kwargs: Additional arguments for model inference. + + Returns: + a dict containing the following + - 'pred_x_prev': x_{t-1}. + - 'pred_x_0': a prediction of x_0. + """ + pred_x_0, pred_eps, pred_v = self._get_model_prediction(model, x_t, t, cond, **kwargs) + pred_x_prev = x_t - (t - t_prev) * pred_v + return edict({"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0}) + + @torch.no_grad() + def sample( + self, + model, + noise, + cond: Optional[Any] = None, + steps: int = 50, + rescale_t: float = 1.0, + verbose: bool = True, + **kwargs + ): + """ + Generate samples from the model using Euler method. + + Args: + model: The model to sample from. + noise: The initial noise tensor. + cond: conditional information. + steps: The number of steps to sample. + rescale_t: The rescale factor for t. + verbose: If True, show a progress bar. + **kwargs: Additional arguments for model_inference. + + Returns: + a dict containing the following + - 'samples': the model samples. + - 'pred_x_t': a list of prediction of x_t. + - 'pred_x_0': a list of prediction of x_0. + """ + sample = noise + t_seq = np.linspace(1, 0, 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(steps)) + + t0 = None + t0_high = None + control_high = None + control_high_mask = None + low_control_mask = None + control_high_lat = None + lantent_high_control = None + + if 'control' in kwargs: + control = kwargs['control'] + t0 = t_seq[int(kwargs['t0_idx_value'])] + sample = noise * t0 + control * (1 - t0) # [1, 8, 16, 16, 16] + + if (kwargs.get('control_high') is not None) and (kwargs.get('local_tau_mode', None) == 'guidance'): + control_high = kwargs['control_high'] + t0_high = t_seq[int(kwargs['t0_idx_value_high_control'])] + control_high_mask = None + elif (kwargs.get('control_high') is not None) and (kwargs.get('local_tau_mode', None) == 'masking'): + control_high_mask = kwargs['control_high'] + # control_low_mask = kwargs['control_low_mask'] + + # print(control_low_mask) + # non_zero_ratio = (control_low_mask > 0).float().mean().item() + # print(f"Initial non-zero ratio of control_low_mask: {non_zero_ratio:.4f}", flush=True) + + + # print("Applying morphological dilation to control_low_mask... Shape:", control_low_mask.shape, flush=True) + # control_low_mask = max_pool3d(1 - control_low_mask, kernel_size=5, stride=1, padding=2) + # print("Dilation applied. Shape after dilation:", control_low_mask.shape, flush=True) + # control_low_mask = 1 - control_low_mask + + lantent_high_control = kwargs['latent_high_control'] + t0_high = t_seq[int(kwargs['t0_idx_value_high_control'])] + control_high = None + elif (kwargs.get('control_low_mask') is not None) and (kwargs.get('local_tau_mode', None) == 'low_control_mask'): + low_control_mask = kwargs['control_low_mask'] + non_zero_ratio = (low_control_mask > 0).float().mean().item() + print(f"Initial non-zero ratio of low_control_mask: {non_zero_ratio:.4f}", flush=True) + mean_value_non_zero = low_control_mask[low_control_mask > 0].mean().item() + print(f"Mean value of non-zero elements in low_control_mask: {mean_value_non_zero:.4f}", flush=True) + # Binarize the low_control_mask + low_control_mask = (low_control_mask > 0).float() + non_zero_ratio_after_binarization = (low_control_mask > 0).float().mean().item() + print(f"Non-zero ratio of low_control_mask after binarization: {non_zero_ratio_after_binarization:.4f}", flush=True) + print("mean value of non-zero elements in low_control_mask after binarization: {:.4f}".format(low_control_mask[low_control_mask > 0].mean().item()), flush=True) + + t0_high = t_seq[int(kwargs['t0_idx_value_high_control'])] + control_high_lat = kwargs.get('control_high') + + control_high_mask = None + control_high = None + else: + print("No high control provided, skipping high control adjustment initialization.") + control_high_mask = None + control_high = None + + args = {'neg_cond': kwargs['neg_cond'], + 'cfg_strength': kwargs['cfg_strength'], + 'cfg_interval': kwargs['cfg_interval']} + ret = edict({"samples": None, "pred_x_t": [], "pred_x_0": []}) + + print(f"Starting sampling with {steps} steps...", flush=True) + print(f"Low t0 for control: {t0:.4f}" if t0 is not None else "No low t0 for control.", flush=True) + for t, t_prev in tqdm(t_pairs, desc="Sampling", disable=not verbose): + # print(f"Current t: {t:.4f}, Previous t: {t_prev:.4f}" + (f", t0: {t0:.4f}" if t0 is not None else "") + # + (f", t0_high: {t0_high:.4f}" if t0_high is not None else ""), flush=True) + local_tau_mode = kwargs.get('local_tau_mode', None) + applying_high_guidance = (control_high is not None) and (t0_high is not None) and t > t0_high + applying_high_mask = (control_high_mask is not None) and (t0_high is not None) and t > t0_high + applying_low_mask_blend = ( + local_tau_mode == 'low_control_mask' + and low_control_mask is not None + and control_high_lat is not None + and t0_high is not None + and t > t0_high + ) + if 'control' in kwargs and t0 is not None and t > t0: + print("Skipping step with control due to t > t0...", flush=True) + continue + + out = self.sample_once(model, sample, t, t_prev, cond, **args) + sample = out.pred_x_prev + + if kwargs['polyak_update_tau'] > 0.0: + if (control_high is not None)and t > t0_high: + print("Applying high control adjustment...", flush=True) + # sample_high = noise * t + control_high * (1 - t) # [1, 8, 16, 16, 16] + sample = sample * (1 - kwargs['polyak_update_tau']) + control_high * kwargs['polyak_update_tau'] + elif (control_high_mask is not None) and t > t0_high: + print("Applying high control masking...", flush=True) + noise = torch.randn_like(sample) + sample_gt_t = lantent_high_control # noise * t_prev + lantent_high_control * (1 - t_prev) + + polyak_low = 0.08 + polyak_high = 1.0 + sample_low_control = (1 - polyak_low) * sample * (1-control_high_mask) + polyak_low * sample_gt_t * (1-control_high_mask) # sample * (1 - control_high_mask) + sample_high_control = (1 - polyak_high) * sample * control_high_mask + polyak_high * lantent_high_control * control_high_mask + sample = sample_low_control + sample_high_control + elif applying_low_mask_blend: + print("Applying low-control-mask high-control blend...", flush=True) + sample_gt_t = control_high_lat + + polyak_high = kwargs['polyak_update_tau'] + sample_low_control = sample * low_control_mask + sample_high_control = (1 - polyak_high) * sample + polyak_high * sample_gt_t + sample = sample_low_control + sample_high_control * (1 - low_control_mask) + + resampling_steps = kwargs.get('n_repaint_steps', 10) + print(f"Doing repaint resampling for {resampling_steps} steps to improve blending.", flush=True) + for i in range(resampling_steps): + noise = torch.randn_like(sample) + diff_t = abs(t - t_prev) + noised_sample = sample * (1 - diff_t) + noise * diff_t + out = self.sample_once(model, noised_sample, t_prev+diff_t, t_prev, cond, **args) + sample = out.pred_x_prev + sample_low_control = sample * low_control_mask + sample_high_control = (1 - polyak_high) * sample + polyak_high * sample_gt_t + sample = sample_low_control + sample_high_control * (1 - low_control_mask) + else: + print("No high control adjustment applied for this step.", flush=True) + + else: + print("No high control adjustment applied for this step.", flush=True) + + ret.pred_x_t.append(out.pred_x_prev) + ret.pred_x_0.append(out.pred_x_0) + ret.samples = sample.detach() + return ret + + +class FlowEulerCfgSampler(ClassifierFreeGuidanceSamplerMixin, FlowEulerSampler): + """ + Generate samples from a flow-matching model using Euler sampling with classifier-free guidance. + """ + @torch.no_grad() + def sample( + self, + model, + noise, + cond, + neg_cond, + steps: int = 50, + rescale_t: float = 1.0, + cfg_strength: float = 3.0, + verbose: bool = True, + **kwargs + ): + """ + Generate samples from the model using Euler method. + + Args: + model: The model to sample from. + noise: The initial noise tensor. + cond: conditional information. + neg_cond: negative conditional information. + steps: The number of steps to sample. + rescale_t: The rescale factor for t. + cfg_strength: The strength of classifier-free guidance. + verbose: If True, show a progress bar. + **kwargs: Additional arguments for model_inference. + + Returns: + a dict containing the following + - 'samples': the model samples. + - 'pred_x_t': a list of prediction of x_t. + - 'pred_x_0': a list of prediction of x_0. + """ + return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, **kwargs) + + +class FlowEulerGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, FlowEulerSampler): + """ + Generate samples from a flow-matching model using Euler sampling with classifier-free guidance and interval. + """ + @torch.no_grad() + def sample( + self, + model, + noise, + cond, + neg_cond, + steps: int = 50, + rescale_t: float = 1.0, + cfg_strength: float = 3.0, + cfg_interval: Tuple[float, float] = (0.0, 1.0), + verbose: bool = True, + **kwargs + ): + """ + Generate samples from the model using Euler method. + + Args: + model: The model to sample from. + noise: The initial noise tensor. + cond: conditional information. + neg_cond: negative conditional information. + steps: The number of steps to sample. + rescale_t: The rescale factor for t. + cfg_strength: The strength of classifier-free guidance. + cfg_interval: The interval for classifier-free guidance. + verbose: If True, show a progress bar. + **kwargs: Additional arguments for model_inference. + + Returns: + a dict containing the following + - 'samples': the model samples. + - 'pred_x_t': a list of prediction of x_t. + - 'pred_x_0': a list of prediction of x_0. + """ + return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, cfg_interval=cfg_interval, **kwargs) diff --git a/third_party/TRELLIS/trellis/pipelines/samplers/guidance_interval_mixin.py b/third_party/TRELLIS/trellis/pipelines/samplers/guidance_interval_mixin.py new file mode 100644 index 0000000000000000000000000000000000000000..6c86f9d70908fdb47066819fe69146f4c329df20 --- /dev/null +++ b/third_party/TRELLIS/trellis/pipelines/samplers/guidance_interval_mixin.py @@ -0,0 +1,23 @@ +from typing import * + + +def _without_local_conditioning(kwargs: dict) -> dict: + kwargs = dict(kwargs) + kwargs.pop('cond_list', None) + kwargs.pop('coords_dense_indices', None) + return kwargs + + +class GuidanceIntervalSamplerMixin: + """ + A mixin class for samplers that apply classifier-free guidance with interval. + """ + + def _inference_model(self, model, x_t, t, cond, neg_cond, cfg_strength, cfg_interval, **kwargs): + if cfg_interval[0] <= t <= cfg_interval[1]: + pred = super()._inference_model(model, x_t, t, cond, **kwargs) + neg_pred = super()._inference_model( + model, x_t, t, neg_cond, **_without_local_conditioning(kwargs)) + return (1 + cfg_strength) * pred - cfg_strength * neg_pred + else: + return super()._inference_model(model, x_t, t, cond, **kwargs) diff --git a/third_party/TRELLIS/trellis/pipelines/trellis_image_to_3d.py b/third_party/TRELLIS/trellis/pipelines/trellis_image_to_3d.py new file mode 100644 index 0000000000000000000000000000000000000000..6b1766a7f5a3bc30da9af39d156c1c59c0853a7a --- /dev/null +++ b/third_party/TRELLIS/trellis/pipelines/trellis_image_to_3d.py @@ -0,0 +1,527 @@ +from typing import * +import open3d_pycg as o3d +from contextlib import contextmanager +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from torchvision import transforms +from PIL import Image +import rembg +from .base import Pipeline +from . import samplers +from ..modules import sparse as sp +from lib.util.model_revisions import dinov2_hub_repository + + +class TrellisImageTo3DPipeline(Pipeline): + """ + Pipeline for inferring Trellis image-to-3D models. + + Args: + models (dict[str, nn.Module]): The models to use in the pipeline. + sparse_structure_sampler (samplers.Sampler): The sampler for the sparse structure. + slat_sampler (samplers.Sampler): The sampler for the structured latent. + slat_normalization (dict): The normalization parameters for the structured latent. + image_cond_model (str): The name of the image conditioning model. + """ + def __init__( + self, + models: dict[str, nn.Module] = None, + sparse_structure_sampler: samplers.Sampler = None, + slat_sampler: samplers.Sampler = None, + slat_normalization: dict = None, + image_cond_model: str = None, + ): + if models is None: + return + super().__init__(models) + self.sparse_structure_sampler = sparse_structure_sampler + self.slat_sampler = slat_sampler + self.sparse_structure_sampler_params = {} + self.slat_sampler_params = {} + self.slat_normalization = slat_normalization + self.rembg_session = None + self._init_image_cond_model(image_cond_model) + + @staticmethod + def from_pretrained(path: str) -> "TrellisImageTo3DPipeline": + """ + Load a pretrained model. + + Args: + path (str): The path to the model. Can be either local path or a Hugging Face repository. + """ + pipeline = super(TrellisImageTo3DPipeline, TrellisImageTo3DPipeline).from_pretrained(path) + new_pipeline = TrellisImageTo3DPipeline() + new_pipeline.__dict__ = pipeline.__dict__ + args = pipeline._pretrained_args + + + new_pipeline.sparse_structure_sampler = getattr(samplers, args['sparse_structure_sampler']['name'])(**args['sparse_structure_sampler']['args']) + new_pipeline.sparse_structure_sampler_params = args['sparse_structure_sampler']['params'] + + new_pipeline.slat_sampler = getattr(samplers, args['slat_sampler']['name'])(**args['slat_sampler']['args']) + new_pipeline.slat_sampler_params = args['slat_sampler']['params'] + + new_pipeline.slat_normalization = args['slat_normalization'] + + new_pipeline._init_image_cond_model(args['image_cond_model']) + + return new_pipeline + + def _init_image_cond_model(self, name: str): + """ + Initialize the image conditioning model. + """ + dinov2_model = torch.hub.load(dinov2_hub_repository(), name, pretrained=True, + trust_repo=True, skip_validation=True) + dinov2_model.eval() + self.models['image_cond_model'] = dinov2_model + transform = transforms.Compose([ + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ]) + self.image_cond_model_transform = transform + + def voxelize(self, mesh: o3d.geometry.TriangleMesh) -> torch.Tensor: + """ + Voxelize a mesh. + + Args: + mesh (o3d.geometry.TriangleMesh): The mesh to voxelize. + sha256 (str): The SHA256 hash of the mesh. + output_dir (str): The output directory. + """ + vertices = np.asarray(mesh.vertices) + aabb = np.stack([vertices.min(0), vertices.max(0)]) + center = (aabb[0] + aabb[1]) / 2 + scale = (aabb[1] - aabb[0]).max() + vertices = (vertices - center) / scale + vertices = np.clip(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()]) + return torch.tensor(vertices).int().cuda() + + def preprocess_image(self, input: Image.Image) -> Image.Image: + """ + Preprocess the input image. + """ + # if has alpha channel, use it directly; otherwise, remove background + has_alpha = False + if input.mode == 'RGBA': + alpha = np.array(input)[:, :, 3] + if not np.all(alpha == 255): + has_alpha = True + if has_alpha: + output = input + else: + input = input.convert('RGB') + max_size = max(input.size) + scale = min(1, 1024 / max_size) + if scale < 1: + input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS) + if getattr(self, 'rembg_session', None) is None: + self.rembg_session = rembg.new_session('u2net') + output = rembg.remove(input, session=self.rembg_session) + output_np = np.array(output) + alpha = output_np[:, :, 3] + bbox = np.argwhere(alpha > 0.8 * 255) + bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0]) + center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2 + size = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) + size = int(size * 1.2) + bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2 + output = output.crop(bbox) # type: ignore + output = output.resize((518, 518), Image.Resampling.LANCZOS) + output = np.array(output).astype(np.float32) / 255 + output = output[:, :, :3] * output[:, :, 3:4] + output = Image.fromarray((output * 255).astype(np.uint8)) + return output + + @torch.no_grad() + def encode_image(self, image: Union[torch.Tensor, list[Image.Image]]) -> torch.Tensor: + """ + Encode the image. + + Args: + image (Union[torch.Tensor, list[Image.Image]]): The image to encode + + Returns: + torch.Tensor: The encoded features. + """ + if isinstance(image, torch.Tensor): + assert image.ndim == 4, "Image tensor should be batched (B, C, H, W)" + elif isinstance(image, list): + assert all(isinstance(i, Image.Image) for i in image), "Image list should be list of PIL images" + image = [i.resize((518, 518), Image.LANCZOS) for i in image] + image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image] + image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image] + image = torch.stack(image).to(self.device) + else: + raise ValueError(f"Unsupported type of image: {type(image)}") + + image = self.image_cond_model_transform(image).to(self.device) + features = self.models['image_cond_model'](image, is_training=True)['x_prenorm'] + patchtokens = F.layer_norm(features, features.shape[-1:]) + return patchtokens + + def get_cond(self, image: Union[torch.Tensor, list[Image.Image]]) -> dict: + """ + Get the conditioning information for the model. + + Args: + image (Union[torch.Tensor, list[Image.Image]]): The image prompts. + + Returns: + dict: The conditioning information + """ + cond = self.encode_image(image) + neg_cond = torch.zeros_like(cond) + return { + 'cond': cond, + 'neg_cond': neg_cond, + } + # torch.Size([1, 1374, 1024]) + + + def sample_sparse_structure( + self, + cond: dict, + num_samples: int = 1, + sampler_params: dict = {}, + ) -> torch.Tensor: + """ + Sample sparse structures with the given conditioning. + + Args: + cond (dict): The conditioning information. + num_samples (int): The number of samples to generate. + sampler_params (dict): Additional parameters for the sampler. + """ + # Sample occupancy latent + flow_model = self.models['sparse_structure_flow_model'] + reso = flow_model.resolution + noise = torch.randn(num_samples, flow_model.in_channels, reso, reso, reso).to(self.device) + sampler_params = {**self.sparse_structure_sampler_params, **sampler_params} + z_s = self.sparse_structure_sampler.sample( + flow_model, + noise, + **cond, + **sampler_params, + verbose=True + ).samples + + # Decode occupancy latent + decoder = self.models['sparse_structure_decoder'] + coords = torch.argwhere(decoder(z_s)>0)[:, [0, 2, 3, 4]].int() + + return coords + + def decode_slat( + self, + slat: sp.SparseTensor, + formats: List[str] = ['mesh', 'gaussian', 'radiance_field'], + ) -> dict: + """ + Decode the structured latent. + + Args: + slat (sp.SparseTensor): The structured latent. + formats (List[str]): The formats to decode the structured latent to. + + Returns: + dict: The decoded structured latent. + """ + ret = {} + if 'mesh' in formats: + ret['mesh'] = self.models['slat_decoder_mesh'](slat) + if 'gaussian' in formats: + ret['gaussian'] = self.models['slat_decoder_gs'](slat) + if 'radiance_field' in formats: + ret['radiance_field'] = self.models['slat_decoder_rf'](slat) + return ret + + def sample_slat( + self, + cond: dict, + coords: torch.Tensor, + sampler_params: dict = {}, + ) -> sp.SparseTensor: + """ + Sample structured latent with the given conditioning. + + Args: + cond (dict): The conditioning information. + coords (torch.Tensor): The coordinates of the sparse structure. + sampler_params (dict): Additional parameters for the sampler. + """ + # Sample structured latent + flow_model = self.models['slat_flow_model'] + noise = sp.SparseTensor( + feats=torch.randn(coords.shape[0], flow_model.in_channels).to(self.device), + coords=coords, + ) + sampler_params = {**self.slat_sampler_params, **sampler_params} + slat = self.slat_sampler.sample( + flow_model, + noise, + **cond, + **sampler_params, + verbose=True + ).samples + + std = torch.tensor(self.slat_normalization['std'])[None].to(slat.device) + mean = torch.tensor(self.slat_normalization['mean'])[None].to(slat.device) + slat = slat * std + mean + + return slat + + @torch.no_grad() + def run_diffusion_step( + self, + seed: int = 42, + noise: sp.SparseTensor = None, + cond: dict = None, + sampler_params: dict = {}, + **kwargs, + ) -> dict: + """ + Run the texture generation(2nd stage) pipeline. + + Args: + binary_voxel (np.ndarray): The input binary voxel. + image (Image.Image or a list of Image.Image): The image prompt(s). + num_samples (int): The number of samples to generate. + sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler. + slat_sampler_params (dict): Additional parameters for the structured latent sampler. + preprocess_image (bool): Whether to preprocess the image. + """ + + torch.manual_seed(seed) + + # Sample structured latent + flow_model = self.models['slat_flow_model'] + sampler_params = {**self.slat_sampler_params, **sampler_params} + slat = self.slat_sampler.sample( + flow_model, + noise, + **cond, + **sampler_params, + verbose=True + ).samples + + std = torch.tensor(self.slat_normalization['std'])[None].to(slat.device) + mean = torch.tensor(self.slat_normalization['mean'])[None].to(slat.device) + slat = slat * std + mean + return slat + + @torch.no_grad() + def run_partial_denoising( + self, + seed: int = 42, + noise: sp.SparseTensor = None, + image: Image.Image = None, + sampler_params: dict = {}, + formats: List[str] = ["mesh", "gaussian", "radiance_field"], + **kwargs, + ) -> dict: + """ + Run the texture generation(2nd stage) pipeline. + + Args: + binary_voxel (np.ndarray): The input binary voxel. + image (Image.Image or a list of Image.Image): The image prompt(s). + num_samples (int): The number of samples to generate. + sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler. + slat_sampler_params (dict): Additional parameters for the structured latent sampler. + preprocess_image (bool): Whether to preprocess the image. + """ + + torch.manual_seed(seed) + + cond = torch.zeros((1, 1374, 1024)).cuda() + neg_cond = torch.zeros_like(cond) + cond = { + 'cond': cond, + 'neg_cond': neg_cond, + } + + image = self.preprocess_image(image) + cond = self.get_cond([image]) + + # Sample structured latent + flow_model = self.models['slat_flow_model'] + sampler_params = {**self.slat_sampler_params, **sampler_params} + slat = self.slat_sampler.sample( + flow_model, + noise, + **cond, + **sampler_params, + verbose=True + ).samples + + std = torch.tensor(self.slat_normalization['std'])[None].to(slat.device) + mean = torch.tensor(self.slat_normalization['mean'])[None].to(slat.device) + slat = slat * std + mean + return self.decode_slat(slat, formats) + + @torch.no_grad() + def run_detail_variation( + self, + mesh: o3d.geometry.TriangleMesh, + image: Image.Image, + num_samples: int = 1, + seed: int = 42, + slat_sampler_params: dict = {}, + formats: List[str] = ["mesh", "gaussian", "radiance_field"], + preprocess_image: bool = True, + **kwargs, + ) -> dict: + """ + Run the texture generation(2nd stage) pipeline. + + Args: + binary_voxel (np.ndarray): The input binary voxel. + image (Image.Image or a list of Image.Image): The image prompt(s). + num_samples (int): The number of samples to generate. + sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler. + slat_sampler_params (dict): Additional parameters for the structured latent sampler. + preprocess_image (bool): Whether to preprocess the image. + """ + + if preprocess_image: + image = self.preprocess_image(image) + cond = self.get_cond([image]) + + coords = self.voxelize(mesh) + coords = torch.cat([ + torch.arange(num_samples).repeat_interleave(coords.shape[0], 0)[:, None].int().cuda(), + coords.repeat(num_samples, 1) + ], 1) + torch.manual_seed(seed) + slat = self.sample_slat(cond, coords, slat_sampler_params) + return self.decode_slat(slat, formats) + + @torch.no_grad() + def run( + self, + image: Image.Image, + num_samples: int = 1, + seed: int = 42, + sparse_structure_sampler_params: dict = {}, + slat_sampler_params: dict = {}, + formats: List[str] = ['mesh', 'gaussian', 'radiance_field'], + preprocess_image: bool = True, + ) -> dict: + """ + Run the pipeline. + + Args: + image (Image.Image): The image prompt. + num_samples (int): The number of samples to generate. + seed (int): The random seed. + sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler. + slat_sampler_params (dict): Additional parameters for the structured latent sampler. + formats (List[str]): The formats to decode the structured latent to. + preprocess_image (bool): Whether to preprocess the image. + """ + if preprocess_image: + image = self.preprocess_image(image) + cond = self.get_cond([image]) + torch.manual_seed(seed) + coords = self.sample_sparse_structure(cond, num_samples, sparse_structure_sampler_params) + slat = self.sample_slat(cond, coords, slat_sampler_params) + return self.decode_slat(slat, formats) + + @contextmanager + def inject_sampler_multi_image( + self, + sampler_name: str, + num_images: int, + num_steps: int, + mode: Literal['stochastic', 'multidiffusion'] = 'stochastic', + ): + """ + Inject a sampler with multiple images as condition. + + Args: + sampler_name (str): The name of the sampler to inject. + num_images (int): The number of images to condition on. + num_steps (int): The number of steps to run the sampler for. + """ + sampler = getattr(self, sampler_name) + setattr(sampler, f'_old_inference_model', sampler._inference_model) + + if mode == 'stochastic': + if num_images > num_steps: + print(f"\033[93mWarning: number of conditioning images is greater than number of steps for {sampler_name}. " + "This may lead to performance degradation.\033[0m") + + cond_indices = (np.arange(num_steps) % num_images).tolist() + def _new_inference_model(self, model, x_t, t, cond, **kwargs): + cond_idx = cond_indices.pop(0) + cond_i = cond[cond_idx:cond_idx+1] + return self._old_inference_model(model, x_t, t, cond=cond_i, **kwargs) + + elif mode =='multidiffusion': + from .samplers import FlowEulerSampler + def _new_inference_model(self, model, x_t, t, cond, neg_cond, cfg_strength, cfg_interval, **kwargs): + if cfg_interval[0] <= t <= cfg_interval[1]: + preds = [] + for i in range(len(cond)): + preds.append(FlowEulerSampler._inference_model(self, model, x_t, t, cond[i:i+1], **kwargs)) + pred = sum(preds) / len(preds) + neg_pred = FlowEulerSampler._inference_model(self, model, x_t, t, neg_cond, **kwargs) + return (1 + cfg_strength) * pred - cfg_strength * neg_pred + else: + preds = [] + for i in range(len(cond)): + preds.append(FlowEulerSampler._inference_model(self, model, x_t, t, cond[i:i+1], **kwargs)) + pred = sum(preds) / len(preds) + return pred + + else: + raise ValueError(f"Unsupported mode: {mode}") + + sampler._inference_model = _new_inference_model.__get__(sampler, type(sampler)) + + yield + + sampler._inference_model = sampler._old_inference_model + delattr(sampler, f'_old_inference_model') + + @torch.no_grad() + def run_multi_image( + self, + images: List[Image.Image], + num_samples: int = 1, + seed: int = 42, + sparse_structure_sampler_params: dict = {}, + slat_sampler_params: dict = {}, + formats: List[str] = ['mesh', 'gaussian', 'radiance_field'], + preprocess_image: bool = True, + mode: Literal['stochastic', 'multidiffusion'] = 'stochastic', + ) -> dict: + """ + Run the pipeline with multiple images as condition + + Args: + images (List[Image.Image]): The multi-view images of the assets + num_samples (int): The number of samples to generate. + sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler. + slat_sampler_params (dict): Additional parameters for the structured latent sampler. + preprocess_image (bool): Whether to preprocess the image. + """ + if preprocess_image: + images = [self.preprocess_image(image) for image in images] + cond = self.get_cond(images) + cond['neg_cond'] = cond['neg_cond'][:1] + torch.manual_seed(seed) + ss_steps = {**self.sparse_structure_sampler_params, **sparse_structure_sampler_params}.get('steps') + with self.inject_sampler_multi_image('sparse_structure_sampler', len(images), ss_steps, mode=mode): + coords = self.sample_sparse_structure(cond, num_samples, sparse_structure_sampler_params) + slat_steps = {**self.slat_sampler_params, **slat_sampler_params}.get('steps') + with self.inject_sampler_multi_image('slat_sampler', len(images), slat_steps, mode=mode): + slat = self.sample_slat(cond, coords, slat_sampler_params) + return self.decode_slat(slat, formats) diff --git a/third_party/TRELLIS/trellis/pipelines/trellis_text_to_3d.py b/third_party/TRELLIS/trellis/pipelines/trellis_text_to_3d.py new file mode 100644 index 0000000000000000000000000000000000000000..e6bef380a63d70a35632e2b96cbf4e29462cd517 --- /dev/null +++ b/third_party/TRELLIS/trellis/pipelines/trellis_text_to_3d.py @@ -0,0 +1,732 @@ +from typing import * +import torch +import torch.nn as nn +import numpy as np +import copy +from transformers import CLIPTextModel, AutoTokenizer +from torchvision import transforms +import torch.nn.functional as F +import rembg +from PIL import Image +import open3d_pycg as o3d +import os +import time +import imageio +import logging +from .base import Pipeline +from . import samplers +from ..modules import sparse as sp +from ..representations import Gaussian +from ..utils import render_utils +# from gui import utils (Directly copied the functions here because of coliding dependencies) +from pathlib import Path +from sklearn.decomposition import PCA +from lib.util.model_revisions import dinov2_hub_repository + +log = logging.getLogger(__name__) + + +def merge_meshes(mesh_list): + merged = o3d.geometry.TriangleMesh() + v_offset = 0 + + for mesh in mesh_list: + vertices = np.asarray(mesh.vertices) + triangles = np.asarray(mesh.triangles) + v_offset + + merged.vertices.extend(o3d.utility.Vector3dVector(vertices)) + merged.triangles.extend(o3d.utility.Vector3iVector(triangles)) + + v_offset += len(vertices) + + return merged + + +def voxelgrid_to_open3d(voxels: np.ndarray, threshold=0.5): + if len(voxels.shape) > 3: + C, D, H, W = voxels.shape + flat_feats = voxels.reshape(C, -1).transpose(1,0) + pca = PCA(n_components=3) + reduced = pca.fit_transform(flat_feats) + # Compute feature norm and PCA color std + + # Normalize for RGB + reduced -= reduced.min(0) + reduced /= reduced.max(0) + 1e-6 + + # Compute norms and color std + norms = np.linalg.norm(flat_feats, axis=1) + color_std = np.std(reduced, axis=1) + + # Filter: active voxels with non-trivial color + mask = (norms > threshold) & (color_std > 1e-3) + + # zz, yy, xx = np.meshgrid(np.arange(D), np.arange(H), np.arange(W), indexing='ij') + xx, yy, zz = np.meshgrid(np.arange(D), np.arange(H), np.arange(W), indexing='ij') + coords = np.stack([xx, yy, zz], axis=-1).reshape(-1, 3) + valid_coords = coords[mask] + valid_colors = reduced[mask] + + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(valid_coords.astype(np.float32)) + pcd.colors = o3d.utility.Vector3dVector(valid_colors.astype(np.float32)) + else: + coords = np.argwhere(voxels > threshold) + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(coords) + return pcd + + +def save_voxelgrid_as_ply(voxels: np.ndarray, filename: str, threshold=0.5): + pcd = voxelgrid_to_open3d(voxels, threshold) + o3d.io.write_point_cloud(filename, pcd) + + +def voxelize_sq_francis(file_name): + superquadric_mesh = o3d.io.read_triangle_mesh(file_name) + vertices = np.clip(np.asarray(superquadric_mesh.vertices), -0.5 + 1e-6, 0.5 - 1e-6) + superquadric_mesh.vertices = o3d.utility.Vector3dVector(vertices) + theta = np.pi / 2 + #superquadric_mesh.rotate(R_x, center=(0, 0, 0)) + voxel_grid = o3d.geometry.VoxelGrid.create_from_triangle_mesh_within_bounds( + superquadric_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()]) + unique_points = np.unique(vertices, axis=0) + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(unique_points) + o3d.io.write_point_cloud("merged_mesh_voxelized.ply", pcd) + + zeros = np.zeros((unique_points.shape[0], 1)) + unique_points_4d = np.hstack((zeros, unique_points)) # shape [N, 4] + unique_points_4d_torch = torch.from_numpy(unique_points_4d).to(dtype=torch.int32, device='cpu') + my_coords_orig = unique_points_4d_torch + + coords_dense = torch.ones(1, 1, 64, 64, 64).to(device='cpu', dtype=torch.float32) * 0.0 + for i in range(my_coords_orig.shape[0]): + x, y, z = my_coords_orig[i, 1], my_coords_orig[i, 2], my_coords_orig[i, 3] + coords_dense[0, 0, x, y, z] = 1.0 + return coords_dense + + +class TrellisTextTo3DPipeline(Pipeline): + """ + Pipeline for inferring Trellis text-to-3D models. + + Args: + models (dict[str, nn.Module]): The models to use in the pipeline. + sparse_structure_sampler (samplers.Sampler): The sampler for the sparse structure. + slat_sampler (samplers.Sampler): The sampler for the structured latent. + slat_normalization (dict): The normalization parameters for the structured latent. + text_cond_model (str): The name of the text conditioning model. + """ + def __init__( + self, + models: dict[str, nn.Module] = None, + sparse_structure_sampler: samplers.Sampler = None, + slat_sampler: samplers.Sampler = None, + slat_normalization: dict = None, + text_cond_model: str = None, + image_cond_model: str = None, + ): + if models is None: + return + super().__init__(models) + self.sparse_structure_sampler = sparse_structure_sampler + self.slat_sampler = slat_sampler + self.sparse_structure_sampler_params = {} + self.slat_sampler_params = {} + self.slat_normalization = slat_normalization + self._init_text_cond_model(text_cond_model) + self.image_cond_model_name = image_cond_model + self.image_cond_model_transform = None + if os.environ.get("TRELLIS_EAGER_IMAGE_COND", "").strip().lower() in {"1", "true", "yes", "on"}: + self._init_image_cond_model(image_cond_model) + + @staticmethod + def from_pretrained(path: str) -> "TrellisTextTo3DPipeline": + """ + Load a pretrained model. + + Args: + path (str): The path to the model. Can be either local path or a Hugging Face repository. + """ + pipeline = super(TrellisTextTo3DPipeline, TrellisTextTo3DPipeline).from_pretrained(path) + new_pipeline = TrellisTextTo3DPipeline() + new_pipeline.__dict__ = pipeline.__dict__ + args = pipeline._pretrained_args + + new_pipeline.sparse_structure_sampler = getattr(samplers, args['sparse_structure_sampler']['name'])(**args['sparse_structure_sampler']['args']) + new_pipeline.sparse_structure_sampler_params = args['sparse_structure_sampler']['params'] + + new_pipeline.slat_sampler = getattr(samplers, args['slat_sampler']['name'])(**args['slat_sampler']['args']) + new_pipeline.slat_sampler_params = args['slat_sampler']['params'] + + new_pipeline.slat_normalization = args['slat_normalization'] + + new_pipeline._init_text_cond_model(args['text_cond_model']) + + new_pipeline.image_cond_model_name = args.get('image_cond_model') + new_pipeline.image_cond_model_transform = None + if ( + new_pipeline.image_cond_model_name + and os.environ.get("TRELLIS_EAGER_IMAGE_COND", "").strip().lower() in {"1", "true", "yes", "on"} + ): + new_pipeline._init_image_cond_model(args['image_cond_model']) + + return new_pipeline + + def _init_text_cond_model(self, name: str): + """ + Initialize the text conditioning model. + """ + # load model + model = CLIPTextModel.from_pretrained(name) + tokenizer = AutoTokenizer.from_pretrained(name) + model.eval() + model = model.cuda() + self.text_cond_model = { + 'model': model, + 'tokenizer': tokenizer, + } + self.text_cond_model['null_cond'] = self.encode_text(['']) + + def _init_image_cond_model(self, name: str): + """ + Initialize the image conditioning model. + """ + if not name: + raise ValueError("No image conditioning model configured for this pipeline") + if ( + 'image_cond_model' in self.models + and getattr(self, 'image_cond_model_name', None) == name + and getattr(self, 'image_cond_model_transform', None) is not None + ): + return + dinov2_model = torch.hub.load(dinov2_hub_repository(), name, pretrained=True, + trust_repo=True, skip_validation=True) + target_device = self.device + dinov2_model.eval().to(target_device) + self.models['image_cond_model'] = dinov2_model + self.image_cond_model_name = name + transform = transforms.Compose([ + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ]) + self.image_cond_model_transform = transform + + def _ensure_image_cond_model(self) -> None: + if 'image_cond_model' not in self.models or getattr(self, 'image_cond_model_transform', None) is None: + self._init_image_cond_model(getattr(self, 'image_cond_model_name', None)) + + @torch.no_grad() + def encode_text(self, text: List[str]) -> torch.Tensor: + """ + Encode the text. + """ + assert isinstance(text, list) and all(isinstance(t, str) for t in text), "text must be a list of strings" + encoding = self.text_cond_model['tokenizer'](text, max_length=77, padding='max_length', truncation=True, return_tensors='pt') + tokens = encoding['input_ids'].cuda() + embeddings = self.text_cond_model['model'](input_ids=tokens).last_hidden_state + + return embeddings + + @torch.no_grad() + def encode_image(self, image: Union[torch.Tensor, list[Image.Image]]) -> torch.Tensor: + """ + Encode the image. + + Args: + image (Union[torch.Tensor, list[Image.Image]]): The image to encode + + Returns: + torch.Tensor: The encoded features. + """ + if isinstance(image, torch.Tensor): + assert image.ndim == 4, "Image tensor should be batched (B, C, H, W)" + elif isinstance(image, list): + assert all(isinstance(i, Image.Image) for i in image), "Image list should be list of PIL images" + image = [i.resize((518, 518), Image.LANCZOS) for i in image] + image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image] + image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image] + image = torch.stack(image).to(self.device) + else: + raise ValueError(f"Unsupported type of image: {type(image)}") + + self._ensure_image_cond_model() + image = self.image_cond_model_transform(image).to(self.device) + features = self.models['image_cond_model'](image, is_training=True)['x_prenorm'] + patchtokens = F.layer_norm(features, features.shape[-1:]) + return patchtokens + + def preprocess_image(self, input: Image.Image) -> Image.Image: + """ + Preprocess the input image. + """ + # if has alpha channel, use it directly; otherwise, remove background + has_alpha = False + if input.mode == 'RGBA': + alpha = np.array(input)[:, :, 3] + if not np.all(alpha == 255): + has_alpha = True + if has_alpha: + output = input + else: + input = input.convert('RGB') + max_size = max(input.size) + scale = min(1, 1024 / max_size) + if scale < 1: + input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS) + if getattr(self, 'rembg_session', None) is None: + self.rembg_session = rembg.new_session('u2net') + output = rembg.remove(input, session=self.rembg_session) + output_np = np.array(output) + alpha = output_np[:, :, 3] + bbox = np.argwhere(alpha > 0.8 * 255) + bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0]) + center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2 + size = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) + size = int(size * 1.2) + bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2 + output = output.crop(bbox) # type: ignore + output = output.resize((518, 518), Image.Resampling.LANCZOS) + output = np.array(output).astype(np.float32) / 255 + output = output[:, :, :3] * output[:, :, 3:4] + output = Image.fromarray((output * 255).astype(np.uint8)) + return output + + @torch.no_grad() + def encode_spatial_control(self, spatial_control_path: str) -> torch.Tensor: + """ + Encode the spatial control. + """ + spatial_control = voxelize_sq_francis(spatial_control_path).to(device=self.device) # [1, 1, 64, 64, 64] + spatial_control_latent = self.models['sparse_structure_encoder'](spatial_control) # [1, 8, 16, 16, 16] + # Only for debugging: + save_voxelgrid_as_ply(spatial_control[0, 0].cpu().numpy(), Path(spatial_control_path).parent / "spatial_control_voxlized.ply") + save_voxelgrid_as_ply(spatial_control_latent[0].cpu().numpy(), Path(spatial_control_path).parent / "spatial_control_latent.ply") + return spatial_control_latent + + @torch.no_grad() + def load_mesh_high_control(self, mesh_path: str) -> torch.Tensor: + """ + Load the high-control mesh and encode it to latent space. + + Args: + mesh_path (str): The path to the high-control mesh. + + Returns: + torch.Tensor: Latent space mask. + """ + + spatial_control = voxelize_sq_francis(mesh_path).to(device=self.device) # [1, 1, 64, 64, 64] + print(f"Loaded high-control mesh from {mesh_path}, voxelized shape: {spatial_control.shape}, non-zero voxels: {(spatial_control > 0).sum().item()}") + + # Reduce resolution to 16^3 + # spatial_control = F.interpolate(spatial_control, size=(16, 16, 16), mode='trilinear', align_corners=False) # [1, 1, 16, 16, 16] + spatial_control = F.avg_pool3d(spatial_control, kernel_size=4, stride=4) # [1, 1, 16, 16, 16] + + return spatial_control + + def get_cond_text(self, prompt: List[str]) -> dict: + """ + Get the conditioning information for the model. + + Args: + prompt (List[str]): The text prompt. + + Returns: + dict: The conditioning information + """ + cond = self.encode_text(prompt) + neg_cond = self.text_cond_model['null_cond'] + return { + 'cond': cond, + 'neg_cond': neg_cond, + } + + def get_cond_image(self, image: Union[torch.Tensor, list[Image.Image]]) -> dict: + """ + Get the conditioning information for the model. + + Args: + image (Union[torch.Tensor, list[Image.Image]]): The image prompts. + + Returns: + dict: The conditioning information + """ + cond = self.encode_image(image) + neg_cond = torch.zeros_like(cond) + return { + 'cond': cond, + 'neg_cond': neg_cond, + } + @torch.no_grad() + def sample_sparse_structure( + self, + cond: dict, + num_samples: int = 1, + sampler_params: dict = {}, + vis_output_dir: str = None, + ) -> torch.Tensor: + """ + Sample sparse structures with the given conditioning. + + Args: + cond (dict): The conditioning information. + num_samples (int): The number of samples to generate. + sampler_params (dict): Additional parameters for the sampler. + vis_output_dir (str): The directory to save visualization outputs. + """ + # Sample occupancy latent + flow_model = self.models['sparse_structure_flow_model'] + reso = flow_model.resolution + noise = torch.randn(num_samples, flow_model.in_channels, reso, reso, reso).to(self.device) + sampler_params = {**self.sparse_structure_sampler_params, **sampler_params} + + sample_start = time.perf_counter() + log.info( + "Sampling TRELLIS sparse structure: samples=%d, resolution=%d, steps=%s", + num_samples, + reso, + sampler_params.get("steps", "default"), + ) + ret = self.sparse_structure_sampler.sample( + flow_model, + noise, + **cond, + **sampler_params, + verbose=True + ) + z_s = ret.samples + log.info("Sampled TRELLIS sparse structure latent in %.2fs", time.perf_counter() - sample_start) + + # Decode occupancy latent + decode_start = time.perf_counter() + decoder = self.models['sparse_structure_decoder'] + occupancy = decoder(z_s) + coords = torch.argwhere(occupancy > 0)[:, [0, 2, 3, 4]].int() + log.info( + "Decoded TRELLIS sparse structure occupancy in %.2fs (%d voxels)", + time.perf_counter() - decode_start, + coords.shape[0], + ) + + # Save intermediate voxel structure for visualization + save_voxelgrid_as_ply( + occupancy[0, 0].cpu().numpy(), "debug/structure_fm_output.ply" + ) + + if (vis_output_dir is not None) and (len(ret.pred_x_0) > 0): + video_path = os.path.join(vis_output_dir, 'denoising_evolution.mp4') + self._render_denoising_evolution(ret.pred_x_0, video_path) + + + return coords + + def decode_slat( + self, + slat: sp.SparseTensor, + formats: List[str] = ['mesh', 'gaussian', 'radiance_field'], + ) -> dict: + """ + Decode the structured latent. + + Args: + slat (sp.SparseTensor): The structured latent. + formats (List[str]): The formats to decode the structured latent to. + + Returns: + dict: The decoded structured latent. + """ + ret = {} + if 'mesh' in formats: + mesh_start = time.perf_counter() + ret['mesh'] = self.models['slat_decoder_mesh'](slat) + log.info("Decoded TRELLIS mesh SLAT in %.2fs", time.perf_counter() - mesh_start) + if 'gaussian' in formats: + gaussian_start = time.perf_counter() + ret['gaussian'] = self.models['slat_decoder_gs'](slat) + log.info("Decoded TRELLIS gaussian SLAT in %.2fs", time.perf_counter() - gaussian_start) + if 'radiance_field' in formats: + rf_start = time.perf_counter() + ret['radiance_field'] = self.models['slat_decoder_rf'](slat) + log.info("Decoded TRELLIS radiance-field SLAT in %.2fs", time.perf_counter() - rf_start) + return ret + + def sample_slat( + self, + cond: dict, + coords: torch.Tensor, + sampler_params: dict = {}, + ) -> sp.SparseTensor: + """ + Sample structured latent with the given conditioning. + + Args: + cond (dict): The conditioning information. + coords (torch.Tensor): The coordinates of the sparse structure. + sampler_params (dict): Additional parameters for the sampler. + """ + # Sample structured latent + if cond['cond'].shape[-1] == 768: + flow_model_key = 'slat_flow_model_text' + flow_model = self.models[flow_model_key] + else: + flow_model_key = 'slat_flow_model_image' + flow_model = self.models[flow_model_key] + noise = sp.SparseTensor( + feats=torch.randn(coords.shape[0], flow_model.in_channels).to(self.device), + coords=coords, + ) + sampler_params = {**self.slat_sampler_params, **sampler_params} + sample_start = time.perf_counter() + log.info( + "Sampling TRELLIS SLAT with %s: coords=%d, steps=%s", + flow_model_key, + coords.shape[0], + sampler_params.get("steps", "default"), + ) + slat = self.slat_sampler.sample( + flow_model, + noise, + **cond, + **sampler_params, + verbose=True + ).samples + log.info("Sampled TRELLIS SLAT in %.2fs", time.perf_counter() - sample_start) + + norm_start = time.perf_counter() + std = torch.tensor(self.slat_normalization['std'])[None].to(slat.device) + mean = torch.tensor(self.slat_normalization['mean'])[None].to(slat.device) + slat = slat * std + mean + log.info("Applied TRELLIS SLAT normalization in %.2fs", time.perf_counter() - norm_start) + return slat + + @torch.no_grad() + def run( + self, + prompt: str, + image: Image.Image = None, + num_samples: int = 1, + seed: int = 42, + sparse_structure_sampler_params: dict = {}, + slat_sampler_params: dict = {}, + formats: List[str] = ['mesh', 'gaussian', 'radiance_field'], + preprocess_image: bool = True, + ) -> dict: + """ + Run the pipeline. + + Args: + prompt (str): The text prompt. + num_samples (int): The number of samples to generate. + seed (int): The random seed. + sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler. + slat_sampler_params (dict): Additional parameters for the structured latent sampler. + formats (List[str]): The formats to decode the structured latent to. + """ + cond_text = self.get_cond_text([prompt]) + torch.manual_seed(seed) + spatial_control_latent = self.encode_spatial_control(sparse_structure_sampler_params['spatial_control_mesh_path']) + cond_text = {**cond_text, 'control': spatial_control_latent} + coords = self.sample_sparse_structure(cond_text, num_samples, sparse_structure_sampler_params) + cond_text.pop('control') + + if preprocess_image and image is not None: + image = self.preprocess_image(image) + cond_image = self.get_cond_image([image]) + cond = cond_image + else: + cond = cond_text + + + slat = self.sample_slat(cond, coords, slat_sampler_params) + return self.decode_slat(slat, formats) + + + @torch.no_grad() + def gen_structure( + self, + prompt: str, + num_samples: int = 1, + seed: int = 42, + sparse_structure_sampler_params: dict = {}, + slat_sampler_params: dict = {}, + ) -> dict: + """ + Run the pipeline. + + Args: + prompt (str): The text prompt. + num_samples (int): The number of samples to generate. + seed (int): The random seed. + sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler. + slat_sampler_params (dict): Additional parameters for the structured latent sampler. + """ + cond_text = self.get_cond_text([prompt]) + torch.manual_seed(seed) + spatial_control_latent = self.encode_spatial_control(sparse_structure_sampler_params['spatial_control_mesh_path']) + cond_text = {**cond_text, 'control': spatial_control_latent} + coords = self.sample_sparse_structure(cond_text, num_samples, sparse_structure_sampler_params) + + return coords + + def gen_structure_v2( + self, + prompt: str, + num_samples: int = 1, + seed: int = 42, + sparse_structure_sampler_params: dict = {}, + slat_sampler_params: dict = {}, + vis_output_dir: str = None, + ) -> dict: + """ + Run the pipeline. + + Args: + prompt (str): The text prompt. + num_samples (int): The number of samples to generate. + seed (int): The random seed. + sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler. + slat_sampler_params (dict): Additional parameters for the structured latent sampler. + """ + conditioning_start = time.perf_counter() + text_start = time.perf_counter() + cond_text = self.get_cond_text([prompt]) + log.info(f"Encoded TRELLIS text condition in {time.perf_counter() - text_start:.2f}s") + + torch.manual_seed(seed) + spatial_start = time.perf_counter() + spatial_control_latent = self.encode_spatial_control(sparse_structure_sampler_params['spatial_control_mesh_path']) + log.info(f"Encoded spatial control latent in {time.perf_counter() - spatial_start:.2f}s") + + high_control_spatial_control = None + low_control_spatial_control = None + lantent_high_control = None + if (sparse_structure_sampler_params.get('high_control_spatial_control_mesh_path', None) is not None) and (sparse_structure_sampler_params.get('local_tau_mode', None) == 'guidance'): + high_start = time.perf_counter() + high_control_spatial_control = self.encode_spatial_control(sparse_structure_sampler_params['high_control_spatial_control_mesh_path']) + log.info(f"Encoded high-control spatial latent in {time.perf_counter() - high_start:.2f}s") + elif (sparse_structure_sampler_params.get('high_control_spatial_control_mesh_path', None) is not None) and (sparse_structure_sampler_params.get('local_tau_mode', None) == 'masking'): + high_mask_start = time.perf_counter() + high_control_spatial_control = self.load_mesh_high_control(sparse_structure_sampler_params['high_control_spatial_control_mesh_path']) + log.info(f"Prepared high-control mask in {time.perf_counter() - high_mask_start:.2f}s") + # low_control_spatial_control = self.load_mesh_high_control(sparse_structure_sampler_params['spatial_control_mesh_path']) + print("High control sum:", high_control_spatial_control.sum().item()) + # print("Low control sum:", low_control_spatial_control.sum().item()) + # low_control_spatial_control = low_control_spatial_control - high_control_spatial_control + # print("Low control after subtracting high control sum:", low_control_spatial_control.sum().item()) + high_latent_start = time.perf_counter() + lantent_high_control = self.encode_spatial_control(sparse_structure_sampler_params['high_control_spatial_control_mesh_path']) + log.info(f"Encoded high-control spatial latent in {time.perf_counter() - high_latent_start:.2f}s") + elif (sparse_structure_sampler_params.get('low_control_superquadric_mask_path', None) is not None) and (sparse_structure_sampler_params.get('local_tau_mode', None) == 'low_control_mask'): + low_mask_start = time.perf_counter() + low_control_spatial_control = self.load_mesh_high_control(sparse_structure_sampler_params['low_control_superquadric_mask_path']) + log.info(f"Prepared low-control mask in {time.perf_counter() - low_mask_start:.2f}s") + high_start = time.perf_counter() + high_control_spatial_control = self.encode_spatial_control(sparse_structure_sampler_params['high_control_spatial_control_mesh_path']) + log.info(f"Encoded high-control spatial latent in {time.perf_counter() - high_start:.2f}s") + + cond_text = {**cond_text, 'control': spatial_control_latent, 'control_high': high_control_spatial_control, 'control_low_mask': low_control_spatial_control, 'latent_high_control': lantent_high_control} + log.info(f"Prepared TRELLIS structure conditioning in {time.perf_counter() - conditioning_start:.2f}s") + coords = self.sample_sparse_structure(cond_text, num_samples, sparse_structure_sampler_params, vis_output_dir=vis_output_dir) + + return coords + + def voxelize(self, mesh: o3d.geometry.TriangleMesh) -> torch.Tensor: + """ + Voxelize a mesh. + + Args: + mesh (o3d.geometry.TriangleMesh): The mesh to voxelize. + sha256 (str): The SHA256 hash of the mesh. + output_dir (str): The output directory. + """ + vertices = np.asarray(mesh.vertices) + aabb = np.stack([vertices.min(0), vertices.max(0)]) + center = (aabb[0] + aabb[1]) / 2 + scale = (aabb[1] - aabb[0]).max() + vertices = (vertices - center) / scale + vertices = np.clip(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()]) + return torch.tensor(vertices).int().cuda() + + @torch.no_grad() + def run_variant( + self, + mesh: o3d.geometry.TriangleMesh, + prompt: str, + num_samples: int = 1, + seed: int = 42, + slat_sampler_params: dict = {}, + formats: List[str] = ['mesh', 'gaussian', 'radiance_field'], + ) -> dict: + """ + Run the pipeline for making variants of an asset. + + Args: + mesh (o3d.geometry.TriangleMesh): The base mesh. + prompt (str): The text prompt. + num_samples (int): The number of samples to generate. + seed (int): The random seed + slat_sampler_params (dict): Additional parameters for the structured latent sampler. + formats (List[str]): The formats to decode the structured latent to. + """ + cond = self.get_cond([prompt]) + coords = self.voxelize(mesh) + coords = torch.cat([ + torch.arange(num_samples).repeat_interleave(coords.shape[0], 0)[:, None].int().cuda(), + coords.repeat(num_samples, 1) + ], 1) + torch.manual_seed(seed) + slat = self.sample_slat(cond, coords, slat_sampler_params) + return self.decode_slat(slat, formats) + + @torch.no_grad() + def _render_denoising_evolution( + self, + pred_x_0_list: list, + output_path: str, + total_frames: int = 1800, + resolution: int = 512, + ) -> None: + """Render each denoising step's pred_x_0 as a voxel cloud and compile into one video.""" + decoder = self.models['sparse_structure_decoder'] + num_steps = len(pred_x_0_list) + if num_steps == 0: + return + frames_per_step = max(1, total_frames // num_steps) + + all_frames = [] + for step_idx, latent in enumerate(pred_x_0_list): + occupancy = (decoder(latent) > 0).float() # (B, 1, 64, 64, 64) + occupied_idx = torch.argwhere(occupancy[0, 0]) # (M, 3) + + if occupied_idx.shape[0] == 0: + blank = np.zeros((resolution, resolution, 3), dtype=np.uint8) + all_frames.extend([blank] * frames_per_step) + continue + + positions = (occupied_idx.float() + 0.5) / 64.0 - 0.5 # (M, 3) in [-0.5, 0.5] + n = positions.shape[0] + gs = Gaussian(aabb=[-0.5, -0.5, -0.5, 1.0, 1.0, 1.0], sh_degree=0, device='cuda') + gs.from_xyz(positions) + identity = torch.zeros(n, 4, device=self.device) + identity[:, 0] = 1.0 + gs.from_rotation(identity) + gs.from_scaling(torch.full((n, 3), 1.0 / 128, device=self.device)) + gs.from_opacity(torch.full((n, 1), 0.9, device=self.device)) + gs._features_dc = torch.zeros(n, 1, 3, device=self.device) + gs._features_rest = None + + step_frames = render_utils.render_video( + gs, num_frames=frames_per_step, resolution=resolution, r=2, fov=40, + )['color'] + all_frames.extend(step_frames) + print(f"[evolution] step {step_idx+1}/{num_steps}: {n} voxels") + + os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) + imageio.mimsave(output_path, all_frames, fps=30) + print(f"[evolution] saved {output_path} ({len(all_frames)} frames)") diff --git a/third_party/TRELLIS/trellis/renderers/__init__.py b/third_party/TRELLIS/trellis/renderers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0339355c56b8d17f72e926650d140a658452fbe9 --- /dev/null +++ b/third_party/TRELLIS/trellis/renderers/__init__.py @@ -0,0 +1,31 @@ +import importlib + +__attributes = { + 'OctreeRenderer': 'octree_renderer', + 'GaussianRenderer': 'gaussian_render', + 'MeshRenderer': 'mesh_renderer', +} + +__submodules = [] + +__all__ = list(__attributes.keys()) + __submodules + +def __getattr__(name): + if name not in globals(): + if name in __attributes: + module_name = __attributes[name] + module = importlib.import_module(f".{module_name}", __name__) + globals()[name] = getattr(module, name) + elif name in __submodules: + module = importlib.import_module(f".{name}", __name__) + globals()[name] = module + else: + raise AttributeError(f"module {__name__} has no attribute {name}") + return globals()[name] + + +# For Pylance +if __name__ == '__main__': + from .octree_renderer import OctreeRenderer + from .gaussian_render import GaussianRenderer + from .mesh_renderer import MeshRenderer \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/renderers/gaussian_render.py b/third_party/TRELLIS/trellis/renderers/gaussian_render.py new file mode 100644 index 0000000000000000000000000000000000000000..57108e3cccf6aab8e3059431557c461de46aff1a --- /dev/null +++ b/third_party/TRELLIS/trellis/renderers/gaussian_render.py @@ -0,0 +1,231 @@ +# +# Copyright (C) 2023, Inria +# GRAPHDECO research group, https://team.inria.fr/graphdeco +# All rights reserved. +# +# This software is free for non-commercial, research and evaluation use +# under the terms of the LICENSE.md file. +# +# For inquiries contact george.drettakis@inria.fr +# + +import torch +import math +from easydict import EasyDict as edict +import numpy as np +from ..representations.gaussian import Gaussian +from .sh_utils import eval_sh +import torch.nn.functional as F +from easydict import EasyDict as edict + + +def intrinsics_to_projection( + intrinsics: torch.Tensor, + near: float, + far: float, + ) -> torch.Tensor: + """ + OpenCV intrinsics to OpenGL perspective matrix + + Args: + intrinsics (torch.Tensor): [3, 3] OpenCV intrinsics matrix + near (float): near plane to clip + far (float): far plane to clip + Returns: + (torch.Tensor): [4, 4] OpenGL perspective matrix + """ + fx, fy = intrinsics[0, 0], intrinsics[1, 1] + cx, cy = intrinsics[0, 2], intrinsics[1, 2] + ret = torch.zeros((4, 4), dtype=intrinsics.dtype, device=intrinsics.device) + ret[0, 0] = 2 * fx + ret[1, 1] = 2 * fy + ret[0, 2] = 2 * cx - 1 + ret[1, 2] = - 2 * cy + 1 + ret[2, 2] = far / (far - near) + ret[2, 3] = near * far / (near - far) + ret[3, 2] = 1. + return ret + + +def render(viewpoint_camera, pc : Gaussian, pipe, bg_color : torch.Tensor, scaling_modifier = 1.0, override_color = None): + """ + Render the scene. + + Background tensor (bg_color) must be on GPU! + """ + # lazy import + if 'GaussianRasterizer' not in globals(): + from diff_gaussian_rasterization import GaussianRasterizer, GaussianRasterizationSettings + + # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + # Set up rasterization configuration + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + kernel_size = pipe.kernel_size + subpixel_offset = torch.zeros((int(viewpoint_camera.image_height), int(viewpoint_camera.image_width), 2), dtype=torch.float32, device="cuda") + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + kernel_size=kernel_size, + subpixel_offset=subpixel_offset, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + # If precomputed 3d covariance is provided, use it. If not, then it will be computed from + # scaling / rotation by the rasterizer. + scales = None + rotations = None + cov3D_precomp = None + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors + # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer. + shs = None + colors_precomp = None + if override_color is None: + if pipe.convert_SHs_python: + shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree+1)**2) + dir_pp = (pc.get_xyz - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1)) + dir_pp_normalized = dir_pp/dir_pp.norm(dim=1, keepdim=True) + sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized) + colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0) + else: + shs = pc.get_features + else: + colors_precomp = override_color + + # Rasterize visible Gaussians to image, obtain their radii (on screen). + rendered_image, radii = rasterizer( + means3D = means3D, + means2D = means2D, + shs = shs, + colors_precomp = colors_precomp, + opacities = opacity, + scales = scales, + rotations = rotations, + cov3D_precomp = cov3D_precomp + ) + + # Those Gaussians that were frustum culled or had a radius of 0 were not visible. + # They will be excluded from value updates used in the splitting criteria. + return edict({"render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter" : radii > 0, + "radii": radii}) + + +class GaussianRenderer: + """ + Renderer for the Voxel representation. + + Args: + rendering_options (dict): Rendering options. + """ + + def __init__(self, rendering_options={}) -> None: + self.pipe = edict({ + "kernel_size": 0.1, + "convert_SHs_python": False, + "compute_cov3D_python": False, + "scale_modifier": 1.0, + "debug": False + }) + self.rendering_options = edict({ + "resolution": None, + "near": None, + "far": None, + "ssaa": 1, + "bg_color": 'random', + }) + self.rendering_options.update(rendering_options) + self.bg_color = None + + def render( + self, + gausssian: Gaussian, + extrinsics: torch.Tensor, + intrinsics: torch.Tensor, + colors_overwrite: torch.Tensor = None + ) -> edict: + """ + Render the gausssian. + + Args: + gaussian : gaussianmodule + extrinsics (torch.Tensor): (4, 4) camera extrinsics + intrinsics (torch.Tensor): (3, 3) camera intrinsics + colors_overwrite (torch.Tensor): (N, 3) override color + + Returns: + edict containing: + color (torch.Tensor): (3, H, W) rendered color image + """ + resolution = self.rendering_options["resolution"] + near = self.rendering_options["near"] + far = self.rendering_options["far"] + ssaa = self.rendering_options["ssaa"] + + if self.rendering_options["bg_color"] == 'random': + self.bg_color = torch.zeros(3, dtype=torch.float32, device="cuda") + if np.random.rand() < 0.5: + self.bg_color += 1 + else: + self.bg_color = torch.tensor(self.rendering_options["bg_color"], dtype=torch.float32, device="cuda") + + view = extrinsics + perspective = intrinsics_to_projection(intrinsics, near, far) + camera = torch.inverse(view)[:3, 3] + focalx = intrinsics[0, 0] + focaly = intrinsics[1, 1] + fovx = 2 * torch.atan(0.5 / focalx) + fovy = 2 * torch.atan(0.5 / focaly) + + camera_dict = edict({ + "image_height": resolution * ssaa, + "image_width": resolution * ssaa, + "FoVx": fovx, + "FoVy": fovy, + "znear": near, + "zfar": far, + "world_view_transform": view.T.contiguous(), + "projection_matrix": perspective.T.contiguous(), + "full_proj_transform": (perspective @ view).T.contiguous(), + "camera_center": camera + }) + + # Render + render_ret = render(camera_dict, gausssian, self.pipe, self.bg_color, override_color=colors_overwrite, scaling_modifier=self.pipe.scale_modifier) + + if ssaa > 1: + render_ret.render = F.interpolate(render_ret.render[None], size=(resolution, resolution), mode='bilinear', align_corners=False, antialias=True).squeeze() + + ret = edict({ + 'color': render_ret['render'] + }) + return ret diff --git a/third_party/TRELLIS/trellis/renderers/mesh_renderer.py b/third_party/TRELLIS/trellis/renderers/mesh_renderer.py new file mode 100644 index 0000000000000000000000000000000000000000..b504fa4d140c68ef3c611669ea075000d9723a04 --- /dev/null +++ b/third_party/TRELLIS/trellis/renderers/mesh_renderer.py @@ -0,0 +1,133 @@ +import torch +import nvdiffrast.torch as dr +from easydict import EasyDict as edict +from ..representations.mesh import MeshExtractResult +import torch.nn.functional as F + + +def intrinsics_to_projection( + intrinsics: torch.Tensor, + near: float, + far: float, + ) -> torch.Tensor: + """ + OpenCV intrinsics to OpenGL perspective matrix + + Args: + intrinsics (torch.Tensor): [3, 3] OpenCV intrinsics matrix + near (float): near plane to clip + far (float): far plane to clip + Returns: + (torch.Tensor): [4, 4] OpenGL perspective matrix + """ + fx, fy = intrinsics[0, 0], intrinsics[1, 1] + cx, cy = intrinsics[0, 2], intrinsics[1, 2] + ret = torch.zeros((4, 4), dtype=intrinsics.dtype, device=intrinsics.device) + ret[0, 0] = 2 * fx + ret[1, 1] = 2 * fy + ret[0, 2] = 2 * cx - 1 + ret[1, 2] = - 2 * cy + 1 + ret[2, 2] = far / (far - near) + ret[2, 3] = near * far / (near - far) + ret[3, 2] = 1. + return ret + + +class MeshRenderer: + """ + Renderer for the Mesh representation. + + Args: + rendering_options (dict): Rendering options. + glctx (nvdiffrast.torch.RasterizeGLContext): RasterizeGLContext object for CUDA/OpenGL interop. + """ + def __init__(self, rendering_options={}, device='cuda'): + self.rendering_options = edict({ + "resolution": None, + "near": None, + "far": None, + "ssaa": 1 + }) + self.rendering_options.update(rendering_options) + self.glctx = dr.RasterizeCudaContext(device=device) + self.device=device + + def render( + self, + mesh : MeshExtractResult, + extrinsics: torch.Tensor, + intrinsics: torch.Tensor, + return_types = ["mask", "normal", "depth"] + ) -> edict: + """ + Render the mesh. + + Args: + mesh : meshmodel + extrinsics (torch.Tensor): (4, 4) camera extrinsics + intrinsics (torch.Tensor): (3, 3) camera intrinsics + return_types (list): list of return types, can be "mask", "depth", "normal_map", "normal", "color" + + Returns: + edict based on return_types containing: + color (torch.Tensor): [3, H, W] rendered color image + depth (torch.Tensor): [H, W] rendered depth image + normal (torch.Tensor): [3, H, W] rendered normal image + normal_map (torch.Tensor): [3, H, W] rendered normal map image + mask (torch.Tensor): [H, W] rendered mask image + """ + resolution = self.rendering_options["resolution"] + near = self.rendering_options["near"] + far = self.rendering_options["far"] + ssaa = self.rendering_options["ssaa"] + + if mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0: + default_img = torch.zeros((1, resolution, resolution, 3), dtype=torch.float32, device=self.device) + ret_dict = {k : default_img if k in ['normal', 'normal_map', 'color'] else default_img[..., :1] for k in return_types} + return ret_dict + + perspective = intrinsics_to_projection(intrinsics, near, far) + + RT = extrinsics.unsqueeze(0) + full_proj = (perspective @ extrinsics).unsqueeze(0) + + vertices = mesh.vertices.unsqueeze(0) + + vertices_homo = torch.cat([vertices, torch.ones_like(vertices[..., :1])], dim=-1) + vertices_camera = torch.bmm(vertices_homo, RT.transpose(-1, -2)) + vertices_clip = torch.bmm(vertices_homo, full_proj.transpose(-1, -2)) + faces_int = mesh.faces.int() + rast, _ = dr.rasterize( + self.glctx, vertices_clip, faces_int, (resolution * ssaa, resolution * ssaa)) + + out_dict = edict() + for type in return_types: + img = None + if type == "mask" : + img = dr.antialias((rast[..., -1:] > 0).float(), rast, vertices_clip, faces_int) + elif type == "depth": + img = dr.interpolate(vertices_camera[..., 2:3].contiguous(), rast, faces_int)[0] + img = dr.antialias(img, rast, vertices_clip, faces_int) + elif type == "normal" : + img = dr.interpolate( + mesh.face_normal.reshape(1, -1, 3), rast, + torch.arange(mesh.faces.shape[0] * 3, device=self.device, dtype=torch.int).reshape(-1, 3) + )[0] + img = dr.antialias(img, rast, vertices_clip, faces_int) + # normalize norm pictures + img = (img + 1) / 2 + elif type == "normal_map" : + img = dr.interpolate(mesh.vertex_attrs[:, 3:].contiguous(), rast, faces_int)[0] + img = dr.antialias(img, rast, vertices_clip, faces_int) + elif type == "color" : + img = dr.interpolate(mesh.vertex_attrs[:, :3].contiguous(), rast, faces_int)[0] + img = dr.antialias(img, rast, vertices_clip, faces_int) + + if ssaa > 1: + img = F.interpolate(img.permute(0, 3, 1, 2), (resolution, resolution), mode='bilinear', align_corners=False, antialias=True) + img = img.squeeze() + else: + img = img.permute(0, 3, 1, 2).squeeze() + out_dict[type] = img + + return out_dict diff --git a/third_party/TRELLIS/trellis/renderers/octree_renderer.py b/third_party/TRELLIS/trellis/renderers/octree_renderer.py new file mode 100644 index 0000000000000000000000000000000000000000..136069cdb0645b5759d5d17f7815612a1dfc7bea --- /dev/null +++ b/third_party/TRELLIS/trellis/renderers/octree_renderer.py @@ -0,0 +1,300 @@ +import numpy as np +import torch +import torch.nn.functional as F +import math +import cv2 +from scipy.stats import qmc +from easydict import EasyDict as edict +from ..representations.octree import DfsOctree + + +def intrinsics_to_projection( + intrinsics: torch.Tensor, + near: float, + far: float, + ) -> torch.Tensor: + """ + OpenCV intrinsics to OpenGL perspective matrix + + Args: + intrinsics (torch.Tensor): [3, 3] OpenCV intrinsics matrix + near (float): near plane to clip + far (float): far plane to clip + Returns: + (torch.Tensor): [4, 4] OpenGL perspective matrix + """ + fx, fy = intrinsics[0, 0], intrinsics[1, 1] + cx, cy = intrinsics[0, 2], intrinsics[1, 2] + ret = torch.zeros((4, 4), dtype=intrinsics.dtype, device=intrinsics.device) + ret[0, 0] = 2 * fx + ret[1, 1] = 2 * fy + ret[0, 2] = 2 * cx - 1 + ret[1, 2] = - 2 * cy + 1 + ret[2, 2] = far / (far - near) + ret[2, 3] = near * far / (near - far) + ret[3, 2] = 1. + return ret + + +def render(viewpoint_camera, octree : DfsOctree, pipe, bg_color : torch.Tensor, scaling_modifier = 1.0, used_rank = None, colors_overwrite = None, aux=None, halton_sampler=None): + """ + Render the scene. + + Background tensor (bg_color) must be on GPU! + """ + # lazy import + if 'OctreeTrivecRasterizer' not in globals(): + from diffoctreerast import OctreeVoxelRasterizer, OctreeGaussianRasterizer, OctreeTrivecRasterizer, OctreeDecoupolyRasterizer + + # Set up rasterization configuration + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = edict( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=octree.active_sh_degree, + campos=viewpoint_camera.camera_center, + with_distloss=pipe.with_distloss, + jitter=pipe.jitter, + debug=pipe.debug, + ) + + positions = octree.get_xyz + if octree.primitive == "voxel": + densities = octree.get_density + elif octree.primitive == "gaussian": + opacities = octree.get_opacity + elif octree.primitive == "trivec": + trivecs = octree.get_trivec + densities = octree.get_density + raster_settings.density_shift = octree.density_shift + elif octree.primitive == "decoupoly": + decoupolys_V, decoupolys_g = octree.get_decoupoly + densities = octree.get_density + raster_settings.density_shift = octree.density_shift + else: + raise ValueError(f"Unknown primitive {octree.primitive}") + depths = octree.get_depth + + # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors + # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer. + colors_precomp = None + shs = octree.get_features + if octree.primitive in ["voxel", "gaussian"] and colors_overwrite is not None: + colors_precomp = colors_overwrite + shs = None + + ret = edict() + + if octree.primitive == "voxel": + renderer = OctreeVoxelRasterizer(raster_settings=raster_settings) + rgb, depth, alpha, distloss = renderer( + positions = positions, + densities = densities, + shs = shs, + colors_precomp = colors_precomp, + depths = depths, + aabb = octree.aabb, + aux = aux, + ) + ret['rgb'] = rgb + ret['depth'] = depth + ret['alpha'] = alpha + ret['distloss'] = distloss + elif octree.primitive == "gaussian": + renderer = OctreeGaussianRasterizer(raster_settings=raster_settings) + rgb, depth, alpha = renderer( + positions = positions, + opacities = opacities, + shs = shs, + colors_precomp = colors_precomp, + depths = depths, + aabb = octree.aabb, + aux = aux, + ) + ret['rgb'] = rgb + ret['depth'] = depth + ret['alpha'] = alpha + elif octree.primitive == "trivec": + raster_settings.used_rank = used_rank if used_rank is not None else trivecs.shape[1] + renderer = OctreeTrivecRasterizer(raster_settings=raster_settings) + rgb, depth, alpha, percent_depth = renderer( + positions = positions, + trivecs = trivecs, + densities = densities, + shs = shs, + colors_precomp = colors_precomp, + colors_overwrite = colors_overwrite, + depths = depths, + aabb = octree.aabb, + aux = aux, + halton_sampler = halton_sampler, + ) + ret['percent_depth'] = percent_depth + ret['rgb'] = rgb + ret['depth'] = depth + ret['alpha'] = alpha + elif octree.primitive == "decoupoly": + raster_settings.used_rank = used_rank if used_rank is not None else decoupolys_V.shape[1] + renderer = OctreeDecoupolyRasterizer(raster_settings=raster_settings) + rgb, depth, alpha = renderer( + positions = positions, + decoupolys_V = decoupolys_V, + decoupolys_g = decoupolys_g, + densities = densities, + shs = shs, + colors_precomp = colors_precomp, + depths = depths, + aabb = octree.aabb, + aux = aux, + ) + ret['rgb'] = rgb + ret['depth'] = depth + ret['alpha'] = alpha + + return ret + + +class OctreeRenderer: + """ + Renderer for the Voxel representation. + + Args: + rendering_options (dict): Rendering options. + """ + + def __init__(self, rendering_options={}) -> None: + try: + import diffoctreerast + except ImportError: + print("\033[93m[WARNING] diffoctreerast is not installed. The renderer will be disabled.\033[0m") + self.unsupported = True + else: + self.unsupported = False + + self.pipe = edict({ + "with_distloss": False, + "with_aux": False, + "scale_modifier": 1.0, + "used_rank": None, + "jitter": False, + "debug": False, + }) + self.rendering_options = edict({ + "resolution": None, + "near": None, + "far": None, + "ssaa": 1, + "bg_color": 'random', + }) + self.halton_sampler = qmc.Halton(2, scramble=False) + self.rendering_options.update(rendering_options) + self.bg_color = None + + def render( + self, + octree: DfsOctree, + extrinsics: torch.Tensor, + intrinsics: torch.Tensor, + colors_overwrite: torch.Tensor = None, + ) -> edict: + """ + Render the octree. + + Args: + octree (Octree): octree + extrinsics (torch.Tensor): (4, 4) camera extrinsics + intrinsics (torch.Tensor): (3, 3) camera intrinsics + colors_overwrite (torch.Tensor): (N, 3) override color + + Returns: + edict containing: + color (torch.Tensor): (3, H, W) rendered color + depth (torch.Tensor): (H, W) rendered depth + alpha (torch.Tensor): (H, W) rendered alpha + distloss (Optional[torch.Tensor]): (H, W) rendered distance loss + percent_depth (Optional[torch.Tensor]): (H, W) rendered percent depth + aux (Optional[edict]): auxiliary tensors + """ + resolution = self.rendering_options["resolution"] + near = self.rendering_options["near"] + far = self.rendering_options["far"] + ssaa = self.rendering_options["ssaa"] + + if self.unsupported: + image = np.zeros((512, 512, 3), dtype=np.uint8) + text_bbox = cv2.getTextSize("Unsupported", cv2.FONT_HERSHEY_SIMPLEX, 2, 3)[0] + origin = (512 - text_bbox[0]) // 2, (512 - text_bbox[1]) // 2 + image = cv2.putText(image, "Unsupported", origin, cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 3, cv2.LINE_AA) + return { + 'color': torch.tensor(image, dtype=torch.float32).permute(2, 0, 1) / 255, + } + + if self.rendering_options["bg_color"] == 'random': + self.bg_color = torch.zeros(3, dtype=torch.float32, device="cuda") + if np.random.rand() < 0.5: + self.bg_color += 1 + else: + self.bg_color = torch.tensor(self.rendering_options["bg_color"], dtype=torch.float32, device="cuda") + + if self.pipe["with_aux"]: + aux = { + 'grad_color2': torch.zeros((octree.num_leaf_nodes, 3), dtype=torch.float32, requires_grad=True, device="cuda") + 0, + 'contributions': torch.zeros((octree.num_leaf_nodes, 1), dtype=torch.float32, requires_grad=True, device="cuda") + 0, + } + for k in aux.keys(): + aux[k].requires_grad_() + aux[k].retain_grad() + else: + aux = None + + view = extrinsics + perspective = intrinsics_to_projection(intrinsics, near, far) + camera = torch.inverse(view)[:3, 3] + focalx = intrinsics[0, 0] + focaly = intrinsics[1, 1] + fovx = 2 * torch.atan(0.5 / focalx) + fovy = 2 * torch.atan(0.5 / focaly) + + camera_dict = edict({ + "image_height": resolution * ssaa, + "image_width": resolution * ssaa, + "FoVx": fovx, + "FoVy": fovy, + "znear": near, + "zfar": far, + "world_view_transform": view.T.contiguous(), + "projection_matrix": perspective.T.contiguous(), + "full_proj_transform": (perspective @ view).T.contiguous(), + "camera_center": camera + }) + + # Render + render_ret = render(camera_dict, octree, self.pipe, self.bg_color, aux=aux, colors_overwrite=colors_overwrite, scaling_modifier=self.pipe.scale_modifier, used_rank=self.pipe.used_rank, halton_sampler=self.halton_sampler) + + if ssaa > 1: + render_ret.rgb = F.interpolate(render_ret.rgb[None], size=(resolution, resolution), mode='bilinear', align_corners=False, antialias=True).squeeze() + render_ret.depth = F.interpolate(render_ret.depth[None, None], size=(resolution, resolution), mode='bilinear', align_corners=False, antialias=True).squeeze() + render_ret.alpha = F.interpolate(render_ret.alpha[None, None], size=(resolution, resolution), mode='bilinear', align_corners=False, antialias=True).squeeze() + if hasattr(render_ret, 'percent_depth'): + render_ret.percent_depth = F.interpolate(render_ret.percent_depth[None, None], size=(resolution, resolution), mode='bilinear', align_corners=False, antialias=True).squeeze() + + ret = edict({ + 'color': render_ret.rgb, + 'depth': render_ret.depth, + 'alpha': render_ret.alpha, + }) + if self.pipe["with_distloss"] and 'distloss' in render_ret: + ret['distloss'] = render_ret.distloss + if self.pipe["with_aux"]: + ret['aux'] = aux + if hasattr(render_ret, 'percent_depth'): + ret['percent_depth'] = render_ret.percent_depth + return ret diff --git a/third_party/TRELLIS/trellis/renderers/sh_utils.py b/third_party/TRELLIS/trellis/renderers/sh_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..bbca7d192aa3a7edf8c5b2d24dee535eac765785 --- /dev/null +++ b/third_party/TRELLIS/trellis/renderers/sh_utils.py @@ -0,0 +1,118 @@ +# Copyright 2021 The PlenOctree Authors. +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. + +import torch + +C0 = 0.28209479177387814 +C1 = 0.4886025119029199 +C2 = [ + 1.0925484305920792, + -1.0925484305920792, + 0.31539156525252005, + -1.0925484305920792, + 0.5462742152960396 +] +C3 = [ + -0.5900435899266435, + 2.890611442640554, + -0.4570457994644658, + 0.3731763325901154, + -0.4570457994644658, + 1.445305721320277, + -0.5900435899266435 +] +C4 = [ + 2.5033429417967046, + -1.7701307697799304, + 0.9461746957575601, + -0.6690465435572892, + 0.10578554691520431, + -0.6690465435572892, + 0.47308734787878004, + -1.7701307697799304, + 0.6258357354491761, +] + + +def eval_sh(deg, sh, dirs): + """ + Evaluate spherical harmonics at unit directions + using hardcoded SH polynomials. + Works with torch/np/jnp. + ... Can be 0 or more batch dimensions. + Args: + deg: int SH deg. Currently, 0-3 supported + sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2] + dirs: jnp.ndarray unit directions [..., 3] + Returns: + [..., C] + """ + assert deg <= 4 and deg >= 0 + coeff = (deg + 1) ** 2 + assert sh.shape[-1] >= coeff + + result = C0 * sh[..., 0] + if deg > 0: + x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3] + result = (result - + C1 * y * sh[..., 1] + + C1 * z * sh[..., 2] - + C1 * x * sh[..., 3]) + + if deg > 1: + xx, yy, zz = x * x, y * y, z * z + xy, yz, xz = x * y, y * z, x * z + result = (result + + C2[0] * xy * sh[..., 4] + + C2[1] * yz * sh[..., 5] + + C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] + + C2[3] * xz * sh[..., 7] + + C2[4] * (xx - yy) * sh[..., 8]) + + if deg > 2: + result = (result + + C3[0] * y * (3 * xx - yy) * sh[..., 9] + + C3[1] * xy * z * sh[..., 10] + + C3[2] * y * (4 * zz - xx - yy)* sh[..., 11] + + C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] + + C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] + + C3[5] * z * (xx - yy) * sh[..., 14] + + C3[6] * x * (xx - 3 * yy) * sh[..., 15]) + + if deg > 3: + result = (result + C4[0] * xy * (xx - yy) * sh[..., 16] + + C4[1] * yz * (3 * xx - yy) * sh[..., 17] + + C4[2] * xy * (7 * zz - 1) * sh[..., 18] + + C4[3] * yz * (7 * zz - 3) * sh[..., 19] + + C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] + + C4[5] * xz * (7 * zz - 3) * sh[..., 21] + + C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] + + C4[7] * xz * (xx - 3 * yy) * sh[..., 23] + + C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) * sh[..., 24]) + return result + +def RGB2SH(rgb): + return (rgb - 0.5) / C0 + +def SH2RGB(sh): + return sh * C0 + 0.5 \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/representations/__init__.py b/third_party/TRELLIS/trellis/representations/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..549ffdb97e87181552e9b3e086766f873e4bfb5e --- /dev/null +++ b/third_party/TRELLIS/trellis/representations/__init__.py @@ -0,0 +1,4 @@ +from .radiance_field import Strivec +from .octree import DfsOctree as Octree +from .gaussian import Gaussian +from .mesh import MeshExtractResult diff --git a/third_party/TRELLIS/trellis/representations/gaussian/__init__.py b/third_party/TRELLIS/trellis/representations/gaussian/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e3de6e180bd732836af876d748255595be2d4d74 --- /dev/null +++ b/third_party/TRELLIS/trellis/representations/gaussian/__init__.py @@ -0,0 +1 @@ +from .gaussian_model import Gaussian \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/representations/gaussian/gaussian_model.py b/third_party/TRELLIS/trellis/representations/gaussian/gaussian_model.py new file mode 100644 index 0000000000000000000000000000000000000000..54ba16f1550e8edb1728605202cc31b6dd805d90 --- /dev/null +++ b/third_party/TRELLIS/trellis/representations/gaussian/gaussian_model.py @@ -0,0 +1,209 @@ +import torch +import numpy as np +from plyfile import PlyData, PlyElement +from .general_utils import inverse_sigmoid, strip_symmetric, build_scaling_rotation +import utils3d + + +class Gaussian: + def __init__( + self, + aabb : list, + sh_degree : int = 0, + mininum_kernel_size : float = 0.0, + scaling_bias : float = 0.01, + opacity_bias : float = 0.1, + scaling_activation : str = "exp", + device='cuda' + ): + self.init_params = { + 'aabb': aabb, + 'sh_degree': sh_degree, + 'mininum_kernel_size': mininum_kernel_size, + 'scaling_bias': scaling_bias, + 'opacity_bias': opacity_bias, + 'scaling_activation': scaling_activation, + } + + self.sh_degree = sh_degree + self.active_sh_degree = sh_degree + self.mininum_kernel_size = mininum_kernel_size + self.scaling_bias = scaling_bias + self.opacity_bias = opacity_bias + self.scaling_activation_type = scaling_activation + self.device = device + self.aabb = torch.tensor(aabb, dtype=torch.float32, device=device) + self.setup_functions() + + self._xyz = None + self._features_dc = None + self._features_rest = None + self._scaling = None + self._rotation = None + self._opacity = None + + def setup_functions(self): + def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation): + L = build_scaling_rotation(scaling_modifier * scaling, rotation) + actual_covariance = L @ L.transpose(1, 2) + symm = strip_symmetric(actual_covariance) + return symm + + if self.scaling_activation_type == "exp": + self.scaling_activation = torch.exp + self.inverse_scaling_activation = torch.log + elif self.scaling_activation_type == "softplus": + self.scaling_activation = torch.nn.functional.softplus + self.inverse_scaling_activation = lambda x: x + torch.log(-torch.expm1(-x)) + + self.covariance_activation = build_covariance_from_scaling_rotation + + self.opacity_activation = torch.sigmoid + self.inverse_opacity_activation = inverse_sigmoid + + self.rotation_activation = torch.nn.functional.normalize + + self.scale_bias = self.inverse_scaling_activation(torch.tensor(self.scaling_bias)).cuda() + self.rots_bias = torch.zeros((4)).cuda() + self.rots_bias[0] = 1 + self.opacity_bias = self.inverse_opacity_activation(torch.tensor(self.opacity_bias)).cuda() + + @property + def get_scaling(self): + scales = self.scaling_activation(self._scaling + self.scale_bias) + scales = torch.square(scales) + self.mininum_kernel_size ** 2 + scales = torch.sqrt(scales) + return scales + + @property + def get_rotation(self): + return self.rotation_activation(self._rotation + self.rots_bias[None, :]) + + @property + def get_xyz(self): + return self._xyz * self.aabb[None, 3:] + self.aabb[None, :3] + + @property + def get_features(self): + return torch.cat((self._features_dc, self._features_rest), dim=2) if self._features_rest is not None else self._features_dc + + @property + def get_opacity(self): + return self.opacity_activation(self._opacity + self.opacity_bias) + + def get_covariance(self, scaling_modifier = 1): + return self.covariance_activation(self.get_scaling, scaling_modifier, self._rotation + self.rots_bias[None, :]) + + def from_scaling(self, scales): + scales = torch.sqrt(torch.square(scales) - self.mininum_kernel_size ** 2) + self._scaling = self.inverse_scaling_activation(scales) - self.scale_bias + + def from_rotation(self, rots): + self._rotation = rots - self.rots_bias[None, :] + + def from_xyz(self, xyz): + self._xyz = (xyz - self.aabb[None, :3]) / self.aabb[None, 3:] + + def from_features(self, features): + self._features_dc = features + + def from_opacity(self, opacities): + self._opacity = self.inverse_opacity_activation(opacities) - self.opacity_bias + + def construct_list_of_attributes(self): + l = ['x', 'y', 'z', 'nx', 'ny', 'nz'] + # All channels except the 3 DC + for i in range(self._features_dc.shape[1]*self._features_dc.shape[2]): + l.append('f_dc_{}'.format(i)) + l.append('opacity') + for i in range(self._scaling.shape[1]): + l.append('scale_{}'.format(i)) + for i in range(self._rotation.shape[1]): + l.append('rot_{}'.format(i)) + return l + + def save_ply(self, path, transform=[[1, 0, 0], [0, 0, -1], [0, 1, 0]]): + xyz = self.get_xyz.detach().cpu().numpy() + normals = np.zeros_like(xyz) + f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy() + opacities = inverse_sigmoid(self.get_opacity).detach().cpu().numpy() + scale = torch.log(self.get_scaling).detach().cpu().numpy() + rotation = (self._rotation + self.rots_bias[None, :]).detach().cpu().numpy() + + if transform is not None: + transform = np.array(transform) + xyz = np.matmul(xyz, transform.T) + rotation = utils3d.numpy.quaternion_to_matrix(rotation) + rotation = np.matmul(transform, rotation) + rotation = utils3d.numpy.matrix_to_quaternion(rotation) + + dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()] + + elements = np.empty(xyz.shape[0], dtype=dtype_full) + attributes = np.concatenate((xyz, normals, f_dc, opacities, scale, rotation), axis=1) + elements[:] = list(map(tuple, attributes)) + el = PlyElement.describe(elements, 'vertex') + PlyData([el]).write(path) + + def load_ply(self, path, transform=[[1, 0, 0], [0, 0, -1], [0, 1, 0]]): + plydata = PlyData.read(path) + + xyz = np.stack((np.asarray(plydata.elements[0]["x"]), + np.asarray(plydata.elements[0]["y"]), + np.asarray(plydata.elements[0]["z"])), axis=1) + opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis] + + features_dc = np.zeros((xyz.shape[0], 3, 1)) + features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"]) + features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"]) + features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"]) + + if self.sh_degree > 0: + extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")] + extra_f_names = sorted(extra_f_names, key = lambda x: int(x.split('_')[-1])) + assert len(extra_f_names)==3*(self.sh_degree + 1) ** 2 - 3 + features_extra = np.zeros((xyz.shape[0], len(extra_f_names))) + for idx, attr_name in enumerate(extra_f_names): + features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name]) + # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC) + features_extra = features_extra.reshape((features_extra.shape[0], 3, (self.max_sh_degree + 1) ** 2 - 1)) + + scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")] + scale_names = sorted(scale_names, key = lambda x: int(x.split('_')[-1])) + scales = np.zeros((xyz.shape[0], len(scale_names))) + for idx, attr_name in enumerate(scale_names): + scales[:, idx] = np.asarray(plydata.elements[0][attr_name]) + + rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")] + rot_names = sorted(rot_names, key = lambda x: int(x.split('_')[-1])) + rots = np.zeros((xyz.shape[0], len(rot_names))) + for idx, attr_name in enumerate(rot_names): + rots[:, idx] = np.asarray(plydata.elements[0][attr_name]) + + if transform is not None: + transform = np.array(transform) + xyz = np.matmul(xyz, transform) + rotation = utils3d.numpy.quaternion_to_matrix(rotation) + rotation = np.matmul(rotation, transform) + rotation = utils3d.numpy.matrix_to_quaternion(rotation) + + # convert to actual gaussian attributes + xyz = torch.tensor(xyz, dtype=torch.float, device=self.device) + features_dc = torch.tensor(features_dc, dtype=torch.float, device=self.device).transpose(1, 2).contiguous() + if self.sh_degree > 0: + features_extra = torch.tensor(features_extra, dtype=torch.float, device=self.device).transpose(1, 2).contiguous() + opacities = torch.sigmoid(torch.tensor(opacities, dtype=torch.float, device=self.device)) + scales = torch.exp(torch.tensor(scales, dtype=torch.float, device=self.device)) + rots = torch.tensor(rots, dtype=torch.float, device=self.device) + + # convert to _hidden attributes + self._xyz = (xyz - self.aabb[None, :3]) / self.aabb[None, 3:] + self._features_dc = features_dc + if self.sh_degree > 0: + self._features_rest = features_extra + else: + self._features_rest = None + self._opacity = self.inverse_opacity_activation(opacities) - self.opacity_bias + self._scaling = self.inverse_scaling_activation(torch.sqrt(torch.square(scales) - self.mininum_kernel_size ** 2)) - self.scale_bias + self._rotation = rots - self.rots_bias[None, :] + \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/representations/gaussian/general_utils.py b/third_party/TRELLIS/trellis/representations/gaussian/general_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..541c0825229a2d86e84460b765879f86f724a59d --- /dev/null +++ b/third_party/TRELLIS/trellis/representations/gaussian/general_utils.py @@ -0,0 +1,133 @@ +# +# Copyright (C) 2023, Inria +# GRAPHDECO research group, https://team.inria.fr/graphdeco +# All rights reserved. +# +# This software is free for non-commercial, research and evaluation use +# under the terms of the LICENSE.md file. +# +# For inquiries contact george.drettakis@inria.fr +# + +import torch +import sys +from datetime import datetime +import numpy as np +import random + +def inverse_sigmoid(x): + return torch.log(x/(1-x)) + +def PILtoTorch(pil_image, resolution): + resized_image_PIL = pil_image.resize(resolution) + resized_image = torch.from_numpy(np.array(resized_image_PIL)) / 255.0 + if len(resized_image.shape) == 3: + return resized_image.permute(2, 0, 1) + else: + return resized_image.unsqueeze(dim=-1).permute(2, 0, 1) + +def get_expon_lr_func( + lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000 +): + """ + Copied from Plenoxels + + Continuous learning rate decay function. Adapted from JaxNeRF + The returned rate is lr_init when step=0 and lr_final when step=max_steps, and + is log-linearly interpolated elsewhere (equivalent to exponential decay). + If lr_delay_steps>0 then the learning rate will be scaled by some smooth + function of lr_delay_mult, such that the initial learning rate is + lr_init*lr_delay_mult at the beginning of optimization but will be eased back + to the normal learning rate when steps>lr_delay_steps. + :param conf: config subtree 'lr' or similar + :param max_steps: int, the number of steps during optimization. + :return HoF which takes step as input + """ + + def helper(step): + if step < 0 or (lr_init == 0.0 and lr_final == 0.0): + # Disable this parameter + return 0.0 + if lr_delay_steps > 0: + # A kind of reverse cosine decay. + delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin( + 0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1) + ) + else: + delay_rate = 1.0 + t = np.clip(step / max_steps, 0, 1) + log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t) + return delay_rate * log_lerp + + return helper + +def strip_lowerdiag(L): + uncertainty = torch.zeros((L.shape[0], 6), dtype=torch.float, device="cuda") + + uncertainty[:, 0] = L[:, 0, 0] + uncertainty[:, 1] = L[:, 0, 1] + uncertainty[:, 2] = L[:, 0, 2] + uncertainty[:, 3] = L[:, 1, 1] + uncertainty[:, 4] = L[:, 1, 2] + uncertainty[:, 5] = L[:, 2, 2] + return uncertainty + +def strip_symmetric(sym): + return strip_lowerdiag(sym) + +def build_rotation(r): + norm = torch.sqrt(r[:,0]*r[:,0] + r[:,1]*r[:,1] + r[:,2]*r[:,2] + r[:,3]*r[:,3]) + + q = r / norm[:, None] + + R = torch.zeros((q.size(0), 3, 3), device='cuda') + + r = q[:, 0] + x = q[:, 1] + y = q[:, 2] + z = q[:, 3] + + R[:, 0, 0] = 1 - 2 * (y*y + z*z) + R[:, 0, 1] = 2 * (x*y - r*z) + R[:, 0, 2] = 2 * (x*z + r*y) + R[:, 1, 0] = 2 * (x*y + r*z) + R[:, 1, 1] = 1 - 2 * (x*x + z*z) + R[:, 1, 2] = 2 * (y*z - r*x) + R[:, 2, 0] = 2 * (x*z - r*y) + R[:, 2, 1] = 2 * (y*z + r*x) + R[:, 2, 2] = 1 - 2 * (x*x + y*y) + return R + +def build_scaling_rotation(s, r): + L = torch.zeros((s.shape[0], 3, 3), dtype=torch.float, device="cuda") + R = build_rotation(r) + + L[:,0,0] = s[:,0] + L[:,1,1] = s[:,1] + L[:,2,2] = s[:,2] + + L = R @ L + return L + +def safe_state(silent): + old_f = sys.stdout + class F: + def __init__(self, silent): + self.silent = silent + + def write(self, x): + if not self.silent: + if x.endswith("\n"): + old_f.write(x.replace("\n", " [{}]\n".format(str(datetime.now().strftime("%d/%m %H:%M:%S"))))) + else: + old_f.write(x) + + def flush(self): + old_f.flush() + + sys.stdout = F(silent) + + random.seed(0) + np.random.seed(0) + torch.manual_seed(0) + torch.cuda.set_device(torch.device("cuda:0")) diff --git a/third_party/TRELLIS/trellis/representations/mesh/__init__.py b/third_party/TRELLIS/trellis/representations/mesh/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..38cf35c0853d11cf09bdc228a87ee9d0b2f34b62 --- /dev/null +++ b/third_party/TRELLIS/trellis/representations/mesh/__init__.py @@ -0,0 +1 @@ +from .cube2mesh import SparseFeatures2Mesh, MeshExtractResult diff --git a/third_party/TRELLIS/trellis/representations/mesh/cube2mesh.py b/third_party/TRELLIS/trellis/representations/mesh/cube2mesh.py new file mode 100644 index 0000000000000000000000000000000000000000..44e8776fafbc21d787e2ba855e4c99bd191a0762 --- /dev/null +++ b/third_party/TRELLIS/trellis/representations/mesh/cube2mesh.py @@ -0,0 +1,143 @@ +import torch +from ...modules.sparse import SparseTensor +from easydict import EasyDict as edict +from .utils_cube import * +from .flexicubes.flexicubes import FlexiCubes + + +class MeshExtractResult: + def __init__(self, + vertices, + faces, + vertex_attrs=None, + res=64 + ): + self.vertices = vertices + self.faces = faces.long() + self.vertex_attrs = vertex_attrs + self.face_normal = self.comput_face_normals(vertices, faces) + self.res = res + self.success = (vertices.shape[0] != 0 and faces.shape[0] != 0) + + # training only + self.tsdf_v = None + self.tsdf_s = None + self.reg_loss = None + + def comput_face_normals(self, verts, faces): + i0 = faces[..., 0].long() + i1 = faces[..., 1].long() + i2 = faces[..., 2].long() + + v0 = verts[i0, :] + v1 = verts[i1, :] + v2 = verts[i2, :] + face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1) + face_normals = torch.nn.functional.normalize(face_normals, dim=1) + # print(face_normals.min(), face_normals.max(), face_normals.shape) + return face_normals[:, None, :].repeat(1, 3, 1) + + def comput_v_normals(self, verts, faces): + i0 = faces[..., 0].long() + i1 = faces[..., 1].long() + i2 = faces[..., 2].long() + + v0 = verts[i0, :] + v1 = verts[i1, :] + v2 = verts[i2, :] + face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1) + v_normals = torch.zeros_like(verts) + v_normals.scatter_add_(0, i0[..., None].repeat(1, 3), face_normals) + v_normals.scatter_add_(0, i1[..., None].repeat(1, 3), face_normals) + v_normals.scatter_add_(0, i2[..., None].repeat(1, 3), face_normals) + + v_normals = torch.nn.functional.normalize(v_normals, dim=1) + return v_normals + + +class SparseFeatures2Mesh: + def __init__(self, device="cuda", res=64, use_color=True): + ''' + a model to generate a mesh from sparse features structures using flexicube + ''' + super().__init__() + self.device=device + self.res = res + self.mesh_extractor = FlexiCubes(device=device) + self.sdf_bias = -1.0 / res + verts, cube = construct_dense_grid(self.res, self.device) + self.reg_c = cube.to(self.device) + self.reg_v = verts.to(self.device) + self.use_color = use_color + self._calc_layout() + + def _calc_layout(self): + LAYOUTS = { + 'sdf': {'shape': (8, 1), 'size': 8}, + 'deform': {'shape': (8, 3), 'size': 8 * 3}, + 'weights': {'shape': (21,), 'size': 21} + } + if self.use_color: + ''' + 6 channel color including normal map + ''' + LAYOUTS['color'] = {'shape': (8, 6,), 'size': 8 * 6} + self.layouts = edict(LAYOUTS) + start = 0 + for k, v in self.layouts.items(): + v['range'] = (start, start + v['size']) + start += v['size'] + self.feats_channels = start + + def get_layout(self, feats : torch.Tensor, name : str): + if name not in self.layouts: + return None + return feats[:, self.layouts[name]['range'][0]:self.layouts[name]['range'][1]].reshape(-1, *self.layouts[name]['shape']) + + def __call__(self, cubefeats : SparseTensor, training=False): + """ + Generates a mesh based on the specified sparse voxel structures. + Args: + cube_attrs [Nx21] : Sparse Tensor attrs about cube weights + verts_attrs [Nx10] : [0:1] SDF [1:4] deform [4:7] color [7:10] normal + Returns: + return the success tag and ni you loss, + """ + # add sdf bias to verts_attrs + coords = cubefeats.coords[:, 1:] + feats = cubefeats.feats + + sdf, deform, color, weights = [self.get_layout(feats, name) for name in ['sdf', 'deform', 'color', 'weights']] + sdf += self.sdf_bias + v_attrs = [sdf, deform, color] if self.use_color else [sdf, deform] + v_pos, v_attrs, reg_loss = sparse_cube2verts(coords, torch.cat(v_attrs, dim=-1), training=training) + v_attrs_d = get_dense_attrs(v_pos, v_attrs, res=self.res+1, sdf_init=True) + weights_d = get_dense_attrs(coords, weights, res=self.res, sdf_init=False) + if self.use_color: + sdf_d, deform_d, colors_d = v_attrs_d[..., 0], v_attrs_d[..., 1:4], v_attrs_d[..., 4:] + else: + sdf_d, deform_d = v_attrs_d[..., 0], v_attrs_d[..., 1:4] + colors_d = None + + x_nx3 = get_defomed_verts(self.reg_v, deform_d, self.res) + + vertices, faces, L_dev, colors = self.mesh_extractor( + voxelgrid_vertices=x_nx3, + scalar_field=sdf_d, + cube_idx=self.reg_c, + resolution=self.res, + beta=weights_d[:, :12], + alpha=weights_d[:, 12:20], + gamma_f=weights_d[:, 20], + voxelgrid_colors=colors_d, + training=training) + + mesh = MeshExtractResult(vertices=vertices, faces=faces, vertex_attrs=colors, res=self.res) + if training: + if mesh.success: + reg_loss += L_dev.mean() * 0.5 + reg_loss += (weights[:,:20]).abs().mean() * 0.2 + mesh.reg_loss = reg_loss + mesh.tsdf_v = get_defomed_verts(v_pos, v_attrs[:, 1:4], self.res) + mesh.tsdf_s = v_attrs[:, 0] + return mesh diff --git a/third_party/TRELLIS/trellis/representations/mesh/flexicubes/DCO.txt b/third_party/TRELLIS/trellis/representations/mesh/flexicubes/DCO.txt new file mode 100644 index 0000000000000000000000000000000000000000..d950df0c3f68ca2457f5e81e6dc31cbf3f279835 --- /dev/null +++ b/third_party/TRELLIS/trellis/representations/mesh/flexicubes/DCO.txt @@ -0,0 +1,34 @@ +Developer Certificate of Origin +Version 1.1 + +Copyright (C) 2004, 2006 The Linux Foundation and its contributors. + +Everyone is permitted to copy and distribute verbatim copies of this +license document, but changing it is not allowed. + + +Developer's Certificate of Origin 1.1 + +By making a contribution to this project, I certify that: + +(a) The contribution was created in whole or in part by me and I + have the right to submit it under the open source license + indicated in the file; or + +(b) The contribution is based upon previous work that, to the best + of my knowledge, is covered under an appropriate open source + license and I have the right under that license to submit that + work with modifications, whether created in whole or in part + by me, under the same open source license (unless I am + permitted to submit under a different license), as indicated + in the file; or + +(c) The contribution was provided directly to me by some other + person who certified (a), (b) or (c) and I have not modified + it. + +(d) I understand and agree that this project and the contribution + are public and that a record of the contribution (including all + personal information I submit with it, including my sign-off) is + maintained indefinitely and may be redistributed consistent with + this project or the open source license(s) involved. \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/representations/mesh/flexicubes/LICENSE.txt b/third_party/TRELLIS/trellis/representations/mesh/flexicubes/LICENSE.txt new file mode 100644 index 0000000000000000000000000000000000000000..40e8f765ee25d88128e7b5cd769389c633ba86bb --- /dev/null +++ b/third_party/TRELLIS/trellis/representations/mesh/flexicubes/LICENSE.txt @@ -0,0 +1,90 @@ +Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + +NVIDIA Source Code License for FlexiCubes + + +======================================================================= + +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 research + or evaluation 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. + +======================================================================= \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/representations/mesh/flexicubes/README.md b/third_party/TRELLIS/trellis/representations/mesh/flexicubes/README.md new file mode 100644 index 0000000000000000000000000000000000000000..8f8b460651edef71c9636d62868c239defaa73ce --- /dev/null +++ b/third_party/TRELLIS/trellis/representations/mesh/flexicubes/README.md @@ -0,0 +1,110 @@ +## Flexible Isosurface Extraction for Gradient-Based Mesh Optimization (FlexiCubes)
Official PyTorch implementation + +![Teaser image]() + +FlexiCubes is a high-quality isosurface representation specifically designed for gradient-based mesh optimization with respect to geometric, visual, or even physical objectives. For more details, please refer to our [paper](https://arxiv.org/abs/2308.05371) and [project page](https://research.nvidia.com/labs/toronto-ai/flexicubes/). + +## Highlights +* [Getting started](https://github.com/nv-tlabs/FlexiCubes#getting-started) +* [Basic workflow](https://github.com/nv-tlabs/FlexiCubes#example-usage) +* [nvdiffrec: image-based reconstruction example](https://github.com/NVlabs/nvdiffrec#news) +* [GET3D: generative AI example](https://github.com/nv-tlabs/GET3D#employing-flexicubes) +* [Bibtex](https://github.com/nv-tlabs/FlexiCubes#citation) + +## Getting Started + +The core functions of FlexiCubes are now in [Kaolin](https://github.com/NVIDIAGameWorks/kaolin/) starting from v0.15.0. See installation instructions [here](https://kaolin.readthedocs.io/en/latest/notes/installation.html) and API documentations [here](https://kaolin.readthedocs.io/en/latest/modules/kaolin.non_commercial.html#kaolin.non_commercial.FlexiCubes) + +The original code of the paper is still visible in `flexicube.py`. + +## Example Usage + +### Gradient-Based Mesh Optimization +We provide examples demonstrating how to use FlexiCubes for reconstructing unknown meshes through gradient-based optimization. Specifically, starting from randomly initialized SDF, we optimize the shape towards the reference mesh by minimizing their geometric difference, measured by multiview mask and depth losses. This workflow is a simplified version of `nvdiffrec` with code largely borrowed from the [nvdiffrec GitHub](https://github.com/NVlabs/nvdiffrec). We use the same pipeline to conduct the analysis in Section 3 and the main experiments described in Section 5 of our paper. We provide a detailed tutorial in `examples/optimization.ipynb`, along with an optimization script in `examples/optimize.py` which accepts command-line arguments. + + +To run the examples, it is suggested to install the Conda environment as detailed below: +```sh +conda create -n flexicubes python=3.9 +conda activate flexicubes +conda install pytorch==1.12.0 torchvision==0.13.0 torchaudio==0.12.0 cudatoolkit=11.3 -c pytorch +pip install imageio trimesh tqdm matplotlib torch_scatter ninja +pip install git+https://github.com/NVlabs/nvdiffrast/ +pip install kaolin==0.15.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-1.12.0_cu113.html +``` + +Then download the dataset collected by [Myles et al.](https://vcg.isti.cnr.it/Publications/2014/MPZ14/) as follows. We include one shape in 'examples/data/inputmodels/block.obj' if you want to test without downloading the full dataset. + +```sh +cd examples +python download_data.py +``` + +After downloading the data, run shape optimization with the following example command: +```sh +python optimize.py --ref_mesh data/inputmodels/block.obj --out_dir out/block +``` +You can find visualization and output meshes in the `out/block`. Below, we show the initial and final shapes during optimization, with the reference shape on the right. + +block_init + +block_final + + +To further demonstrate the flexibility of our FlexiCubes representation, which can accommodates both reconstruction objectives and regularizers defined on the extracted mesh, you can add a developability regularizer (proposed by [Stein et al.](https://www.cs.cmu.edu/~kmcrane/Projects/DiscreteDevelopable/)) to the previous reconstruction pipeline to encourage fabricability from panels: +```sh +python optimize.py --ref_mesh data/inputmodels/david.obj --out_dir out/david_dev --develop_reg True --iter=1250 +``` + +### Extract mesh from known signed distance field +While not its designated use case, our function can extract a mesh from a known Signed Distance Field (SDF) without optimization. Please refer to the tutorial found in `examples/extraction.ipynb` for details. + +## Tips for using FlexiCubes +### Regularization losses: +We commonly use three regularizers in our mesh optimization pipelines, referenced in lines `L104-L106` in `examples/optimize.py`. The weights of these regularizers should be scaled according to the your application objectives. Initially, it is suggested to employ low weights because strong regularization can hinder convergence. You can incrementally increase the weights if you notice artifacts appearing in the optimized meshes. Specifically: + +* The loss function at `L104` helps to remove floaters in areas of the shape that are not supervised by the application objective, such as internal faces when using image supervision only. +* The L_dev loss at `L105` can be increased if you observe artifacts in flat areas, as illustrated in the image below. +* Generally, the L1 regularizer on flexible weights at `L106` does not have a significant impact during the optimization of a single shape. However, we found it to be effective in stabilizing training in generative pipelines such as GET3D. +Ablating L_dev + +### Resolution of voxel grid vs. tetrahedral grid: +If you are switching from our previous work, DMTet, it's important to note the difference in grid resolution when compared to FlexiCubes. In both implementations, the resolution is defined by the edge length: a grid resolution of `n` means the grid edge length is 1/n for both the voxel and tetrahedral grids. However, a tetrahedral grid with a resolution of `n` contains only `(n/2+1)³` grid vertices, in contrast to the `(n+1)³` vertices in a voxel grid. Consequently, if you are switching from DMTet to FlexiCubes while maintaining the same resolution, you will notice not only a denser output mesh but also a substantial increase in computational cost. To align the triangle count in the output meshes more closely, we recommend adopting a 4:5 resolution ratio between the voxel grid and the tetrahedral grid. For instance, in our paper, `64³` FlexiCubes generate approximately the same number of triangles as `80³` DMTet. + +## Applications +FlexiCubes is now integrated into NVIDIA applications as a drop-in replacement for DMTet. You can visit their GitHub pages to see how FlexiCubes is used in advanced photogrammetry and 3D generative pipelines. + +[Extracting Triangular 3D Models, Materials, and Lighting From Images (nvdiffrec)](https://github.com/NVlabs/nvdiffrec#news) + +[GET3D: A Generative Model of High Quality 3D Textured Shapes Learned from Images](https://github.com/nv-tlabs/GET3D#employing-flexicubes) + + + +## License +Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +This work is made available under the [Nvidia Source Code License](LICENSE.txt). + +For business inquiries, please visit our website and submit the form: [NVIDIA Research Licensing](https://www.nvidia.com/en-us/research/inquiries/). + +## Citation +```bibtex +@article{shen2023flexicubes, +author = {Shen, Tianchang and Munkberg, Jacob and Hasselgren, Jon and Yin, Kangxue and Wang, Zian + and Chen, Wenzheng and Gojcic, Zan and Fidler, Sanja and Sharp, Nicholas and Gao, Jun}, +title = {Flexible Isosurface Extraction for Gradient-Based Mesh Optimization}, +year = {2023}, +issue_date = {August 2023}, +publisher = {Association for Computing Machinery}, +address = {New York, NY, USA}, +volume = {42}, +number = {4}, +issn = {0730-0301}, +url = {https://doi.org/10.1145/3592430}, +doi = {10.1145/3592430}, +journal = {ACM Trans. Graph.}, +month = {jul}, +articleno = {37}, +numpages = {16} +} +``` diff --git a/third_party/TRELLIS/trellis/representations/mesh/flexicubes/flexicubes.py b/third_party/TRELLIS/trellis/representations/mesh/flexicubes/flexicubes.py new file mode 100644 index 0000000000000000000000000000000000000000..c539b533b398f9b2b9acf21ae8441a4477305998 --- /dev/null +++ b/third_party/TRELLIS/trellis/representations/mesh/flexicubes/flexicubes.py @@ -0,0 +1,397 @@ +# Copyright (c) 2023, 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. + +import torch +from .tables import * +try: + from kaolin.utils.testing import check_tensor +except ImportError: # kaolin is only used for this shape check; avoid the heavy CUDA dependency + def check_tensor(tensor, shape=None, dtype=None, device=None, throw=True): + ok = True + if shape is not None: + ok = tensor.dim() == len(shape) and all(s is None or s == t for s, t in zip(shape, tensor.shape)) + if ok and dtype is not None: + ok = tensor.dtype == dtype + if ok and device is not None: + ok = tensor.device == device + if not ok and throw: + raise ValueError("check_tensor failed") + return ok + +__all__ = [ + 'FlexiCubes' +] + + +class FlexiCubes: + def __init__(self, device="cuda"): + + self.device = device + self.dmc_table = torch.tensor(dmc_table, dtype=torch.long, device=device, requires_grad=False) + self.num_vd_table = torch.tensor(num_vd_table, + dtype=torch.long, device=device, requires_grad=False) + self.check_table = torch.tensor( + check_table, + dtype=torch.long, device=device, requires_grad=False) + + self.tet_table = torch.tensor(tet_table, dtype=torch.long, device=device, requires_grad=False) + self.quad_split_1 = torch.tensor([0, 1, 2, 0, 2, 3], dtype=torch.long, device=device, requires_grad=False) + self.quad_split_2 = torch.tensor([0, 1, 3, 3, 1, 2], dtype=torch.long, device=device, requires_grad=False) + self.quad_split_train = torch.tensor( + [0, 1, 1, 2, 2, 3, 3, 0], dtype=torch.long, device=device, requires_grad=False) + + self.cube_corners = torch.tensor([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [0, 0, 1], [ + 1, 0, 1], [0, 1, 1], [1, 1, 1]], dtype=torch.float, device=device) + self.cube_corners_idx = torch.pow(2, torch.arange(8, requires_grad=False)) + self.cube_edges = torch.tensor([0, 1, 1, 5, 4, 5, 0, 4, 2, 3, 3, 7, 6, 7, 2, 6, + 2, 0, 3, 1, 7, 5, 6, 4], dtype=torch.long, device=device, requires_grad=False) + + self.edge_dir_table = torch.tensor([0, 2, 0, 2, 0, 2, 0, 2, 1, 1, 1, 1], + dtype=torch.long, device=device) + self.dir_faces_table = torch.tensor([ + [[5, 4], [3, 2], [4, 5], [2, 3]], + [[5, 4], [1, 0], [4, 5], [0, 1]], + [[3, 2], [1, 0], [2, 3], [0, 1]] + ], dtype=torch.long, device=device) + self.adj_pairs = torch.tensor([0, 1, 1, 3, 3, 2, 2, 0], dtype=torch.long, device=device) + + def __call__(self, voxelgrid_vertices, scalar_field, cube_idx, resolution, qef_reg_scale=1e-3, + weight_scale=0.99, beta=None, alpha=None, gamma_f=None, voxelgrid_colors=None, training=False): + assert torch.is_tensor(voxelgrid_vertices) and \ + check_tensor(voxelgrid_vertices, (None, 3), throw=False), \ + "'voxelgrid_vertices' should be a tensor of shape (num_vertices, 3)" + num_vertices = voxelgrid_vertices.shape[0] + assert torch.is_tensor(scalar_field) and \ + check_tensor(scalar_field, (num_vertices,), throw=False), \ + "'scalar_field' should be a tensor of shape (num_vertices,)" + assert torch.is_tensor(cube_idx) and \ + check_tensor(cube_idx, (None, 8), throw=False), \ + "'cube_idx' should be a tensor of shape (num_cubes, 8)" + num_cubes = cube_idx.shape[0] + assert beta is None or ( + torch.is_tensor(beta) and + check_tensor(beta, (num_cubes, 12), throw=False) + ), "'beta' should be a tensor of shape (num_cubes, 12)" + assert alpha is None or ( + torch.is_tensor(alpha) and + check_tensor(alpha, (num_cubes, 8), throw=False) + ), "'alpha' should be a tensor of shape (num_cubes, 8)" + assert gamma_f is None or ( + torch.is_tensor(gamma_f) and + check_tensor(gamma_f, (num_cubes,), throw=False) + ), "'gamma_f' should be a tensor of shape (num_cubes,)" + + surf_cubes, occ_fx8 = self._identify_surf_cubes(scalar_field, cube_idx) + if surf_cubes.sum() == 0: + return ( + torch.zeros((0, 3), device=self.device), + torch.zeros((0, 3), dtype=torch.long, device=self.device), + torch.zeros((0), device=self.device), + torch.zeros((0, voxelgrid_colors.shape[-1]), device=self.device) if voxelgrid_colors is not None else None + ) + beta, alpha, gamma_f = self._normalize_weights( + beta, alpha, gamma_f, surf_cubes, weight_scale) + + if voxelgrid_colors is not None: + voxelgrid_colors = torch.sigmoid(voxelgrid_colors) + + case_ids = self._get_case_id(occ_fx8, surf_cubes, resolution) + + surf_edges, idx_map, edge_counts, surf_edges_mask = self._identify_surf_edges( + scalar_field, cube_idx, surf_cubes + ) + + vd, L_dev, vd_gamma, vd_idx_map, vd_color = self._compute_vd( + voxelgrid_vertices, cube_idx[surf_cubes], surf_edges, scalar_field, + case_ids, beta, alpha, gamma_f, idx_map, qef_reg_scale, voxelgrid_colors) + vertices, faces, s_edges, edge_indices, vertices_color = self._triangulate( + scalar_field, surf_edges, vd, vd_gamma, edge_counts, idx_map, + vd_idx_map, surf_edges_mask, training, vd_color) + return vertices, faces, L_dev, vertices_color + + def _compute_reg_loss(self, vd, ue, edge_group_to_vd, vd_num_edges): + """ + Regularizer L_dev as in Equation 8 + """ + dist = torch.norm(ue - torch.index_select(input=vd, index=edge_group_to_vd, dim=0), dim=-1) + mean_l2 = torch.zeros_like(vd[:, 0]) + mean_l2 = (mean_l2).index_add_(0, edge_group_to_vd, dist) / vd_num_edges.squeeze(1).float() + mad = (dist - torch.index_select(input=mean_l2, index=edge_group_to_vd, dim=0)).abs() + return mad + + def _normalize_weights(self, beta, alpha, gamma_f, surf_cubes, weight_scale): + """ + Normalizes the given weights to be non-negative. If input weights are None, it creates and returns a set of weights of ones. + """ + n_cubes = surf_cubes.shape[0] + + if beta is not None: + beta = (torch.tanh(beta) * weight_scale + 1) + else: + beta = torch.ones((n_cubes, 12), dtype=torch.float, device=self.device) + + if alpha is not None: + alpha = (torch.tanh(alpha) * weight_scale + 1) + else: + alpha = torch.ones((n_cubes, 8), dtype=torch.float, device=self.device) + + if gamma_f is not None: + gamma_f = torch.sigmoid(gamma_f) * weight_scale + (1 - weight_scale) / 2 + else: + gamma_f = torch.ones((n_cubes), dtype=torch.float, device=self.device) + + return beta[surf_cubes], alpha[surf_cubes], gamma_f[surf_cubes] + + @torch.no_grad() + def _get_case_id(self, occ_fx8, surf_cubes, res): + """ + Obtains the ID of topology cases based on cell corner occupancy. This function resolves the + ambiguity in the Dual Marching Cubes (DMC) configurations as described in Section 1.3 of the + supplementary material. It should be noted that this function assumes a regular grid. + """ + case_ids = (occ_fx8[surf_cubes] * self.cube_corners_idx.to(self.device).unsqueeze(0)).sum(-1) + + problem_config = self.check_table.to(self.device)[case_ids] + to_check = problem_config[..., 0] == 1 + problem_config = problem_config[to_check] + if not isinstance(res, (list, tuple)): + res = [res, res, res] + + # The 'problematic_configs' only contain configurations for surface cubes. Next, we construct a 3D array, + # 'problem_config_full', to store configurations for all cubes (with default config for non-surface cubes). + # This allows efficient checking on adjacent cubes. + problem_config_full = torch.zeros(list(res) + [5], device=self.device, dtype=torch.long) + vol_idx = torch.nonzero(problem_config_full[..., 0] == 0) # N, 3 + vol_idx_problem = vol_idx[surf_cubes][to_check] + problem_config_full[vol_idx_problem[..., 0], vol_idx_problem[..., 1], vol_idx_problem[..., 2]] = problem_config + vol_idx_problem_adj = vol_idx_problem + problem_config[..., 1:4] + + within_range = ( + vol_idx_problem_adj[..., 0] >= 0) & ( + vol_idx_problem_adj[..., 0] < res[0]) & ( + vol_idx_problem_adj[..., 1] >= 0) & ( + vol_idx_problem_adj[..., 1] < res[1]) & ( + vol_idx_problem_adj[..., 2] >= 0) & ( + vol_idx_problem_adj[..., 2] < res[2]) + + vol_idx_problem = vol_idx_problem[within_range] + vol_idx_problem_adj = vol_idx_problem_adj[within_range] + problem_config = problem_config[within_range] + problem_config_adj = problem_config_full[vol_idx_problem_adj[..., 0], + vol_idx_problem_adj[..., 1], vol_idx_problem_adj[..., 2]] + # If two cubes with cases C16 and C19 share an ambiguous face, both cases are inverted. + to_invert = (problem_config_adj[..., 0] == 1) + idx = torch.arange(case_ids.shape[0], device=self.device)[to_check][within_range][to_invert] + case_ids.index_put_((idx,), problem_config[to_invert][..., -1]) + return case_ids + + @torch.no_grad() + def _identify_surf_edges(self, scalar_field, cube_idx, surf_cubes): + """ + Identifies grid edges that intersect with the underlying surface by checking for opposite signs. As each edge + can be shared by multiple cubes, this function also assigns a unique index to each surface-intersecting edge + and marks the cube edges with this index. + """ + occ_n = scalar_field < 0 + all_edges = cube_idx[surf_cubes][:, self.cube_edges].reshape(-1, 2) + unique_edges, _idx_map, counts = torch.unique(all_edges, dim=0, return_inverse=True, return_counts=True) + + unique_edges = unique_edges.long() + mask_edges = occ_n[unique_edges.reshape(-1)].reshape(-1, 2).sum(-1) == 1 + + surf_edges_mask = mask_edges[_idx_map] + counts = counts[_idx_map] + + mapping = torch.ones((unique_edges.shape[0]), dtype=torch.long, device=cube_idx.device) * -1 + mapping[mask_edges] = torch.arange(mask_edges.sum(), device=cube_idx.device) + # Shaped as [number of cubes x 12 edges per cube]. This is later used to map a cube edge to the unique index + # for a surface-intersecting edge. Non-surface-intersecting edges are marked with -1. + idx_map = mapping[_idx_map] + surf_edges = unique_edges[mask_edges] + return surf_edges, idx_map, counts, surf_edges_mask + + @torch.no_grad() + def _identify_surf_cubes(self, scalar_field, cube_idx): + """ + Identifies grid cubes that intersect with the underlying surface by checking if the signs at + all corners are not identical. + """ + occ_n = scalar_field < 0 + occ_fx8 = occ_n[cube_idx.reshape(-1)].reshape(-1, 8) + _occ_sum = torch.sum(occ_fx8, -1) + surf_cubes = (_occ_sum > 0) & (_occ_sum < 8) + return surf_cubes, occ_fx8 + + def _linear_interp(self, edges_weight, edges_x): + """ + Computes the location of zero-crossings on 'edges_x' using linear interpolation with 'edges_weight'. + """ + edge_dim = edges_weight.dim() - 2 + assert edges_weight.shape[edge_dim] == 2 + edges_weight = torch.cat([torch.index_select(input=edges_weight, index=torch.tensor(1, device=self.device), dim=edge_dim), - + torch.index_select(input=edges_weight, index=torch.tensor(0, device=self.device), dim=edge_dim)] + , edge_dim) + denominator = edges_weight.sum(edge_dim) + ue = (edges_x * edges_weight).sum(edge_dim) / denominator + return ue + + def _solve_vd_QEF(self, p_bxnx3, norm_bxnx3, c_bx3, qef_reg_scale): + p_bxnx3 = p_bxnx3.reshape(-1, 7, 3) + norm_bxnx3 = norm_bxnx3.reshape(-1, 7, 3) + c_bx3 = c_bx3.reshape(-1, 3) + A = norm_bxnx3 + B = ((p_bxnx3) * norm_bxnx3).sum(-1, keepdims=True) + + A_reg = (torch.eye(3, device=p_bxnx3.device) * qef_reg_scale).unsqueeze(0).repeat(p_bxnx3.shape[0], 1, 1) + B_reg = (qef_reg_scale * c_bx3).unsqueeze(-1) + A = torch.cat([A, A_reg], 1) + B = torch.cat([B, B_reg], 1) + dual_verts = torch.linalg.lstsq(A, B).solution.squeeze(-1) + return dual_verts + + def _compute_vd(self, voxelgrid_vertices, surf_cubes_fx8, surf_edges, scalar_field, + case_ids, beta, alpha, gamma_f, idx_map, qef_reg_scale, voxelgrid_colors): + """ + Computes the location of dual vertices as described in Section 4.2 + """ + alpha_nx12x2 = torch.index_select(input=alpha, index=self.cube_edges, dim=1).reshape(-1, 12, 2) + surf_edges_x = torch.index_select(input=voxelgrid_vertices, index=surf_edges.reshape(-1), dim=0).reshape(-1, 2, 3) + surf_edges_s = torch.index_select(input=scalar_field, index=surf_edges.reshape(-1), dim=0).reshape(-1, 2, 1) + zero_crossing = self._linear_interp(surf_edges_s, surf_edges_x) + + if voxelgrid_colors is not None: + C = voxelgrid_colors.shape[-1] + surf_edges_c = torch.index_select(input=voxelgrid_colors, index=surf_edges.reshape(-1), dim=0).reshape(-1, 2, C) + + idx_map = idx_map.reshape(-1, 12) + num_vd = torch.index_select(input=self.num_vd_table, index=case_ids, dim=0) + edge_group, edge_group_to_vd, edge_group_to_cube, vd_num_edges, vd_gamma = [], [], [], [], [] + + # if color is not None: + # vd_color = [] + + total_num_vd = 0 + vd_idx_map = torch.zeros((case_ids.shape[0], 12), dtype=torch.long, device=self.device, requires_grad=False) + + for num in torch.unique(num_vd): + cur_cubes = (num_vd == num) # consider cubes with the same numbers of vd emitted (for batching) + curr_num_vd = cur_cubes.sum() * num + curr_edge_group = self.dmc_table[case_ids[cur_cubes], :num].reshape(-1, num * 7) + curr_edge_group_to_vd = torch.arange( + curr_num_vd, device=self.device).unsqueeze(-1).repeat(1, 7) + total_num_vd + total_num_vd += curr_num_vd + curr_edge_group_to_cube = torch.arange(idx_map.shape[0], device=self.device)[ + cur_cubes].unsqueeze(-1).repeat(1, num * 7).reshape_as(curr_edge_group) + + curr_mask = (curr_edge_group != -1) + edge_group.append(torch.masked_select(curr_edge_group, curr_mask)) + edge_group_to_vd.append(torch.masked_select(curr_edge_group_to_vd.reshape_as(curr_edge_group), curr_mask)) + edge_group_to_cube.append(torch.masked_select(curr_edge_group_to_cube, curr_mask)) + vd_num_edges.append(curr_mask.reshape(-1, 7).sum(-1, keepdims=True)) + vd_gamma.append(torch.masked_select(gamma_f, cur_cubes).unsqueeze(-1).repeat(1, num).reshape(-1)) + # if color is not None: + # vd_color.append(color[cur_cubes].unsqueeze(1).repeat(1, num, 1).reshape(-1, 3)) + + edge_group = torch.cat(edge_group) + edge_group_to_vd = torch.cat(edge_group_to_vd) + edge_group_to_cube = torch.cat(edge_group_to_cube) + vd_num_edges = torch.cat(vd_num_edges) + vd_gamma = torch.cat(vd_gamma) + # if color is not None: + # vd_color = torch.cat(vd_color) + # else: + # vd_color = None + + vd = torch.zeros((total_num_vd, 3), device=self.device) + beta_sum = torch.zeros((total_num_vd, 1), device=self.device) + + idx_group = torch.gather(input=idx_map.reshape(-1), dim=0, index=edge_group_to_cube * 12 + edge_group) + + x_group = torch.index_select(input=surf_edges_x, index=idx_group.reshape(-1), dim=0).reshape(-1, 2, 3) + s_group = torch.index_select(input=surf_edges_s, index=idx_group.reshape(-1), dim=0).reshape(-1, 2, 1) + + + zero_crossing_group = torch.index_select( + input=zero_crossing, index=idx_group.reshape(-1), dim=0).reshape(-1, 3) + + alpha_group = torch.index_select(input=alpha_nx12x2.reshape(-1, 2), dim=0, + index=edge_group_to_cube * 12 + edge_group).reshape(-1, 2, 1) + ue_group = self._linear_interp(s_group * alpha_group, x_group) + + beta_group = torch.gather(input=beta.reshape(-1), dim=0, + index=edge_group_to_cube * 12 + edge_group).reshape(-1, 1) + beta_sum = beta_sum.index_add_(0, index=edge_group_to_vd, source=beta_group) + vd = vd.index_add_(0, index=edge_group_to_vd, source=ue_group * beta_group) / beta_sum + + ''' + interpolate colors use the same method as dual vertices + ''' + if voxelgrid_colors is not None: + vd_color = torch.zeros((total_num_vd, C), device=self.device) + c_group = torch.index_select(input=surf_edges_c, index=idx_group.reshape(-1), dim=0).reshape(-1, 2, C) + uc_group = self._linear_interp(s_group * alpha_group, c_group) + vd_color = vd_color.index_add_(0, index=edge_group_to_vd, source=uc_group * beta_group) / beta_sum + else: + vd_color = None + + L_dev = self._compute_reg_loss(vd, zero_crossing_group, edge_group_to_vd, vd_num_edges) + + v_idx = torch.arange(vd.shape[0], device=self.device) # + total_num_vd + + vd_idx_map = (vd_idx_map.reshape(-1)).scatter(dim=0, index=edge_group_to_cube * + 12 + edge_group, src=v_idx[edge_group_to_vd]) + + return vd, L_dev, vd_gamma, vd_idx_map, vd_color + + def _triangulate(self, scalar_field, surf_edges, vd, vd_gamma, edge_counts, idx_map, vd_idx_map, surf_edges_mask, training, vd_color): + """ + Connects four neighboring dual vertices to form a quadrilateral. The quadrilaterals are then split into + triangles based on the gamma parameter, as described in Section 4.3. + """ + with torch.no_grad(): + group_mask = (edge_counts == 4) & surf_edges_mask # surface edges shared by 4 cubes. + group = idx_map.reshape(-1)[group_mask] + vd_idx = vd_idx_map[group_mask] + edge_indices, indices = torch.sort(group, stable=True) + quad_vd_idx = vd_idx[indices].reshape(-1, 4) + + # Ensure all face directions point towards the positive SDF to maintain consistent winding. + s_edges = scalar_field[surf_edges[edge_indices.reshape(-1, 4)[:, 0]].reshape(-1)].reshape(-1, 2) + flip_mask = s_edges[:, 0] > 0 + quad_vd_idx = torch.cat((quad_vd_idx[flip_mask][:, [0, 1, 3, 2]], + quad_vd_idx[~flip_mask][:, [2, 3, 1, 0]])) + + quad_gamma = torch.index_select(input=vd_gamma, index=quad_vd_idx.reshape(-1), dim=0).reshape(-1, 4) + gamma_02 = quad_gamma[:, 0] * quad_gamma[:, 2] + gamma_13 = quad_gamma[:, 1] * quad_gamma[:, 3] + if not training: + mask = (gamma_02 > gamma_13) + faces = torch.zeros((quad_gamma.shape[0], 6), dtype=torch.long, device=quad_vd_idx.device) + faces[mask] = quad_vd_idx[mask][:, self.quad_split_1] + faces[~mask] = quad_vd_idx[~mask][:, self.quad_split_2] + faces = faces.reshape(-1, 3) + else: + vd_quad = torch.index_select(input=vd, index=quad_vd_idx.reshape(-1), dim=0).reshape(-1, 4, 3) + vd_02 = (vd_quad[:, 0] + vd_quad[:, 2]) / 2 + vd_13 = (vd_quad[:, 1] + vd_quad[:, 3]) / 2 + weight_sum = (gamma_02 + gamma_13) + 1e-8 + vd_center = (vd_02 * gamma_02.unsqueeze(-1) + vd_13 * gamma_13.unsqueeze(-1)) / weight_sum.unsqueeze(-1) + + if vd_color is not None: + color_quad = torch.index_select(input=vd_color, index=quad_vd_idx.reshape(-1), dim=0).reshape(-1, 4, vd_color.shape[-1]) + color_02 = (color_quad[:, 0] + color_quad[:, 2]) / 2 + color_13 = (color_quad[:, 1] + color_quad[:, 3]) / 2 + color_center = (color_02 * gamma_02.unsqueeze(-1) + color_13 * gamma_13.unsqueeze(-1)) / weight_sum.unsqueeze(-1) + vd_color = torch.cat([vd_color, color_center]) + + + vd_center_idx = torch.arange(vd_center.shape[0], device=self.device) + vd.shape[0] + vd = torch.cat([vd, vd_center]) + faces = quad_vd_idx[:, self.quad_split_train].reshape(-1, 4, 2) + faces = torch.cat([faces, vd_center_idx.reshape(-1, 1, 1).repeat(1, 4, 1)], -1).reshape(-1, 3) + return vd, faces, s_edges, edge_indices, vd_color \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/representations/mesh/flexicubes/tables.py b/third_party/TRELLIS/trellis/representations/mesh/flexicubes/tables.py new file mode 100644 index 0000000000000000000000000000000000000000..5873e7727b5595a1e4fbc3bd10ae5be8f3d06cca --- /dev/null +++ b/third_party/TRELLIS/trellis/representations/mesh/flexicubes/tables.py @@ -0,0 +1,791 @@ +# Copyright (c) 2023, 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. +dmc_table = [ +[[-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 8, 9, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 4, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 4, 7, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 4, 5, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 4, 5, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[5, 7, 8, 9, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 5, 7, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 5, 7, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 5, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 8, 9, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 7, 8, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 4, 7, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [4, 7, 8, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 4, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 5, 9, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 8, 11, -1, -1, -1], [4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 4, 5, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 4, 5, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[5, 7, 8, 9, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 5, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 5, 7, 8, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 5, 7, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 8, 9, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 7, 8, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 4, 7, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 9, 10, -1, -1, -1], [4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 4, 7, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 5, 9, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [4, 5, 9, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 4, 5, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 4, 5, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[5, 7, 8, 9, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 5, 7, 9, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 5, 7, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 5, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 8, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 9, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[8, 9, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 7, 8, -1, -1, -1, -1], [1, 3, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 4, 7, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 9, 10, 11, -1, -1], [4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 7, 9, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 5, 9, -1, -1, -1, -1], [1, 3, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 8, 10, 11, -1, -1], [4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 4, 5, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 5, 8, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[5, 7, 8, 9, -1, -1, -1], [1, 3, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 5, 7, 9, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 5, 7, 8, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[5, 7, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 8, 9, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 6, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 4, 6, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [4, 6, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 4, 6, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 5, 9, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [4, 5, 9, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 4, 5, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 4, 5, 8, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[5, 6, 8, 9, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 5, 6, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 5, 6, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 5, 6, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 6, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 6, 7, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [2, 3, 6, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 6, 7, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 4, 6, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 4, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [2, 3, 4, 6, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 4, 6, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 5, 9, -1, -1, -1, -1], [2, 3, 6, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 6, 7, 8, -1, -1], [4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 4, 5, -1, -1, -1], [2, 3, 6, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 4, 5, 6, 7, 8], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 5, 6, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 5, 6, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 2, 3, 5, 6, 8], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 5, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 10, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 9, 10, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 8, 9, 10, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 6, 8, 11, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 4, 6, 11, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 9, 10, -1, -1, -1], [4, 6, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 4, 6, 9, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 5, 9, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [4, 5, 9, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1]], +[[0, 2, 4, 5, 10, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 4, 5, 8, 10, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[5, 6, 8, 9, 11, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 5, 6, 9, 11, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 5, 6, 8, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 5, 6, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 6, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 6, 7, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 6, 7, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[6, 7, 8, 9, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 4, 6, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 4, 6, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 4, 6, 8, 9, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 6, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 5, 9, -1, -1, -1, -1], [1, 3, 6, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 6, 7, 8, 10, -1], [4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 4, 5, 6, 7, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 5, 6, 7, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 5, 6, 8, 9, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 5, 6, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 8, 9, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 7, 8, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 4, 7, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [4, 7, 8, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 4, 7, 9, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 6, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [4, 6, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 4, 6, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 4, 6, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[6, 7, 8, 9, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 6, 7, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 6, 7, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 6, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 11, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 8, 11, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 8, 9, 11, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 7, 8, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 4, 7, 11, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [4, 7, 8, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1]], +[[1, 2, 4, 7, 9, 11, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 6, 9, 10, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 8, 11, -1, -1, -1], [4, 6, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 4, 6, 10, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 4, 6, 8, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[6, 7, 8, 9, 10, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 6, 7, 9, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 6, 7, 8, 10, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 6, 7, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 5, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [1, 2, 5, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 5, 6, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 5, 6, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 7, 8, -1, -1, -1, -1], [1, 2, 5, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 4, 7, -1, -1, -1], [1, 2, 5, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 5, 6, 9, -1, -1], [4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 4, 5, 6, 7, 9], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 4, 6, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [1, 2, 4, 6, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 4, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 4, 6, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 6, 7, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 2, 3, 6, 7, 9], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 6, 7, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 6, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 5, 6, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 5, 6, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 5, 6, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[5, 6, 8, 9, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 7, 8, -1, -1, -1, -1], [1, 3, 5, 6, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 4, 5, 6, 7, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 5, 6, 9, 11, -1], [4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 5, 6, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 4, 6, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 4, 6, 8, 9, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 4, 6, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 6, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 6, 7, 8, 9, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 6, 7, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[5, 7, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [5, 7, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [5, 7, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 8, 9, -1, -1, -1], [5, 7, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 5, 8, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 4, 5, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [4, 5, 8, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 4, 5, 9, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 7, 9, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [4, 7, 9, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 4, 7, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 4, 7, 8, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[8, 9, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 9, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 8, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 5, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 5, 7, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [2, 3, 5, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 5, 7, 8, 9, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 4, 5, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 4, 5, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [2, 3, 4, 5, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 4, 5, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 4, 7, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 4, 7, 8, 9, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 2, 3, 4, 7, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 7, 8, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 8, 9, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 2, 3, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 5, 7, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [1, 2, 5, 7, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 5, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 5, 7, 8, 9, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 4, 5, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 2, 3, 4, 5, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 4, 5, 8, 9, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 5, 9, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 4, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [1, 2, 4, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 4, 7, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 4, 7, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 2, 8, 9, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 2, 3, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 2, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 5, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 5, 7, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 5, 7, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[5, 7, 8, 9, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 4, 5, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 4, 5, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 4, 5, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 4, 7, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 4, 7, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 4, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[1, 3, 8, 9, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 1, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[0, 3, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]], +[[-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]] +] +num_vd_table = [0, 1, 1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 2, 1, 3, 1, 2, 2, +2, 1, 2, 1, 2, 1, 1, 2, 1, 1, 2, 2, 2, 1, 2, 3, 1, 1, 2, 2, 1, 1, 1, 1, 1, 1, 2, +1, 2, 1, 2, 2, 1, 1, 2, 1, 1, 1, 1, 2, 2, 2, 1, 1, 2, 1, 2, 3, 2, 2, 1, 1, 1, 1, +1, 1, 2, 1, 1, 1, 2, 1, 2, 2, 2, 1, 1, 1, 1, 1, 2, 3, 2, 2, 2, 2, 2, 1, 3, 4, 2, +2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 2, 1, 1, 2, 2, 2, 2, 2, +3, 2, 1, 2, 1, 1, 1, 1, 1, 1, 2, 2, 3, 2, 3, 2, 4, 2, 2, 2, 2, 1, 2, 1, 2, 1, 1, +2, 1, 1, 2, 2, 2, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 2, 1, 1, 1, 1, 1, +1, 2, 1, 1, 1, 2, 2, 2, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 1, 2, +1, 1, 1, 1, 2, 1, 1, 1, 1, 1, 2, 1, 1, 1, 1, 1, 2, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, +1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0] +check_table = [ +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 1, 0, 0, 194], +[1, -1, 0, 0, 193], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 0, 1, 0, 164], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 0, -1, 0, 161], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 0, 0, 1, 152], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 0, 0, 1, 145], +[1, 0, 0, 1, 144], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 0, 0, -1, 137], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 0, 1, 0, 133], +[1, 0, 1, 0, 132], +[1, 1, 0, 0, 131], +[1, 1, 0, 0, 130], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 0, 0, 1, 100], +[0, 0, 0, 0, 0], +[1, 0, 0, 1, 98], +[0, 0, 0, 0, 0], +[1, 0, 0, 1, 96], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 0, 1, 0, 88], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 0, -1, 0, 82], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 0, 1, 0, 74], +[0, 0, 0, 0, 0], +[1, 0, 1, 0, 72], +[0, 0, 0, 0, 0], +[1, 0, 0, -1, 70], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, -1, 0, 0, 67], +[0, 0, 0, 0, 0], +[1, -1, 0, 0, 65], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 1, 0, 0, 56], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, -1, 0, 0, 52], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 1, 0, 0, 44], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 1, 0, 0, 40], +[0, 0, 0, 0, 0], +[1, 0, 0, -1, 38], +[1, 0, -1, 0, 37], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 0, -1, 0, 33], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, -1, 0, 0, 28], +[0, 0, 0, 0, 0], +[1, 0, -1, 0, 26], +[1, 0, 0, -1, 25], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, -1, 0, 0, 20], +[0, 0, 0, 0, 0], +[1, 0, -1, 0, 18], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 0, 0, -1, 9], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[1, 0, 0, -1, 6], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0], +[0, 0, 0, 0, 0] +] +tet_table = [ +[-1, -1, -1, -1, -1, -1], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[1, 1, 1, 1, 1, 1], +[4, 4, 4, 4, 4, 4], +[0, 0, 0, 0, 0, 0], +[4, 0, 0, 4, 4, -1], +[1, 1, 1, 1, 1, 1], +[4, 4, 4, 4, 4, 4], +[0, 4, 0, 4, 4, -1], +[0, 0, 0, 0, 0, 0], +[1, 1, 1, 1, 1, 1], +[5, 5, 5, 5, 5, 5], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[1, 1, 1, 1, 1, 1], +[2, 2, 2, 2, 2, 2], +[0, 0, 0, 0, 0, 0], +[2, 0, 2, -1, 0, 2], +[1, 1, 1, 1, 1, 1], +[2, -1, 2, 4, 4, 2], +[0, 0, 0, 0, 0, 0], +[2, 0, 2, 4, 4, 2], +[1, 1, 1, 1, 1, 1], +[2, 4, 2, 4, 4, 2], +[0, 4, 0, 4, 4, 0], +[2, 0, 2, 0, 0, 2], +[1, 1, 1, 1, 1, 1], +[2, 5, 2, 5, 5, 2], +[0, 0, 0, 0, 0, 0], +[2, 0, 2, 0, 0, 2], +[1, 1, 1, 1, 1, 1], +[1, 1, 1, 1, 1, 1], +[0, 1, 1, -1, 0, 1], +[0, 0, 0, 0, 0, 0], +[2, 2, 2, 2, 2, 2], +[4, 1, 1, 4, 4, 1], +[0, 1, 1, 0, 0, 1], +[4, 0, 0, 4, 4, 0], +[2, 2, 2, 2, 2, 2], +[-1, 1, 1, 4, 4, 1], +[0, 1, 1, 4, 4, 1], +[0, 0, 0, 0, 0, 0], +[2, 2, 2, 2, 2, 2], +[5, 1, 1, 5, 5, 1], +[0, 1, 1, 0, 0, 1], +[0, 0, 0, 0, 0, 0], +[2, 2, 2, 2, 2, 2], +[1, 1, 1, 1, 1, 1], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[8, 8, 8, 8, 8, 8], +[1, 1, 1, 4, 4, 1], +[0, 0, 0, 0, 0, 0], +[4, 0, 0, 4, 4, 0], +[4, 4, 4, 4, 4, 4], +[1, 1, 1, 4, 4, 1], +[0, 4, 0, 4, 4, 0], +[0, 0, 0, 0, 0, 0], +[4, 4, 4, 4, 4, 4], +[1, 1, 1, 5, 5, 1], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[5, 5, 5, 5, 5, 5], +[6, 6, 6, 6, 6, 6], +[6, -1, 0, 6, 0, 6], +[6, 0, 0, 6, 0, 6], +[6, 1, 1, 6, 1, 6], +[4, 4, 4, 4, 4, 4], +[0, 0, 0, 0, 0, 0], +[4, 0, 0, 4, 4, 4], +[1, 1, 1, 1, 1, 1], +[6, 4, -1, 6, 4, 6], +[6, 4, 0, 6, 4, 6], +[6, 0, 0, 6, 0, 6], +[6, 1, 1, 6, 1, 6], +[5, 5, 5, 5, 5, 5], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[1, 1, 1, 1, 1, 1], +[2, 2, 2, 2, 2, 2], +[0, 0, 0, 0, 0, 0], +[2, 0, 2, 2, 0, 2], +[1, 1, 1, 1, 1, 1], +[2, 2, 2, 2, 2, 2], +[0, 0, 0, 0, 0, 0], +[2, 0, 2, 2, 2, 2], +[1, 1, 1, 1, 1, 1], +[2, 4, 2, 2, 4, 2], +[0, 4, 0, 4, 4, 0], +[2, 0, 2, 2, 0, 2], +[1, 1, 1, 1, 1, 1], +[2, 2, 2, 2, 2, 2], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[1, 1, 1, 1, 1, 1], +[6, 1, 1, 6, -1, 6], +[6, 1, 1, 6, 0, 6], +[6, 0, 0, 6, 0, 6], +[6, 2, 2, 6, 2, 6], +[4, 1, 1, 4, 4, 1], +[0, 1, 1, 0, 0, 1], +[4, 0, 0, 4, 4, 4], +[2, 2, 2, 2, 2, 2], +[6, 1, 1, 6, 4, 6], +[6, 1, 1, 6, 4, 6], +[6, 0, 0, 6, 0, 6], +[6, 2, 2, 6, 2, 6], +[5, 1, 1, 5, 5, 1], +[0, 1, 1, 0, 0, 1], +[0, 0, 0, 0, 0, 0], +[2, 2, 2, 2, 2, 2], +[1, 1, 1, 1, 1, 1], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[6, 6, 6, 6, 6, 6], +[1, 1, 1, 1, 1, 1], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[4, 4, 4, 4, 4, 4], +[1, 1, 1, 1, 4, 1], +[0, 4, 0, 4, 4, 0], +[0, 0, 0, 0, 0, 0], +[4, 4, 4, 4, 4, 4], +[1, 1, 1, 1, 1, 1], +[0, 0, 0, 0, 0, 0], +[0, 5, 0, 5, 0, 5], +[5, 5, 5, 5, 5, 5], +[5, 5, 5, 5, 5, 5], +[0, 5, 0, 5, 0, 5], +[-1, 5, 0, 5, 0, 5], +[1, 5, 1, 5, 1, 5], +[4, 5, -1, 5, 4, 5], +[0, 5, 0, 5, 0, 5], +[4, 5, 0, 5, 4, 5], +[1, 5, 1, 5, 1, 5], +[4, 4, 4, 4, 4, 4], +[0, 4, 0, 4, 4, 4], +[0, 0, 0, 0, 0, 0], +[1, 1, 1, 1, 1, 1], +[6, 6, 6, 6, 6, 6], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[1, 1, 1, 1, 1, 1], +[2, 5, 2, 5, -1, 5], +[0, 5, 0, 5, 0, 5], +[2, 5, 2, 5, 0, 5], +[1, 5, 1, 5, 1, 5], +[2, 5, 2, 5, 4, 5], +[0, 5, 0, 5, 0, 5], +[2, 5, 2, 5, 4, 5], +[1, 5, 1, 5, 1, 5], +[2, 4, 2, 4, 4, 2], +[0, 4, 0, 4, 4, 4], +[2, 0, 2, 0, 0, 2], +[1, 1, 1, 1, 1, 1], +[2, 6, 2, 6, 6, 2], +[0, 0, 0, 0, 0, 0], +[2, 0, 2, 0, 0, 2], +[1, 1, 1, 1, 1, 1], +[1, 1, 1, 1, 1, 1], +[0, 1, 1, 1, 0, 1], +[0, 0, 0, 0, 0, 0], +[2, 2, 2, 2, 2, 2], +[4, 1, 1, 1, 4, 1], +[0, 1, 1, 1, 0, 1], +[4, 0, 0, 4, 4, 0], +[2, 2, 2, 2, 2, 2], +[1, 1, 1, 1, 1, 1], +[0, 1, 1, 1, 1, 1], +[0, 0, 0, 0, 0, 0], +[2, 2, 2, 2, 2, 2], +[1, 1, 1, 1, 1, 1], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[2, 2, 2, 2, 2, 2], +[1, 1, 1, 1, 1, 1], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[5, 5, 5, 5, 5, 5], +[1, 1, 1, 1, 4, 1], +[0, 0, 0, 0, 0, 0], +[4, 0, 0, 4, 4, 0], +[4, 4, 4, 4, 4, 4], +[1, 1, 1, 1, 1, 1], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[4, 4, 4, 4, 4, 4], +[1, 1, 1, 1, 1, 1], +[6, 0, 0, 6, 0, 6], +[0, 0, 0, 0, 0, 0], +[6, 6, 6, 6, 6, 6], +[5, 5, 5, 5, 5, 5], +[5, 5, 0, 5, 0, 5], +[5, 5, 0, 5, 0, 5], +[5, 5, 1, 5, 1, 5], +[4, 4, 4, 4, 4, 4], +[0, 0, 0, 0, 0, 0], +[4, 4, 0, 4, 4, 4], +[1, 1, 1, 1, 1, 1], +[4, 4, 4, 4, 4, 4], +[4, 4, 0, 4, 4, 4], +[0, 0, 0, 0, 0, 0], +[1, 1, 1, 1, 1, 1], +[8, 8, 8, 8, 8, 8], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[1, 1, 1, 1, 1, 1], +[2, 2, 2, 2, 2, 2], +[0, 0, 0, 0, 0, 0], +[2, 2, 2, 2, 0, 2], +[1, 1, 1, 1, 1, 1], +[2, 2, 2, 2, 2, 2], +[0, 0, 0, 0, 0, 0], +[2, 2, 2, 2, 2, 2], +[1, 1, 1, 1, 1, 1], +[2, 2, 2, 2, 2, 2], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[4, 1, 1, 4, 4, 1], +[2, 2, 2, 2, 2, 2], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[1, 1, 1, 1, 1, 1], +[1, 1, 1, 1, 1, 1], +[1, 1, 1, 1, 0, 1], +[0, 0, 0, 0, 0, 0], +[2, 2, 2, 2, 2, 2], +[1, 1, 1, 1, 1, 1], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[2, 4, 2, 4, 4, 2], +[1, 1, 1, 1, 1, 1], +[1, 1, 1, 1, 1, 1], +[0, 0, 0, 0, 0, 0], +[2, 2, 2, 2, 2, 2], +[1, 1, 1, 1, 1, 1], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[2, 2, 2, 2, 2, 2], +[1, 1, 1, 1, 1, 1], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[5, 5, 5, 5, 5, 5], +[1, 1, 1, 1, 1, 1], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[4, 4, 4, 4, 4, 4], +[1, 1, 1, 1, 1, 1], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[4, 4, 4, 4, 4, 4], +[1, 1, 1, 1, 1, 1], +[0, 0, 0, 0, 0, 0], +[0, 0, 0, 0, 0, 0], +[12, 12, 12, 12, 12, 12] +] diff --git a/third_party/TRELLIS/trellis/representations/mesh/utils_cube.py b/third_party/TRELLIS/trellis/representations/mesh/utils_cube.py new file mode 100644 index 0000000000000000000000000000000000000000..23913c97bb2d57dfa0384667c69f9860ea0a4155 --- /dev/null +++ b/third_party/TRELLIS/trellis/representations/mesh/utils_cube.py @@ -0,0 +1,61 @@ +import torch +cube_corners = torch.tensor([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [0, 0, 1], [ + 1, 0, 1], [0, 1, 1], [1, 1, 1]], dtype=torch.int) +cube_neighbor = torch.tensor([[1, 0, 0], [-1, 0, 0], [0, 1, 0], [0, -1, 0], [0, 0, 1], [0, 0, -1]]) +cube_edges = torch.tensor([0, 1, 1, 5, 4, 5, 0, 4, 2, 3, 3, 7, 6, 7, 2, 6, + 2, 0, 3, 1, 7, 5, 6, 4], dtype=torch.long, requires_grad=False) + +def construct_dense_grid(res, device='cuda'): + '''construct a dense grid based on resolution''' + res_v = res + 1 + vertsid = torch.arange(res_v ** 3, device=device) + coordsid = vertsid.reshape(res_v, res_v, res_v)[:res, :res, :res].flatten() + cube_corners_bias = (cube_corners[:, 0] * res_v + cube_corners[:, 1]) * res_v + cube_corners[:, 2] + cube_fx8 = (coordsid.unsqueeze(1) + cube_corners_bias.unsqueeze(0).to(device)) + verts = torch.stack([vertsid // (res_v ** 2), (vertsid // res_v) % res_v, vertsid % res_v], dim=1) + return verts, cube_fx8 + + +def construct_voxel_grid(coords): + verts = (cube_corners.unsqueeze(0).to(coords) + coords.unsqueeze(1)).reshape(-1, 3) + verts_unique, inverse_indices = torch.unique(verts, dim=0, return_inverse=True) + cubes = inverse_indices.reshape(-1, 8) + return verts_unique, cubes + + +def cubes_to_verts(num_verts, cubes, value, reduce='mean'): + """ + Args: + cubes [Vx8] verts index for each cube + value [Vx8xM] value to be scattered + Operation: + reduced[cubes[i][j]][k] += value[i][k] + """ + M = value.shape[2] # number of channels + reduced = torch.zeros(num_verts, M, device=cubes.device) + return torch.scatter_reduce(reduced, 0, + cubes.unsqueeze(-1).expand(-1, -1, M).flatten(0, 1), + value.flatten(0, 1), reduce=reduce, include_self=False) + +def sparse_cube2verts(coords, feats, training=True): + new_coords, cubes = construct_voxel_grid(coords) + new_feats = cubes_to_verts(new_coords.shape[0], cubes, feats) + if training: + con_loss = torch.mean((feats - new_feats[cubes]) ** 2) + else: + con_loss = 0.0 + return new_coords, new_feats, con_loss + + +def get_dense_attrs(coords : torch.Tensor, feats : torch.Tensor, res : int, sdf_init=True): + F = feats.shape[-1] + dense_attrs = torch.zeros([res] * 3 + [F], device=feats.device) + if sdf_init: + dense_attrs[..., 0] = 1 # initial outside sdf value + dense_attrs[coords[:, 0], coords[:, 1], coords[:, 2], :] = feats + return dense_attrs.reshape(-1, F) + + +def get_defomed_verts(v_pos : torch.Tensor, deform : torch.Tensor, res): + return v_pos / res - 0.5 + (1 - 1e-8) / (res * 2) * torch.tanh(deform) + \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/representations/octree/__init__.py b/third_party/TRELLIS/trellis/representations/octree/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f66a39a5a7498e2e99fe9d94d663796b3bc157b5 --- /dev/null +++ b/third_party/TRELLIS/trellis/representations/octree/__init__.py @@ -0,0 +1 @@ +from .octree_dfs import DfsOctree \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/representations/octree/octree_dfs.py b/third_party/TRELLIS/trellis/representations/octree/octree_dfs.py new file mode 100644 index 0000000000000000000000000000000000000000..9d1f7898f30414f304953cfb2d51d00511ec8325 --- /dev/null +++ b/third_party/TRELLIS/trellis/representations/octree/octree_dfs.py @@ -0,0 +1,362 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +DEFAULT_TRIVEC_CONFIG = { + 'dim': 8, + 'rank': 8, +} + +DEFAULT_VOXEL_CONFIG = { + 'solid': False, +} + +DEFAULT_DECOPOLY_CONFIG = { + 'degree': 8, + 'rank': 16, +} + + +class DfsOctree: + """ + Sparse Voxel Octree (SVO) implementation for PyTorch. + Using Depth-First Search (DFS) order to store the octree. + DFS order suits rendering and ray tracing. + + The structure and data are separatedly stored. + Structure is stored as a continuous array, each element is a 3*32 bits descriptor. + |-----------------------------------------| + | 0:3 bits | 4:31 bits | + | leaf num | unused | + |-----------------------------------------| + | 0:31 bits | + | child ptr | + |-----------------------------------------| + | 0:31 bits | + | data ptr | + |-----------------------------------------| + Each element represents a non-leaf node in the octree. + The valid mask is used to indicate whether the children are valid. + The leaf mask is used to indicate whether the children are leaf nodes. + The child ptr is used to point to the first non-leaf child. Non-leaf children descriptors are stored continuously from the child ptr. + The data ptr is used to point to the data of leaf children. Leaf children data are stored continuously from the data ptr. + + There are also auxiliary arrays to store the additional structural information to facilitate parallel processing. + - Position: the position of the octree nodes. + - Depth: the depth of the octree nodes. + + Args: + depth (int): the depth of the octree. + """ + + def __init__( + self, + depth, + aabb=[0,0,0,1,1,1], + sh_degree=2, + primitive='voxel', + primitive_config={}, + device='cuda', + ): + self.max_depth = depth + self.aabb = torch.tensor(aabb, dtype=torch.float32, device=device) + self.device = device + self.sh_degree = sh_degree + self.active_sh_degree = sh_degree + self.primitive = primitive + self.primitive_config = primitive_config + + self.structure = torch.tensor([[8, 1, 0]], dtype=torch.int32, device=self.device) + self.position = torch.zeros((8, 3), dtype=torch.float32, device=self.device) + self.depth = torch.zeros((8, 1), dtype=torch.uint8, device=self.device) + self.position[:, 0] = torch.tensor([0.25, 0.75, 0.25, 0.75, 0.25, 0.75, 0.25, 0.75], device=self.device) + self.position[:, 1] = torch.tensor([0.25, 0.25, 0.75, 0.75, 0.25, 0.25, 0.75, 0.75], device=self.device) + self.position[:, 2] = torch.tensor([0.25, 0.25, 0.25, 0.25, 0.75, 0.75, 0.75, 0.75], device=self.device) + self.depth[:, 0] = 1 + + self.data = ['position', 'depth'] + self.param_names = [] + + if primitive == 'voxel': + self.features_dc = torch.zeros((8, 1, 3), dtype=torch.float32, device=self.device) + self.features_ac = torch.zeros((8, (sh_degree+1)**2-1, 3), dtype=torch.float32, device=self.device) + self.data += ['features_dc', 'features_ac'] + self.param_names += ['features_dc', 'features_ac'] + if not primitive_config.get('solid', False): + self.density = torch.zeros((8, 1), dtype=torch.float32, device=self.device) + self.data.append('density') + self.param_names.append('density') + elif primitive == 'gaussian': + self.features_dc = torch.zeros((8, 1, 3), dtype=torch.float32, device=self.device) + self.features_ac = torch.zeros((8, (sh_degree+1)**2-1, 3), dtype=torch.float32, device=self.device) + self.opacity = torch.zeros((8, 1), dtype=torch.float32, device=self.device) + self.data += ['features_dc', 'features_ac', 'opacity'] + self.param_names += ['features_dc', 'features_ac', 'opacity'] + elif primitive == 'trivec': + self.trivec = torch.zeros((8, primitive_config['rank'], 3, primitive_config['dim']), dtype=torch.float32, device=self.device) + self.density = torch.zeros((8, primitive_config['rank']), dtype=torch.float32, device=self.device) + self.features_dc = torch.zeros((8, primitive_config['rank'], 1, 3), dtype=torch.float32, device=self.device) + self.features_ac = torch.zeros((8, primitive_config['rank'], (sh_degree+1)**2-1, 3), dtype=torch.float32, device=self.device) + self.density_shift = 0 + self.data += ['trivec', 'density', 'features_dc', 'features_ac'] + self.param_names += ['trivec', 'density', 'features_dc', 'features_ac'] + elif primitive == 'decoupoly': + self.decoupoly_V = torch.zeros((8, primitive_config['rank'], 3), dtype=torch.float32, device=self.device) + self.decoupoly_g = torch.zeros((8, primitive_config['rank'], primitive_config['degree']), dtype=torch.float32, device=self.device) + self.density = torch.zeros((8, primitive_config['rank']), dtype=torch.float32, device=self.device) + self.features_dc = torch.zeros((8, primitive_config['rank'], 1, 3), dtype=torch.float32, device=self.device) + self.features_ac = torch.zeros((8, primitive_config['rank'], (sh_degree+1)**2-1, 3), dtype=torch.float32, device=self.device) + self.density_shift = 0 + self.data += ['decoupoly_V', 'decoupoly_g', 'density', 'features_dc', 'features_ac'] + self.param_names += ['decoupoly_V', 'decoupoly_g', 'density', 'features_dc', 'features_ac'] + + self.setup_functions() + + def setup_functions(self): + self.density_activation = (lambda x: torch.exp(x - 2)) if self.primitive != 'trivec' else (lambda x: x) + self.opacity_activation = lambda x: torch.sigmoid(x - 6) + self.inverse_opacity_activation = lambda x: torch.log(x / (1 - x)) + 6 + self.color_activation = lambda x: torch.sigmoid(x) + + @property + def num_non_leaf_nodes(self): + return self.structure.shape[0] + + @property + def num_leaf_nodes(self): + return self.depth.shape[0] + + @property + def cur_depth(self): + return self.depth.max().item() + + @property + def occupancy(self): + return self.num_leaf_nodes / 8 ** self.cur_depth + + @property + def get_xyz(self): + return self.position + + @property + def get_depth(self): + return self.depth + + @property + def get_density(self): + if self.primitive == 'voxel' and self.voxel_config['solid']: + return torch.full((self.position.shape[0], 1), 1000, dtype=torch.float32, device=self.device) + return self.density_activation(self.density) + + @property + def get_opacity(self): + return self.opacity_activation(self.density) + + @property + def get_trivec(self): + return self.trivec + + @property + def get_decoupoly(self): + return F.normalize(self.decoupoly_V, dim=-1), self.decoupoly_g + + @property + def get_color(self): + return self.color_activation(self.colors) + + @property + def get_features(self): + if self.sh_degree == 0: + return self.features_dc + return torch.cat([self.features_dc, self.features_ac], dim=-2) + + def state_dict(self): + ret = {'structure': self.structure, 'position': self.position, 'depth': self.depth, 'sh_degree': self.sh_degree, 'active_sh_degree': self.active_sh_degree, 'trivec_config': self.trivec_config, 'voxel_config': self.voxel_config, 'primitive': self.primitive} + if hasattr(self, 'density_shift'): + ret['density_shift'] = self.density_shift + for data in set(self.data + self.param_names): + if not isinstance(getattr(self, data), nn.Module): + ret[data] = getattr(self, data) + else: + ret[data] = getattr(self, data).state_dict() + return ret + + def load_state_dict(self, state_dict): + keys = list(set(self.data + self.param_names + list(state_dict.keys()) + ['structure', 'position', 'depth'])) + for key in keys: + if key not in state_dict: + print(f"Warning: key {key} not found in the state_dict.") + continue + try: + if not isinstance(getattr(self, key), nn.Module): + setattr(self, key, state_dict[key]) + else: + getattr(self, key).load_state_dict(state_dict[key]) + except Exception as e: + print(e) + raise ValueError(f"Error loading key {key}.") + + def gather_from_leaf_children(self, data): + """ + Gather the data from the leaf children. + + Args: + data (torch.Tensor): the data to gather. The first dimension should be the number of leaf nodes. + """ + leaf_cnt = self.structure[:, 0] + leaf_cnt_masks = [leaf_cnt == i for i in range(1, 9)] + ret = torch.zeros((self.num_non_leaf_nodes,), dtype=data.dtype, device=self.device) + for i in range(8): + if leaf_cnt_masks[i].sum() == 0: + continue + start = self.structure[leaf_cnt_masks[i], 2] + for j in range(i+1): + ret[leaf_cnt_masks[i]] += data[start + j] + return ret + + def gather_from_non_leaf_children(self, data): + """ + Gather the data from the non-leaf children. + + Args: + data (torch.Tensor): the data to gather. The first dimension should be the number of leaf nodes. + """ + non_leaf_cnt = 8 - self.structure[:, 0] + non_leaf_cnt_masks = [non_leaf_cnt == i for i in range(1, 9)] + ret = torch.zeros_like(data, device=self.device) + for i in range(8): + if non_leaf_cnt_masks[i].sum() == 0: + continue + start = self.structure[non_leaf_cnt_masks[i], 1] + for j in range(i+1): + ret[non_leaf_cnt_masks[i]] += data[start + j] + return ret + + def structure_control(self, mask): + """ + Control the structure of the octree. + + Args: + mask (torch.Tensor): the mask to control the structure. 1 for subdivide, -1 for merge, 0 for keep. + """ + # Dont subdivide when the depth is the maximum. + mask[self.depth.squeeze() == self.max_depth] = torch.clamp_max(mask[self.depth.squeeze() == self.max_depth], 0) + # Dont merge when the depth is the minimum. + mask[self.depth.squeeze() == 1] = torch.clamp_min(mask[self.depth.squeeze() == 1], 0) + + # Gather control mask + structre_ctrl = self.gather_from_leaf_children(mask) + structre_ctrl[structre_ctrl==-8] = -1 + + new_leaf_num = self.structure[:, 0].clone() + # Modify the leaf num. + structre_valid = structre_ctrl >= 0 + new_leaf_num[structre_valid] -= structre_ctrl[structre_valid] # Add the new nodes. + structre_delete = structre_ctrl < 0 + merged_nodes = self.gather_from_non_leaf_children(structre_delete.int()) + new_leaf_num += merged_nodes # Delete the merged nodes. + + # Update the structure array to allocate new nodes. + mem_offset = torch.zeros((self.num_non_leaf_nodes + 1,), dtype=torch.int32, device=self.device) + mem_offset.index_add_(0, self.structure[structre_valid, 1], structre_ctrl[structre_valid]) # Add the new nodes. + mem_offset[:-1] -= structre_delete.int() # Delete the merged nodes. + new_structre_idx = torch.arange(0, self.num_non_leaf_nodes + 1, dtype=torch.int32, device=self.device) + mem_offset.cumsum(0) + new_structure_length = new_structre_idx[-1].item() + new_structre_idx = new_structre_idx[:-1] + new_structure = torch.empty((new_structure_length, 3), dtype=torch.int32, device=self.device) + new_structure[new_structre_idx[structre_valid], 0] = new_leaf_num[structre_valid] + + # Initialize the new nodes. + new_node_mask = torch.ones((new_structure_length,), dtype=torch.bool, device=self.device) + new_node_mask[new_structre_idx[structre_valid]] = False + new_structure[new_node_mask, 0] = 8 # Initialize to all leaf nodes. + new_node_num = new_node_mask.sum().item() + + # Rebuild child ptr. + non_leaf_cnt = 8 - new_structure[:, 0] + new_child_ptr = torch.cat([torch.zeros((1,), dtype=torch.int32, device=self.device), non_leaf_cnt.cumsum(0)[:-1]]) + new_structure[:, 1] = new_child_ptr + 1 + + # Rebuild data ptr with old data. + leaf_cnt = torch.zeros((new_structure_length,), dtype=torch.int32, device=self.device) + leaf_cnt.index_add_(0, new_structre_idx, self.structure[:, 0]) + old_data_ptr = torch.cat([torch.zeros((1,), dtype=torch.int32, device=self.device), leaf_cnt.cumsum(0)[:-1]]) + + # Update the data array + subdivide_mask = mask == 1 + merge_mask = mask == -1 + data_valid = ~(subdivide_mask | merge_mask) + mem_offset = torch.zeros((self.num_leaf_nodes + 1,), dtype=torch.int32, device=self.device) + mem_offset.index_add_(0, old_data_ptr[new_node_mask], torch.full((new_node_num,), 8, dtype=torch.int32, device=self.device)) # Add data array for new nodes + mem_offset[:-1] -= subdivide_mask.int() # Delete data elements for subdivide nodes + mem_offset[:-1] -= merge_mask.int() # Delete data elements for merge nodes + mem_offset.index_add_(0, self.structure[structre_valid, 2], merged_nodes[structre_valid]) # Add data elements for merge nodes + new_data_idx = torch.arange(0, self.num_leaf_nodes + 1, dtype=torch.int32, device=self.device) + mem_offset.cumsum(0) + new_data_length = new_data_idx[-1].item() + new_data_idx = new_data_idx[:-1] + new_data = {data: torch.empty((new_data_length,) + getattr(self, data).shape[1:], dtype=getattr(self, data).dtype, device=self.device) for data in self.data} + for data in self.data: + new_data[data][new_data_idx[data_valid]] = getattr(self, data)[data_valid] + + # Rebuild data ptr + leaf_cnt = new_structure[:, 0] + new_data_ptr = torch.cat([torch.zeros((1,), dtype=torch.int32, device=self.device), leaf_cnt.cumsum(0)[:-1]]) + new_structure[:, 2] = new_data_ptr + + # Initialize the new data array + ## For subdivide nodes + if subdivide_mask.sum() > 0: + subdivide_data_ptr = new_structure[new_node_mask, 2] + for data in self.data: + for i in range(8): + if data == 'position': + offset = torch.tensor([i // 4, (i // 2) % 2, i % 2], dtype=torch.float32, device=self.device) - 0.5 + scale = 2 ** (-1.0 - self.depth[subdivide_mask]) + new_data['position'][subdivide_data_ptr + i] = self.position[subdivide_mask] + offset * scale + elif data == 'depth': + new_data['depth'][subdivide_data_ptr + i] = self.depth[subdivide_mask] + 1 + elif data == 'opacity': + new_data['opacity'][subdivide_data_ptr + i] = self.inverse_opacity_activation(torch.sqrt(self.opacity_activation(self.opacity[subdivide_mask]))) + elif data == 'trivec': + offset = torch.tensor([i // 4, (i // 2) % 2, i % 2], dtype=torch.float32, device=self.device) * 0.5 + coord = (torch.linspace(0, 0.5, self.trivec.shape[-1], dtype=torch.float32, device=self.device)[None] + offset[:, None]).reshape(1, 3, self.trivec.shape[-1], 1) + axis = torch.linspace(0, 1, 3, dtype=torch.float32, device=self.device).reshape(1, 3, 1, 1).repeat(1, 1, self.trivec.shape[-1], 1) + coord = torch.stack([coord, axis], dim=3).reshape(1, 3, self.trivec.shape[-1], 2).expand(self.trivec[subdivide_mask].shape[0], -1, -1, -1) * 2 - 1 + new_data['trivec'][subdivide_data_ptr + i] = F.grid_sample(self.trivec[subdivide_mask], coord, align_corners=True) + else: + new_data[data][subdivide_data_ptr + i] = getattr(self, data)[subdivide_mask] + ## For merge nodes + if merge_mask.sum() > 0: + merge_data_ptr = torch.empty((merged_nodes.sum().item(),), dtype=torch.int32, device=self.device) + merge_nodes_cumsum = torch.cat([torch.zeros((1,), dtype=torch.int32, device=self.device), merged_nodes.cumsum(0)[:-1]]) + for i in range(8): + merge_data_ptr[merge_nodes_cumsum[merged_nodes > i] + i] = new_structure[new_structre_idx[merged_nodes > i], 2] + i + old_merge_data_ptr = self.structure[structre_delete, 2] + for data in self.data: + if data == 'position': + scale = 2 ** (1.0 - self.depth[old_merge_data_ptr]) + new_data['position'][merge_data_ptr] = ((self.position[old_merge_data_ptr] + 0.5) / scale).floor() * scale + 0.5 * scale - 0.5 + elif data == 'depth': + new_data['depth'][merge_data_ptr] = self.depth[old_merge_data_ptr] - 1 + elif data == 'opacity': + new_data['opacity'][subdivide_data_ptr + i] = self.inverse_opacity_activation(self.opacity_activation(self.opacity[subdivide_mask])**2) + elif data == 'trivec': + new_data['trivec'][merge_data_ptr] = self.trivec[old_merge_data_ptr] + else: + new_data[data][merge_data_ptr] = getattr(self, data)[old_merge_data_ptr] + + # Update the structure and data array + self.structure = new_structure + for data in self.data: + setattr(self, data, new_data[data]) + + # Save data array control temp variables + self.data_rearrange_buffer = { + 'subdivide_mask': subdivide_mask, + 'merge_mask': merge_mask, + 'data_valid': data_valid, + 'new_data_idx': new_data_idx, + 'new_data_length': new_data_length, + 'new_data': new_data + } diff --git a/third_party/TRELLIS/trellis/representations/radiance_field/__init__.py b/third_party/TRELLIS/trellis/representations/radiance_field/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b72a1b7e76b509ee5a5e6979858eb17b4158a151 --- /dev/null +++ b/third_party/TRELLIS/trellis/representations/radiance_field/__init__.py @@ -0,0 +1 @@ +from .strivec import Strivec \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/representations/radiance_field/strivec.py b/third_party/TRELLIS/trellis/representations/radiance_field/strivec.py new file mode 100644 index 0000000000000000000000000000000000000000..8fc4b749786d934dae82864b560baccd91fcabbc --- /dev/null +++ b/third_party/TRELLIS/trellis/representations/radiance_field/strivec.py @@ -0,0 +1,28 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from ..octree import DfsOctree as Octree + + +class Strivec(Octree): + def __init__( + self, + resolution: int, + aabb: list, + sh_degree: int = 0, + rank: int = 8, + dim: int = 8, + device: str = "cuda", + ): + assert np.log2(resolution) % 1 == 0, "Resolution must be a power of 2" + self.resolution = resolution + depth = int(np.round(np.log2(resolution))) + super().__init__( + depth=depth, + aabb=aabb, + sh_degree=sh_degree, + primitive="trivec", + primitive_config={"rank": rank, "dim": dim}, + device=device, + ) diff --git a/third_party/TRELLIS/trellis/utils/__init__.py b/third_party/TRELLIS/trellis/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/third_party/TRELLIS/trellis/utils/data_utils.py b/third_party/TRELLIS/trellis/utils/data_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..805b6cc118106857e7ef767ab4bfd133dbd78e6f --- /dev/null +++ b/third_party/TRELLIS/trellis/utils/data_utils.py @@ -0,0 +1,226 @@ +from typing import * +import math +import torch +import numpy as np +from torch.utils.data import Sampler, Dataset, DataLoader, DistributedSampler +import torch.distributed as dist + + +def recursive_to_device( + data: Any, + device: torch.device, + non_blocking: bool = False, +) -> Any: + """ + Recursively move all tensors in a data structure to a device. + """ + if hasattr(data, "to"): + return data.to(device, non_blocking=non_blocking) + elif isinstance(data, (list, tuple)): + return type(data)(recursive_to_device(d, device, non_blocking) for d in data) + elif isinstance(data, dict): + return {k: recursive_to_device(v, device, non_blocking) for k, v in data.items()} + else: + return data + + +def load_balanced_group_indices( + load: List[int], + num_groups: int, + equal_size: bool = False, +) -> List[List[int]]: + """ + Split indices into groups with balanced load. + """ + if equal_size: + group_size = len(load) // num_groups + indices = np.argsort(load)[::-1] + groups = [[] for _ in range(num_groups)] + group_load = np.zeros(num_groups) + for idx in indices: + min_group_idx = np.argmin(group_load) + groups[min_group_idx].append(idx) + if equal_size and len(groups[min_group_idx]) == group_size: + group_load[min_group_idx] = float('inf') + else: + group_load[min_group_idx] += load[idx] + return groups + + +def cycle(data_loader: DataLoader) -> Iterator: + while True: + for data in data_loader: + if isinstance(data_loader.sampler, ResumableSampler): + data_loader.sampler.idx += data_loader.batch_size # type: ignore[attr-defined] + yield data + if isinstance(data_loader.sampler, DistributedSampler): + data_loader.sampler.epoch += 1 + if isinstance(data_loader.sampler, ResumableSampler): + data_loader.sampler.epoch += 1 + data_loader.sampler.idx = 0 + + +class ResumableSampler(Sampler): + """ + Distributed sampler that is resumable. + + Args: + dataset: Dataset used for sampling. + rank (int, optional): Rank of the current process within :attr:`num_replicas`. + By default, :attr:`rank` is retrieved from the current distributed + group. + shuffle (bool, optional): If ``True`` (default), sampler will shuffle the + indices. + seed (int, optional): random seed used to shuffle the sampler if + :attr:`shuffle=True`. This number should be identical across all + processes in the distributed group. Default: ``0``. + drop_last (bool, optional): if ``True``, then the sampler will drop the + tail of the data to make it evenly divisible across the number of + replicas. If ``False``, the sampler will add extra indices to make + the data evenly divisible across the replicas. Default: ``False``. + """ + + def __init__( + self, + dataset: Dataset, + shuffle: bool = True, + seed: int = 0, + drop_last: bool = False, + ) -> None: + self.dataset = dataset + self.epoch = 0 + self.idx = 0 + self.drop_last = drop_last + self.world_size = dist.get_world_size() if dist.is_initialized() else 1 + self.rank = dist.get_rank() if dist.is_initialized() else 0 + # If the dataset length is evenly divisible by # of replicas, then there + # is no need to drop any data, since the dataset will be split equally. + if self.drop_last and len(self.dataset) % self.world_size != 0: # type: ignore[arg-type] + # Split to nearest available length that is evenly divisible. + # This is to ensure each rank receives the same amount of data when + # using this Sampler. + self.num_samples = math.ceil( + (len(self.dataset) - self.world_size) / self.world_size # type: ignore[arg-type] + ) + else: + self.num_samples = math.ceil(len(self.dataset) / self.world_size) # type: ignore[arg-type] + self.total_size = self.num_samples * self.world_size + self.shuffle = shuffle + self.seed = seed + + def __iter__(self) -> Iterator: + if self.shuffle: + # deterministically shuffle based on epoch and seed + g = torch.Generator() + g.manual_seed(self.seed + self.epoch) + indices = torch.randperm(len(self.dataset), generator=g).tolist() # type: ignore[arg-type] + else: + indices = list(range(len(self.dataset))) # type: ignore[arg-type] + + if not self.drop_last: + # add extra samples to make it evenly divisible + padding_size = self.total_size - len(indices) + if padding_size <= len(indices): + indices += indices[:padding_size] + else: + indices += (indices * math.ceil(padding_size / len(indices)))[ + :padding_size + ] + else: + # remove tail of data to make it evenly divisible. + indices = indices[: self.total_size] + assert len(indices) == self.total_size + + # subsample + indices = indices[self.rank : self.total_size : self.world_size] + + # resume from previous state + indices = indices[self.idx:] + + return iter(indices) + + def __len__(self) -> int: + return self.num_samples + + def state_dict(self) -> dict[str, int]: + return { + 'epoch': self.epoch, + 'idx': self.idx, + } + + def load_state_dict(self, state_dict): + self.epoch = state_dict['epoch'] + self.idx = state_dict['idx'] + + +class BalancedResumableSampler(ResumableSampler): + """ + Distributed sampler that is resumable and balances the load among the processes. + + Args: + dataset: Dataset used for sampling. + rank (int, optional): Rank of the current process within :attr:`num_replicas`. + By default, :attr:`rank` is retrieved from the current distributed + group. + shuffle (bool, optional): If ``True`` (default), sampler will shuffle the + indices. + seed (int, optional): random seed used to shuffle the sampler if + :attr:`shuffle=True`. This number should be identical across all + processes in the distributed group. Default: ``0``. + drop_last (bool, optional): if ``True``, then the sampler will drop the + tail of the data to make it evenly divisible across the number of + replicas. If ``False``, the sampler will add extra indices to make + the data evenly divisible across the replicas. Default: ``False``. + """ + + def __init__( + self, + dataset: Dataset, + shuffle: bool = True, + seed: int = 0, + drop_last: bool = False, + batch_size: int = 1, + ) -> None: + assert hasattr(dataset, 'loads'), 'Dataset must have "loads" attribute to use BalancedResumableSampler' + super().__init__(dataset, shuffle, seed, drop_last) + self.batch_size = batch_size + self.loads = dataset.loads + + def __iter__(self) -> Iterator: + if self.shuffle: + # deterministically shuffle based on epoch and seed + g = torch.Generator() + g.manual_seed(self.seed + self.epoch) + indices = torch.randperm(len(self.dataset), generator=g).tolist() # type: ignore[arg-type] + else: + indices = list(range(len(self.dataset))) # type: ignore[arg-type] + + if not self.drop_last: + # add extra samples to make it evenly divisible + padding_size = self.total_size - len(indices) + if padding_size <= len(indices): + indices += indices[:padding_size] + else: + indices += (indices * math.ceil(padding_size / len(indices)))[ + :padding_size + ] + else: + # remove tail of data to make it evenly divisible. + indices = indices[: self.total_size] + assert len(indices) == self.total_size + + # balance load among processes + num_batches = len(indices) // (self.batch_size * self.world_size) + balanced_indices = [] + for i in range(num_batches): + start_idx = i * self.batch_size * self.world_size + end_idx = (i + 1) * self.batch_size * self.world_size + batch_indices = indices[start_idx:end_idx] + batch_loads = [self.loads[idx] for idx in batch_indices] + groups = load_balanced_group_indices(batch_loads, self.world_size, equal_size=True) + balanced_indices.extend([batch_indices[j] for j in groups[self.rank]]) + + # resume from previous state + indices = balanced_indices[self.idx:] + + return iter(indices) diff --git a/third_party/TRELLIS/trellis/utils/dist_utils.py b/third_party/TRELLIS/trellis/utils/dist_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..348799c064804ed3ae1a98144249bd9b6bfb9915 --- /dev/null +++ b/third_party/TRELLIS/trellis/utils/dist_utils.py @@ -0,0 +1,93 @@ +import os +import io +from contextlib import contextmanager +import torch +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP + + +def setup_dist(rank, local_rank, world_size, master_addr, master_port): + os.environ['MASTER_ADDR'] = master_addr + os.environ['MASTER_PORT'] = master_port + os.environ['WORLD_SIZE'] = str(world_size) + os.environ['RANK'] = str(rank) + os.environ['LOCAL_RANK'] = str(local_rank) + torch.cuda.set_device(local_rank) + dist.init_process_group('nccl', rank=rank, world_size=world_size) + + +def read_file_dist(path): + """ + Read the binary file distributedly. + File is only read once by the rank 0 process and broadcasted to other processes. + + Returns: + data (io.BytesIO): The binary data read from the file. + """ + if dist.is_initialized() and dist.get_world_size() > 1: + # read file + size = torch.LongTensor(1).cuda() + if dist.get_rank() == 0: + with open(path, 'rb') as f: + data = f.read() + data = torch.ByteTensor( + torch.UntypedStorage.from_buffer(data, dtype=torch.uint8) + ).cuda() + size[0] = data.shape[0] + # broadcast size + dist.broadcast(size, src=0) + if dist.get_rank() != 0: + data = torch.ByteTensor(size[0].item()).cuda() + # broadcast data + dist.broadcast(data, src=0) + # convert to io.BytesIO + data = data.cpu().numpy().tobytes() + data = io.BytesIO(data) + return data + else: + with open(path, 'rb') as f: + data = f.read() + data = io.BytesIO(data) + return data + + +def unwrap_dist(model): + """ + Unwrap the model from distributed training. + """ + if isinstance(model, DDP): + return model.module + return model + + +@contextmanager +def master_first(): + """ + A context manager that ensures master process executes first. + """ + if not dist.is_initialized(): + yield + else: + if dist.get_rank() == 0: + yield + dist.barrier() + else: + dist.barrier() + yield + + +@contextmanager +def local_master_first(): + """ + A context manager that ensures local master process executes first. + """ + if not dist.is_initialized(): + yield + else: + if dist.get_rank() % torch.cuda.device_count() == 0: + yield + dist.barrier() + else: + dist.barrier() + yield + \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/utils/elastic_utils.py b/third_party/TRELLIS/trellis/utils/elastic_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..cba3cf83836e5b58f5bc3333e809ffc932375a04 --- /dev/null +++ b/third_party/TRELLIS/trellis/utils/elastic_utils.py @@ -0,0 +1,228 @@ +from abc import abstractmethod +from contextlib import contextmanager +from typing import Tuple +import torch +import torch.nn as nn +import numpy as np + + +class MemoryController: + """ + Base class for memory management during training. + """ + + _last_input_size = None + _last_mem_ratio = [] + + @contextmanager + def record(self): + pass + + def update_run_states(self, input_size=None, mem_ratio=None): + if self._last_input_size is None: + self._last_input_size = input_size + elif self._last_input_size!= input_size: + raise ValueError(f'Input size should not change for different ElasticModules.') + self._last_mem_ratio.append(mem_ratio) + + @abstractmethod + def get_mem_ratio(self, input_size): + pass + + @abstractmethod + def state_dict(self): + pass + + @abstractmethod + def log(self): + pass + + +class LinearMemoryController(MemoryController): + """ + A simple controller for memory management during training. + The memory usage is modeled as a linear function of: + - the number of input parameters + - the ratio of memory the model use compared to the maximum usage (with no checkpointing) + memory_usage = k * input_size * mem_ratio + b + The controller keeps track of the memory usage and gives the + expected memory ratio to keep the memory usage under a target + """ + def __init__( + self, + buffer_size=1000, + update_every=500, + target_ratio=0.8, + available_memory=None, + max_mem_ratio_start=0.1, + params=None, + device=None + ): + self.buffer_size = buffer_size + self.update_every = update_every + self.target_ratio = target_ratio + self.device = device or torch.cuda.current_device() + self.available_memory = available_memory or torch.cuda.get_device_properties(self.device).total_memory / 1024**3 + + self._memory = np.zeros(buffer_size, dtype=np.float32) + self._input_size = np.zeros(buffer_size, dtype=np.float32) + self._mem_ratio = np.zeros(buffer_size, dtype=np.float32) + self._buffer_ptr = 0 + self._buffer_length = 0 + self._params = tuple(params) if params is not None else (0.0, 0.0) + self._max_mem_ratio = max_mem_ratio_start + self.step = 0 + + def __repr__(self): + return f'LinearMemoryController(target_ratio={self.target_ratio}, available_memory={self.available_memory})' + + def _add_sample(self, memory, input_size, mem_ratio): + self._memory[self._buffer_ptr] = memory + self._input_size[self._buffer_ptr] = input_size + self._mem_ratio[self._buffer_ptr] = mem_ratio + self._buffer_ptr = (self._buffer_ptr + 1) % self.buffer_size + self._buffer_length = min(self._buffer_length + 1, self.buffer_size) + + @contextmanager + def record(self): + torch.cuda.reset_peak_memory_stats(self.device) + self._last_input_size = None + self._last_mem_ratio = [] + yield + self._last_memory = torch.cuda.max_memory_allocated(self.device) / 1024**3 + self._last_mem_ratio = sum(self._last_mem_ratio) / len(self._last_mem_ratio) + self._add_sample(self._last_memory, self._last_input_size, self._last_mem_ratio) + self.step += 1 + if self.step % self.update_every == 0: + self._max_mem_ratio = min(1.0, self._max_mem_ratio + 0.1) + self._fit_params() + + def _fit_params(self): + memory_usage = self._memory[:self._buffer_length] + input_size = self._input_size[:self._buffer_length] + mem_ratio = self._mem_ratio[:self._buffer_length] + + x = input_size * mem_ratio + y = memory_usage + k, b = np.polyfit(x, y, 1) + self._params = (k, b) + # self._visualize() + + def _visualize(self): + import matplotlib.pyplot as plt + memory_usage = self._memory[:self._buffer_length] + input_size = self._input_size[:self._buffer_length] + mem_ratio = self._mem_ratio[:self._buffer_length] + k, b = self._params + + plt.scatter(input_size * mem_ratio, memory_usage, c=mem_ratio, cmap='viridis') + x = np.array([0.0, 20000.0]) + plt.plot(x, k * x + b, c='r') + plt.savefig(f'linear_memory_controller_{self.step}.png') + plt.cla() + + def get_mem_ratio(self, input_size): + k, b = self._params + if k == 0: return np.random.rand() * self._max_mem_ratio + pred = (self.available_memory * self.target_ratio - b) / (k * input_size) + return min(self._max_mem_ratio, max(0.0, pred)) + + def state_dict(self): + return { + 'params': self._params, + } + + def load_state_dict(self, state_dict): + self._params = tuple(state_dict['params']) + + def log(self): + return { + 'params/k': self._params[0], + 'params/b': self._params[1], + 'memory': self._last_memory, + 'input_size': self._last_input_size, + 'mem_ratio': self._last_mem_ratio, + } + + +class ElasticModule(nn.Module): + """ + Module for training with elastic memory management. + """ + def __init__(self): + super().__init__() + self._memory_controller: MemoryController = None + + @abstractmethod + def _get_input_size(self, *args, **kwargs) -> int: + """ + Get the size of the input data. + + Returns: + int: The size of the input data. + """ + pass + + @abstractmethod + def _forward_with_mem_ratio(self, *args, mem_ratio=0.0, **kwargs) -> Tuple[float, Tuple]: + """ + Forward with a given memory ratio. + """ + pass + + def register_memory_controller(self, memory_controller: MemoryController): + self._memory_controller = memory_controller + + def forward(self, *args, **kwargs): + if self._memory_controller is None or not torch.is_grad_enabled() or not self.training: + _, ret = self._forward_with_mem_ratio(*args, **kwargs) + else: + input_size = self._get_input_size(*args, **kwargs) + mem_ratio = self._memory_controller.get_mem_ratio(input_size) + mem_ratio, ret = self._forward_with_mem_ratio(*args, mem_ratio=mem_ratio, **kwargs) + self._memory_controller.update_run_states(input_size, mem_ratio) + return ret + + +class ElasticModuleMixin: + """ + Mixin for training with elastic memory management. + """ + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._memory_controller: MemoryController = None + + @abstractmethod + def _get_input_size(self, *args, **kwargs) -> int: + """ + Get the size of the input data. + + Returns: + int: The size of the input data. + """ + pass + + @abstractmethod + @contextmanager + def with_mem_ratio(self, mem_ratio=1.0) -> float: + """ + Context manager for training with a reduced memory ratio compared to the full memory usage. + + Returns: + float: The exact memory ratio used during the forward pass. + """ + pass + + def register_memory_controller(self, memory_controller: MemoryController): + self._memory_controller = memory_controller + + def forward(self, *args, **kwargs): + if self._memory_controller is None or not torch.is_grad_enabled() or not self.training: + ret = super().forward(*args, **kwargs) + else: + input_size = self._get_input_size(*args, **kwargs) + mem_ratio = self._memory_controller.get_mem_ratio(input_size) + with self.with_mem_ratio(mem_ratio) as exact_mem_ratio: + ret = super().forward(*args, **kwargs) + self._memory_controller.update_run_states(input_size, exact_mem_ratio) + return ret diff --git a/third_party/TRELLIS/trellis/utils/general_utils.py b/third_party/TRELLIS/trellis/utils/general_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3b454d9c75521e33466055fe37c3fc1e37180a79 --- /dev/null +++ b/third_party/TRELLIS/trellis/utils/general_utils.py @@ -0,0 +1,187 @@ +import numpy as np +import cv2 +import torch + + +# Dictionary utils +def _dict_merge(dicta, dictb, prefix=''): + """ + Merge two dictionaries. + """ + assert isinstance(dicta, dict), 'input must be a dictionary' + assert isinstance(dictb, dict), 'input must be a dictionary' + dict_ = {} + all_keys = set(dicta.keys()).union(set(dictb.keys())) + for key in all_keys: + if key in dicta.keys() and key in dictb.keys(): + if isinstance(dicta[key], dict) and isinstance(dictb[key], dict): + dict_[key] = _dict_merge(dicta[key], dictb[key], prefix=f'{prefix}.{key}') + else: + raise ValueError(f'Duplicate key {prefix}.{key} found in both dictionaries. Types: {type(dicta[key])}, {type(dictb[key])}') + elif key in dicta.keys(): + dict_[key] = dicta[key] + else: + dict_[key] = dictb[key] + return dict_ + + +def dict_merge(dicta, dictb): + """ + Merge two dictionaries. + """ + return _dict_merge(dicta, dictb, prefix='') + + +def dict_foreach(dic, func, special_func={}): + """ + Recursively apply a function to all non-dictionary leaf values in a dictionary. + """ + assert isinstance(dic, dict), 'input must be a dictionary' + for key in dic.keys(): + if isinstance(dic[key], dict): + dic[key] = dict_foreach(dic[key], func) + else: + if key in special_func.keys(): + dic[key] = special_func[key](dic[key]) + else: + dic[key] = func(dic[key]) + return dic + + +def dict_reduce(dicts, func, special_func={}): + """ + Reduce a list of dictionaries. Leaf values must be scalars. + """ + assert isinstance(dicts, list), 'input must be a list of dictionaries' + assert all([isinstance(d, dict) for d in dicts]), 'input must be a list of dictionaries' + assert len(dicts) > 0, 'input must be a non-empty list of dictionaries' + all_keys = set([key for dict_ in dicts for key in dict_.keys()]) + reduced_dict = {} + for key in all_keys: + vlist = [dict_[key] for dict_ in dicts if key in dict_.keys()] + if isinstance(vlist[0], dict): + reduced_dict[key] = dict_reduce(vlist, func, special_func) + else: + if key in special_func.keys(): + reduced_dict[key] = special_func[key](vlist) + else: + reduced_dict[key] = func(vlist) + return reduced_dict + + +def dict_any(dic, func): + """ + Recursively apply a function to all non-dictionary leaf values in a dictionary. + """ + assert isinstance(dic, dict), 'input must be a dictionary' + for key in dic.keys(): + if isinstance(dic[key], dict): + if dict_any(dic[key], func): + return True + else: + if func(dic[key]): + return True + return False + + +def dict_all(dic, func): + """ + Recursively apply a function to all non-dictionary leaf values in a dictionary. + """ + assert isinstance(dic, dict), 'input must be a dictionary' + for key in dic.keys(): + if isinstance(dic[key], dict): + if not dict_all(dic[key], func): + return False + else: + if not func(dic[key]): + return False + return True + + +def dict_flatten(dic, sep='.'): + """ + Flatten a nested dictionary into a dictionary with no nested dictionaries. + """ + assert isinstance(dic, dict), 'input must be a dictionary' + flat_dict = {} + for key in dic.keys(): + if isinstance(dic[key], dict): + sub_dict = dict_flatten(dic[key], sep=sep) + for sub_key in sub_dict.keys(): + flat_dict[str(key) + sep + str(sub_key)] = sub_dict[sub_key] + else: + flat_dict[key] = dic[key] + return flat_dict + + +def make_grid(images, nrow=None, ncol=None, aspect_ratio=None): + num_images = len(images) + if nrow is None and ncol is None: + if aspect_ratio is not None: + nrow = int(np.round(np.sqrt(num_images / aspect_ratio))) + else: + nrow = int(np.sqrt(num_images)) + ncol = (num_images + nrow - 1) // nrow + elif nrow is None and ncol is not None: + nrow = (num_images + ncol - 1) // ncol + elif nrow is not None and ncol is None: + ncol = (num_images + nrow - 1) // nrow + else: + assert nrow * ncol >= num_images, 'nrow * ncol must be greater than or equal to the number of images' + + grid = np.zeros((nrow * images[0].shape[0], ncol * images[0].shape[1], images[0].shape[2]), dtype=images[0].dtype) + for i, img in enumerate(images): + row = i // ncol + col = i % ncol + grid[row * img.shape[0]:(row + 1) * img.shape[0], col * img.shape[1]:(col + 1) * img.shape[1]] = img + return grid + + +def notes_on_image(img, notes=None): + img = np.pad(img, ((0, 32), (0, 0), (0, 0)), 'constant', constant_values=0) + img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) + if notes is not None: + img = cv2.putText(img, notes, (0, img.shape[0] - 4), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + return img + + +def save_image_with_notes(img, path, notes=None): + """ + Save an image with notes. + """ + if isinstance(img, torch.Tensor): + img = img.cpu().numpy().transpose(1, 2, 0) + if img.dtype == np.float32 or img.dtype == np.float64: + img = np.clip(img * 255, 0, 255).astype(np.uint8) + img = notes_on_image(img, notes) + cv2.imwrite(path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR)) + + +# debug utils + +def atol(x, y): + """ + Absolute tolerance. + """ + return torch.abs(x - y) + + +def rtol(x, y): + """ + Relative tolerance. + """ + return torch.abs(x - y) / torch.clamp_min(torch.maximum(torch.abs(x), torch.abs(y)), 1e-12) + + +# print utils +def indent(s, n=4): + """ + Indent a string. + """ + lines = s.split('\n') + for i in range(1, len(lines)): + lines[i] = ' ' * n + lines[i] + return '\n'.join(lines) + diff --git a/third_party/TRELLIS/trellis/utils/grad_clip_utils.py b/third_party/TRELLIS/trellis/utils/grad_clip_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..990a4352e24fc73bf732d8eb0f8ca9a07365b49e --- /dev/null +++ b/third_party/TRELLIS/trellis/utils/grad_clip_utils.py @@ -0,0 +1,81 @@ +from typing import * +import torch +import numpy as np +import torch.utils + + +class AdaptiveGradClipper: + """ + Adaptive gradient clipping for training. + """ + def __init__( + self, + max_norm=None, + clip_percentile=95.0, + buffer_size=1000, + ): + self.max_norm = max_norm + self.clip_percentile = clip_percentile + self.buffer_size = buffer_size + + self._grad_norm = np.zeros(buffer_size, dtype=np.float32) + self._max_norm = max_norm + self._buffer_ptr = 0 + self._buffer_length = 0 + + def __repr__(self): + return f'AdaptiveGradClipper(max_norm={self.max_norm}, clip_percentile={self.clip_percentile})' + + def state_dict(self): + return { + 'grad_norm': self._grad_norm, + 'max_norm': self._max_norm, + 'buffer_ptr': self._buffer_ptr, + 'buffer_length': self._buffer_length, + } + + def load_state_dict(self, state_dict): + self._grad_norm = state_dict['grad_norm'] + self._max_norm = state_dict['max_norm'] + self._buffer_ptr = state_dict['buffer_ptr'] + self._buffer_length = state_dict['buffer_length'] + + def log(self): + return { + 'max_norm': self._max_norm, + } + + def __call__(self, parameters, norm_type=2.0, error_if_nonfinite=False, foreach=None): + """Clip the gradient norm of an iterable of parameters. + + The norm is computed over all gradients together, as if they were + concatenated into a single vector. Gradients are modified in-place. + + Args: + parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will have gradients normalized + norm_type (float): type of the used p-norm. Can be ``'inf'`` for + infinity norm. + error_if_nonfinite (bool): if True, an error is thrown if the total + norm of the gradients from :attr:`parameters` is ``nan``, + ``inf``, or ``-inf``. Default: False (will switch to True in the future) + foreach (bool): use the faster foreach-based implementation. + If ``None``, use the foreach implementation for CUDA and CPU native tensors and silently + fall back to the slow implementation for other device types. + Default: ``None`` + + Returns: + Total norm of the parameter gradients (viewed as a single vector). + """ + max_norm = self._max_norm if self._max_norm is not None else float('inf') + grad_norm = torch.nn.utils.clip_grad_norm_(parameters, max_norm=max_norm, norm_type=norm_type, error_if_nonfinite=error_if_nonfinite, foreach=foreach) + + if torch.isfinite(grad_norm): + self._grad_norm[self._buffer_ptr] = grad_norm + self._buffer_ptr = (self._buffer_ptr + 1) % self.buffer_size + self._buffer_length = min(self._buffer_length + 1, self.buffer_size) + if self._buffer_length == self.buffer_size: + self._max_norm = np.percentile(self._grad_norm, self.clip_percentile) + self._max_norm = min(self._max_norm, self.max_norm) if self.max_norm is not None else self._max_norm + + return grad_norm \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/utils/loss_utils.py b/third_party/TRELLIS/trellis/utils/loss_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..52049f69543f2700bc5525b09cbf2fb25c08aa9e --- /dev/null +++ b/third_party/TRELLIS/trellis/utils/loss_utils.py @@ -0,0 +1,92 @@ +import torch +import torch.nn.functional as F +from torch.autograd import Variable +from math import exp +from lpips import LPIPS + + +def smooth_l1_loss(pred, target, beta=1.0): + diff = torch.abs(pred - target) + loss = torch.where(diff < beta, 0.5 * diff ** 2 / beta, diff - 0.5 * beta) + return loss.mean() + + +def l1_loss(network_output, gt): + return torch.abs((network_output - gt)).mean() + + +def l2_loss(network_output, gt): + return ((network_output - gt) ** 2).mean() + + +def gaussian(window_size, sigma): + gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + + +def create_window(window_size, channel): + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def psnr(img1, img2, max_val=1.0): + mse = F.mse_loss(img1, img2) + return 20 * torch.log10(max_val / torch.sqrt(mse)) + + +def ssim(img1, img2, window_size=11, size_average=True): + channel = img1.size(-3) + window = create_window(window_size, channel) + + if img1.is_cuda: + window = window.cuda(img1.get_device()) + window = window.type_as(img1) + + return _ssim(img1, img2, window, window_size, channel, size_average) + +def _ssim(img1, img2, window, window_size, channel, size_average=True): + mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel) + + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2 + + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + + ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + + if size_average: + return ssim_map.mean() + else: + return ssim_map.mean(1).mean(1).mean(1) + + +loss_fn_vgg = None +def lpips(img1, img2, value_range=(0, 1)): + global loss_fn_vgg + if loss_fn_vgg is None: + loss_fn_vgg = LPIPS(net='vgg').cuda().eval() + # normalize to [-1, 1] + img1 = (img1 - value_range[0]) / (value_range[1] - value_range[0]) * 2 - 1 + img2 = (img2 - value_range[0]) / (value_range[1] - value_range[0]) * 2 - 1 + return loss_fn_vgg(img1, img2).mean() + + +def normal_angle(pred, gt): + pred = pred * 2.0 - 1.0 + gt = gt * 2.0 - 1.0 + norms = pred.norm(dim=-1) * gt.norm(dim=-1) + cos_sim = (pred * gt).sum(-1) / (norms + 1e-9) + cos_sim = torch.clamp(cos_sim, -1.0, 1.0) + ang = torch.rad2deg(torch.acos(cos_sim[norms > 1e-9])).mean() + if ang.isnan(): + return -1 + return ang diff --git a/third_party/TRELLIS/trellis/utils/postprocessing_utils.py b/third_party/TRELLIS/trellis/utils/postprocessing_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..37e274ffa17a1f2f3ea7a771d4080925b0a9bfe3 --- /dev/null +++ b/third_party/TRELLIS/trellis/utils/postprocessing_utils.py @@ -0,0 +1,587 @@ +from typing import * +import numpy as np +import torch +import utils3d +import nvdiffrast.torch as dr +from tqdm import tqdm +import trimesh +import trimesh.visual +import xatlas +import pyvista as pv +from pymeshfix import _meshfix +import igraph +import cv2 +from PIL import Image +from .random_utils import sphere_hammersley_sequence +from .render_utils import render_multiview +from ..renderers import GaussianRenderer +from ..representations import Strivec, Gaussian, MeshExtractResult + + +@torch.no_grad() +def _fill_holes( + verts, + faces, + max_hole_size=0.04, + max_hole_nbe=32, + resolution=128, + num_views=500, + debug=False, + verbose=False +): + """ + Rasterize a mesh from multiple views and remove invisible faces. + Also includes postprocessing to: + 1. Remove connected components that are have low visibility. + 2. Mincut to remove faces at the inner side of the mesh connected to the outer side with a small hole. + + Args: + verts (torch.Tensor): Vertices of the mesh. Shape (V, 3). + faces (torch.Tensor): Faces of the mesh. Shape (F, 3). + max_hole_size (float): Maximum area of a hole to fill. + resolution (int): Resolution of the rasterization. + num_views (int): Number of views to rasterize the mesh. + verbose (bool): Whether to print progress. + """ + # Construct cameras + yaws = [] + pitchs = [] + for i in range(num_views): + y, p = sphere_hammersley_sequence(i, num_views) + yaws.append(y) + pitchs.append(p) + yaws = torch.tensor(yaws).cuda() + pitchs = torch.tensor(pitchs).cuda() + radius = 2.0 + fov = torch.deg2rad(torch.tensor(40)).cuda() + projection = utils3d.torch.perspective_from_fov_xy(fov, fov, 1, 3) + views = [] + for (yaw, pitch) in zip(yaws, pitchs): + orig = torch.tensor([ + torch.sin(yaw) * torch.cos(pitch), + torch.cos(yaw) * torch.cos(pitch), + torch.sin(pitch), + ]).cuda().float() * radius + view = utils3d.torch.view_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda()) + views.append(view) + views = torch.stack(views, dim=0) + + # Rasterize + visblity = torch.zeros(faces.shape[0], dtype=torch.int32, device=verts.device) + rastctx = utils3d.torch.RastContext(backend='cuda') + for i in tqdm(range(views.shape[0]), total=views.shape[0], disable=not verbose, desc='Rasterizing'): + view = views[i] + buffers = utils3d.torch.rasterize_triangle_faces( + rastctx, verts[None], faces, resolution, resolution, view=view, projection=projection + ) + face_id = buffers['face_id'][0][buffers['mask'][0] > 0.95] - 1 + face_id = torch.unique(face_id).long() + visblity[face_id] += 1 + visblity = visblity.float() / num_views + + # Mincut + ## construct outer faces + edges, face2edge, edge_degrees = utils3d.torch.compute_edges(faces) + boundary_edge_indices = torch.nonzero(edge_degrees == 1).reshape(-1) + connected_components = utils3d.torch.compute_connected_components(faces, edges, face2edge) + outer_face_indices = torch.zeros(faces.shape[0], dtype=torch.bool, device=faces.device) + for i in range(len(connected_components)): + outer_face_indices[connected_components[i]] = visblity[connected_components[i]] > min(max(visblity[connected_components[i]].quantile(0.75).item(), 0.25), 0.5) + outer_face_indices = outer_face_indices.nonzero().reshape(-1) + + ## construct inner faces + inner_face_indices = torch.nonzero(visblity == 0).reshape(-1) + if verbose: + tqdm.write(f'Found {inner_face_indices.shape[0]} invisible faces') + if inner_face_indices.shape[0] == 0: + return verts, faces + + ## Construct dual graph (faces as nodes, edges as edges) + dual_edges, dual_edge2edge = utils3d.torch.compute_dual_graph(face2edge) + dual_edge2edge = edges[dual_edge2edge] + dual_edges_weights = torch.norm(verts[dual_edge2edge[:, 0]] - verts[dual_edge2edge[:, 1]], dim=1) + if verbose: + tqdm.write(f'Dual graph: {dual_edges.shape[0]} edges') + + ## solve mincut problem + ### construct main graph + g = igraph.Graph() + g.add_vertices(faces.shape[0]) + g.add_edges(dual_edges.cpu().numpy()) + g.es['weight'] = dual_edges_weights.cpu().numpy() + + ### source and target + g.add_vertex('s') + g.add_vertex('t') + + ### connect invisible faces to source + g.add_edges([(f, 's') for f in inner_face_indices], attributes={'weight': torch.ones(inner_face_indices.shape[0], dtype=torch.float32).cpu().numpy()}) + + ### connect outer faces to target + g.add_edges([(f, 't') for f in outer_face_indices], attributes={'weight': torch.ones(outer_face_indices.shape[0], dtype=torch.float32).cpu().numpy()}) + + ### solve mincut + cut = g.mincut('s', 't', (np.array(g.es['weight']) * 1000).tolist()) + remove_face_indices = torch.tensor([v for v in cut.partition[0] if v < faces.shape[0]], dtype=torch.long, device=faces.device) + if verbose: + tqdm.write(f'Mincut solved, start checking the cut') + + ### check if the cut is valid with each connected component + to_remove_cc = utils3d.torch.compute_connected_components(faces[remove_face_indices]) + if debug: + tqdm.write(f'Number of connected components of the cut: {len(to_remove_cc)}') + valid_remove_cc = [] + cutting_edges = [] + for cc in to_remove_cc: + #### check if the connected component has low visibility + visblity_median = visblity[remove_face_indices[cc]].median() + if debug: + tqdm.write(f'visblity_median: {visblity_median}') + if visblity_median > 0.25: + continue + + #### check if the cuting loop is small enough + cc_edge_indices, cc_edges_degree = torch.unique(face2edge[remove_face_indices[cc]], return_counts=True) + cc_boundary_edge_indices = cc_edge_indices[cc_edges_degree == 1] + cc_new_boundary_edge_indices = cc_boundary_edge_indices[~torch.isin(cc_boundary_edge_indices, boundary_edge_indices)] + if len(cc_new_boundary_edge_indices) > 0: + cc_new_boundary_edge_cc = utils3d.torch.compute_edge_connected_components(edges[cc_new_boundary_edge_indices]) + cc_new_boundary_edges_cc_center = [verts[edges[cc_new_boundary_edge_indices[edge_cc]]].mean(dim=1).mean(dim=0) for edge_cc in cc_new_boundary_edge_cc] + cc_new_boundary_edges_cc_area = [] + for i, edge_cc in enumerate(cc_new_boundary_edge_cc): + _e1 = verts[edges[cc_new_boundary_edge_indices[edge_cc]][:, 0]] - cc_new_boundary_edges_cc_center[i] + _e2 = verts[edges[cc_new_boundary_edge_indices[edge_cc]][:, 1]] - cc_new_boundary_edges_cc_center[i] + cc_new_boundary_edges_cc_area.append(torch.norm(torch.cross(_e1, _e2, dim=-1), dim=1).sum() * 0.5) + if debug: + cutting_edges.append(cc_new_boundary_edge_indices) + tqdm.write(f'Area of the cutting loop: {cc_new_boundary_edges_cc_area}') + if any([l > max_hole_size for l in cc_new_boundary_edges_cc_area]): + continue + + valid_remove_cc.append(cc) + + if debug: + face_v = verts[faces].mean(dim=1).cpu().numpy() + vis_dual_edges = dual_edges.cpu().numpy() + vis_colors = np.zeros((faces.shape[0], 3), dtype=np.uint8) + vis_colors[inner_face_indices.cpu().numpy()] = [0, 0, 255] + vis_colors[outer_face_indices.cpu().numpy()] = [0, 255, 0] + vis_colors[remove_face_indices.cpu().numpy()] = [255, 0, 255] + if len(valid_remove_cc) > 0: + vis_colors[remove_face_indices[torch.cat(valid_remove_cc)].cpu().numpy()] = [255, 0, 0] + utils3d.io.write_ply('dbg_dual.ply', face_v, edges=vis_dual_edges, vertex_colors=vis_colors) + + vis_verts = verts.cpu().numpy() + vis_edges = edges[torch.cat(cutting_edges)].cpu().numpy() + utils3d.io.write_ply('dbg_cut.ply', vis_verts, edges=vis_edges) + + + if len(valid_remove_cc) > 0: + remove_face_indices = remove_face_indices[torch.cat(valid_remove_cc)] + mask = torch.ones(faces.shape[0], dtype=torch.bool, device=faces.device) + mask[remove_face_indices] = 0 + faces = faces[mask] + faces, verts = utils3d.torch.remove_unreferenced_vertices(faces, verts) + if verbose: + tqdm.write(f'Removed {(~mask).sum()} faces by mincut') + else: + if verbose: + tqdm.write(f'Removed 0 faces by mincut') + + mesh = _meshfix.PyTMesh() + mesh.load_array(verts.cpu().numpy(), faces.cpu().numpy()) + mesh.fill_small_boundaries(nbe=max_hole_nbe, refine=True) + verts, faces = mesh.return_arrays() + verts, faces = torch.tensor(verts, device='cuda', dtype=torch.float32), torch.tensor(faces, device='cuda', dtype=torch.int32) + + return verts, faces + + +def postprocess_mesh( + vertices: np.array, + faces: np.array, + simplify: bool = True, + simplify_ratio: float = 0.9, + fill_holes: bool = True, + fill_holes_max_hole_size: float = 0.04, + fill_holes_max_hole_nbe: int = 32, + fill_holes_resolution: int = 1024, + fill_holes_num_views: int = 1000, + debug: bool = False, + verbose: bool = False, +): + """ + Postprocess a mesh by simplifying, removing invisible faces, and removing isolated pieces. + + Args: + vertices (np.array): Vertices of the mesh. Shape (V, 3). + faces (np.array): Faces of the mesh. Shape (F, 3). + simplify (bool): Whether to simplify the mesh, using quadric edge collapse. + simplify_ratio (float): Ratio of faces to keep after simplification. + fill_holes (bool): Whether to fill holes in the mesh. + fill_holes_max_hole_size (float): Maximum area of a hole to fill. + fill_holes_max_hole_nbe (int): Maximum number of boundary edges of a hole to fill. + fill_holes_resolution (int): Resolution of the rasterization. + fill_holes_num_views (int): Number of views to rasterize the mesh. + verbose (bool): Whether to print progress. + """ + + if verbose: + tqdm.write(f'Before postprocess: {vertices.shape[0]} vertices, {faces.shape[0]} faces') + + # Simplify + if simplify and simplify_ratio > 0: + mesh = pv.PolyData(vertices, np.concatenate([np.full((faces.shape[0], 1), 3), faces], axis=1)) + mesh = mesh.decimate(simplify_ratio, progress_bar=verbose) + vertices, faces = mesh.points, mesh.faces.reshape(-1, 4)[:, 1:] + if verbose: + tqdm.write(f'After decimate: {vertices.shape[0]} vertices, {faces.shape[0]} faces') + + # Remove invisible faces + if fill_holes: + vertices, faces = torch.tensor(vertices).cuda(), torch.tensor(faces.astype(np.int32)).cuda() + vertices, faces = _fill_holes( + vertices, faces, + max_hole_size=fill_holes_max_hole_size, + max_hole_nbe=fill_holes_max_hole_nbe, + resolution=fill_holes_resolution, + num_views=fill_holes_num_views, + debug=debug, + verbose=verbose, + ) + vertices, faces = vertices.cpu().numpy(), faces.cpu().numpy() + if verbose: + tqdm.write(f'After remove invisible faces: {vertices.shape[0]} vertices, {faces.shape[0]} faces') + + return vertices, faces + + +def parametrize_mesh(vertices: np.array, faces: np.array): + """ + Parametrize a mesh to a texture space, using xatlas. + + Args: + vertices (np.array): Vertices of the mesh. Shape (V, 3). + faces (np.array): Faces of the mesh. Shape (F, 3). + """ + + vmapping, indices, uvs = xatlas.parametrize(vertices, faces) + + vertices = vertices[vmapping] + faces = indices + + return vertices, faces, uvs + + +def bake_texture( + vertices: np.array, + faces: np.array, + uvs: np.array, + observations: List[np.array], + masks: List[np.array], + extrinsics: List[np.array], + intrinsics: List[np.array], + texture_size: int = 2048, + near: float = 0.1, + far: float = 10.0, + mode: Literal['fast', 'opt'] = 'opt', + lambda_tv: float = 1e-2, + verbose: bool = False, +): + """ + Bake texture to a mesh from multiple observations. + + Args: + vertices (np.array): Vertices of the mesh. Shape (V, 3). + faces (np.array): Faces of the mesh. Shape (F, 3). + uvs (np.array): UV coordinates of the mesh. Shape (V, 2). + observations (List[np.array]): List of observations. Each observation is a 2D image. Shape (H, W, 3). + masks (List[np.array]): List of masks. Each mask is a 2D image. Shape (H, W). + extrinsics (List[np.array]): List of extrinsics. Shape (4, 4). + intrinsics (List[np.array]): List of intrinsics. Shape (3, 3). + texture_size (int): Size of the texture. + near (float): Near plane of the camera. + far (float): Far plane of the camera. + mode (Literal['fast', 'opt']): Mode of texture baking. + lambda_tv (float): Weight of total variation loss in optimization. + verbose (bool): Whether to print progress. + """ + vertices = torch.tensor(vertices).cuda() + faces = torch.tensor(faces.astype(np.int32)).cuda() + uvs = torch.tensor(uvs).cuda() + observations = [torch.tensor(obs / 255.0).float().cuda() for obs in observations] + masks = [torch.tensor(m>0).bool().cuda() for m in masks] + views = [utils3d.torch.extrinsics_to_view(torch.tensor(extr).cuda()) for extr in extrinsics] + projections = [utils3d.torch.intrinsics_to_perspective(torch.tensor(intr).cuda(), near, far) for intr in intrinsics] + + if mode == 'fast': + texture = torch.zeros((texture_size * texture_size, 3), dtype=torch.float32).cuda() + texture_weights = torch.zeros((texture_size * texture_size), dtype=torch.float32).cuda() + rastctx = utils3d.torch.RastContext(backend='cuda') + for observation, view, projection in tqdm(zip(observations, views, projections), total=len(observations), disable=not verbose, desc='Texture baking (fast)'): + with torch.no_grad(): + rast = utils3d.torch.rasterize_triangle_faces( + rastctx, vertices[None], faces, observation.shape[1], observation.shape[0], uv=uvs[None], view=view, projection=projection + ) + uv_map = rast['uv'][0].detach().flip(0) + mask = rast['mask'][0].detach().bool() & masks[0] + + # nearest neighbor interpolation + uv_map = (uv_map * texture_size).floor().long() + obs = observation[mask] + uv_map = uv_map[mask] + idx = uv_map[:, 0] + (texture_size - uv_map[:, 1] - 1) * texture_size + texture = texture.scatter_add(0, idx.view(-1, 1).expand(-1, 3), obs) + texture_weights = texture_weights.scatter_add(0, idx, torch.ones((obs.shape[0]), dtype=torch.float32, device=texture.device)) + + mask = texture_weights > 0 + texture[mask] /= texture_weights[mask][:, None] + texture = np.clip(texture.reshape(texture_size, texture_size, 3).cpu().numpy() * 255, 0, 255).astype(np.uint8) + + # inpaint + mask = (texture_weights == 0).cpu().numpy().astype(np.uint8).reshape(texture_size, texture_size) + texture = cv2.inpaint(texture, mask, 3, cv2.INPAINT_TELEA) + + elif mode == 'opt': + rastctx = utils3d.torch.RastContext(backend='cuda') + observations = [observations.flip(0) for observations in observations] + masks = [m.flip(0) for m in masks] + _uv = [] + _uv_dr = [] + for observation, view, projection in tqdm(zip(observations, views, projections), total=len(views), disable=not verbose, desc='Texture baking (opt): UV'): + with torch.no_grad(): + rast = utils3d.torch.rasterize_triangle_faces( + rastctx, vertices[None], faces, observation.shape[1], observation.shape[0], uv=uvs[None], view=view, projection=projection + ) + _uv.append(rast['uv'].detach()) + _uv_dr.append(rast['uv_dr'].detach()) + + texture = torch.nn.Parameter(torch.zeros((1, texture_size, texture_size, 3), dtype=torch.float32).cuda()) + optimizer = torch.optim.Adam([texture], betas=(0.5, 0.9), lr=1e-2) + + def exp_anealing(optimizer, step, total_steps, start_lr, end_lr): + return start_lr * (end_lr / start_lr) ** (step / total_steps) + + def cosine_anealing(optimizer, step, total_steps, start_lr, end_lr): + return end_lr + 0.5 * (start_lr - end_lr) * (1 + np.cos(np.pi * step / total_steps)) + + def tv_loss(texture): + return torch.nn.functional.l1_loss(texture[:, :-1, :, :], texture[:, 1:, :, :]) + \ + torch.nn.functional.l1_loss(texture[:, :, :-1, :], texture[:, :, 1:, :]) + + total_steps = 2500 + with tqdm(total=total_steps, disable=not verbose, desc='Texture baking (opt): optimizing') as pbar: + for step in range(total_steps): + optimizer.zero_grad() + selected = np.random.randint(0, len(views)) + uv, uv_dr, observation, mask = _uv[selected], _uv_dr[selected], observations[selected], masks[selected] + render = dr.texture(texture, uv, uv_dr)[0] + loss = torch.nn.functional.l1_loss(render[mask], observation[mask]) + if lambda_tv > 0: + loss += lambda_tv * tv_loss(texture) + loss.backward() + optimizer.step() + # annealing + optimizer.param_groups[0]['lr'] = cosine_anealing(optimizer, step, total_steps, 1e-2, 1e-5) + pbar.set_postfix({'loss': loss.item()}) + pbar.update() + texture = np.clip(texture[0].flip(0).detach().cpu().numpy() * 255, 0, 255).astype(np.uint8) + mask = 1 - utils3d.torch.rasterize_triangle_faces( + rastctx, (uvs * 2 - 1)[None], faces, texture_size, texture_size + )['mask'][0].detach().cpu().numpy().astype(np.uint8) + texture = cv2.inpaint(texture, mask, 3, cv2.INPAINT_TELEA) + else: + raise VaºlueError(f'Unknown mode: {mode}') + + return texture + + +def to_glb( + app_rep: Union[Strivec, Gaussian], + mesh: MeshExtractResult, + simplify: float = 0.95, + fill_holes: bool = True, + fill_holes_max_size: float = 0.04, + texture_size: int = 1024, + debug: bool = False, + verbose: bool = True, +) -> trimesh.Trimesh: + """ + Convert a generated asset to a glb file. + + Args: + app_rep (Union[Strivec, Gaussian]): Appearance representation. + mesh (MeshExtractResult): Extracted mesh. + simplify (float): Ratio of faces to remove in simplification. + fill_holes (bool): Whether to fill holes in the mesh. + fill_holes_max_size (float): Maximum area of a hole to fill. + texture_size (int): Size of the texture. + debug (bool): Whether to print debug information. + verbose (bool): Whether to print progress. + """ + vertices = mesh.vertices.cpu().numpy() + faces = mesh.faces.cpu().numpy() + + # mesh postprocess + vertices, faces = postprocess_mesh( + vertices, faces, + simplify=simplify > 0, + simplify_ratio=simplify, + fill_holes=fill_holes, + fill_holes_max_hole_size=fill_holes_max_size, + fill_holes_max_hole_nbe=int(250 * np.sqrt(1-simplify)), + fill_holes_resolution=1024, + fill_holes_num_views=1000, + debug=debug, + verbose=verbose, + ) + + # parametrize mesh + vertices, faces, uvs = parametrize_mesh(vertices, faces) + + # bake texture + observations, extrinsics, intrinsics = render_multiview(app_rep, resolution=1024, nviews=100) + masks = [np.any(observation > 0, axis=-1) for observation in observations] + extrinsics = [extrinsics[i].cpu().numpy() for i in range(len(extrinsics))] + intrinsics = [intrinsics[i].cpu().numpy() for i in range(len(intrinsics))] + texture = bake_texture( + vertices, faces, uvs, + observations, masks, extrinsics, intrinsics, + texture_size=texture_size, mode='opt', + lambda_tv=0.01, + verbose=verbose + ) + texture = Image.fromarray(texture) + + material = trimesh.visual.material.PBRMaterial( + roughnessFactor=1.0, + baseColorTexture=texture, + baseColorFactor=np.array([255, 255, 255, 255], dtype=np.uint8) + ) + mesh = trimesh.Trimesh(vertices, faces, visual=trimesh.visual.TextureVisuals(uv=uvs, material=material)) + mesh.vertex_normals + mesh.face_normals + return mesh + + +def simplify_gs( + gs: Gaussian, + simplify: float = 0.95, + verbose: bool = True, +): + """ + Simplify 3D Gaussians + NOTE: this function is not used in the current implementation for the unsatisfactory performance. + + Args: + gs (Gaussian): 3D Gaussian. + simplify (float): Ratio of Gaussians to remove in simplification. + """ + if simplify <= 0: + return gs + + # simplify + observations, extrinsics, intrinsics = render_multiview(gs, resolution=1024, nviews=100) + observations = [torch.tensor(obs / 255.0).float().cuda().permute(2, 0, 1) for obs in observations] + + # Following https://arxiv.org/pdf/2411.06019 + renderer = GaussianRenderer({ + "resolution": 1024, + "near": 0.8, + "far": 1.6, + "ssaa": 1, + "bg_color": (0,0,0), + }) + new_gs = Gaussian(**gs.init_params) + new_gs._features_dc = gs._features_dc.clone() + new_gs._features_rest = gs._features_rest.clone() if gs._features_rest is not None else None + new_gs._opacity = torch.nn.Parameter(gs._opacity.clone()) + new_gs._rotation = torch.nn.Parameter(gs._rotation.clone()) + new_gs._scaling = torch.nn.Parameter(gs._scaling.clone()) + new_gs._xyz = torch.nn.Parameter(gs._xyz.clone()) + + start_lr = [1e-4, 1e-3, 5e-3, 0.025] + end_lr = [1e-6, 1e-5, 5e-5, 0.00025] + optimizer = torch.optim.Adam([ + {"params": new_gs._xyz, "lr": start_lr[0]}, + {"params": new_gs._rotation, "lr": start_lr[1]}, + {"params": new_gs._scaling, "lr": start_lr[2]}, + {"params": new_gs._opacity, "lr": start_lr[3]}, + ], lr=start_lr[0]) + + def exp_anealing(optimizer, step, total_steps, start_lr, end_lr): + return start_lr * (end_lr / start_lr) ** (step / total_steps) + + def cosine_anealing(optimizer, step, total_steps, start_lr, end_lr): + return end_lr + 0.5 * (start_lr - end_lr) * (1 + np.cos(np.pi * step / total_steps)) + + _zeta = new_gs.get_opacity.clone().detach().squeeze() + _lambda = torch.zeros_like(_zeta) + _delta = 1e-7 + _interval = 10 + num_target = int((1 - simplify) * _zeta.shape[0]) + + with tqdm(total=2500, disable=not verbose, desc='Simplifying Gaussian') as pbar: + for i in range(2500): + # prune + if i % 100 == 0: + mask = new_gs.get_opacity.squeeze() > 0.05 + mask = torch.nonzero(mask).squeeze() + new_gs._xyz = torch.nn.Parameter(new_gs._xyz[mask]) + new_gs._rotation = torch.nn.Parameter(new_gs._rotation[mask]) + new_gs._scaling = torch.nn.Parameter(new_gs._scaling[mask]) + new_gs._opacity = torch.nn.Parameter(new_gs._opacity[mask]) + new_gs._features_dc = new_gs._features_dc[mask] + new_gs._features_rest = new_gs._features_rest[mask] if new_gs._features_rest is not None else None + _zeta = _zeta[mask] + _lambda = _lambda[mask] + # update optimizer state + for param_group, new_param in zip(optimizer.param_groups, [new_gs._xyz, new_gs._rotation, new_gs._scaling, new_gs._opacity]): + stored_state = optimizer.state[param_group['params'][0]] + if 'exp_avg' in stored_state: + stored_state['exp_avg'] = stored_state['exp_avg'][mask] + stored_state['exp_avg_sq'] = stored_state['exp_avg_sq'][mask] + del optimizer.state[param_group['params'][0]] + param_group['params'][0] = new_param + optimizer.state[param_group['params'][0]] = stored_state + + opacity = new_gs.get_opacity.squeeze() + + # sparisfy + if i % _interval == 0: + _zeta = _lambda + opacity.detach() + if opacity.shape[0] > num_target: + index = _zeta.topk(num_target)[1] + _m = torch.ones_like(_zeta, dtype=torch.bool) + _m[index] = 0 + _zeta[_m] = 0 + _lambda = _lambda + opacity.detach() - _zeta + + # sample a random view + view_idx = np.random.randint(len(observations)) + observation = observations[view_idx] + extrinsic = extrinsics[view_idx] + intrinsic = intrinsics[view_idx] + + color = renderer.render(new_gs, extrinsic, intrinsic)['color'] + rgb_loss = torch.nn.functional.l1_loss(color, observation) + loss = rgb_loss + \ + _delta * torch.sum(torch.pow(_lambda + opacity - _zeta, 2)) + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + # update lr + for j in range(len(optimizer.param_groups)): + optimizer.param_groups[j]['lr'] = cosine_anealing(optimizer, i, 2500, start_lr[j], end_lr[j]) + + pbar.set_postfix({'loss': rgb_loss.item(), 'num': opacity.shape[0], 'lambda': _lambda.mean().item()}) + pbar.update() + + new_gs._xyz = new_gs._xyz.data + new_gs._rotation = new_gs._rotation.data + new_gs._scaling = new_gs._scaling.data + new_gs._opacity = new_gs._opacity.data + + return new_gs diff --git a/third_party/TRELLIS/trellis/utils/random_utils.py b/third_party/TRELLIS/trellis/utils/random_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5b668c277b51f4930991912a80573adc79364028 --- /dev/null +++ b/third_party/TRELLIS/trellis/utils/random_utils.py @@ -0,0 +1,30 @@ +import numpy as np + +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), remap=False): + u, v = hammersley_sequence(2, n, num_samples) + u += offset[0] / num_samples + v += offset[1] + if remap: + 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] \ No newline at end of file diff --git a/third_party/TRELLIS/trellis/utils/render_utils.py b/third_party/TRELLIS/trellis/utils/render_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..de8e96345f9fe233444c4b901987505e4b77a5ff --- /dev/null +++ b/third_party/TRELLIS/trellis/utils/render_utils.py @@ -0,0 +1,116 @@ +import torch +import numpy as np +from tqdm import tqdm +import utils3d +from PIL import Image + +from ..renderers import OctreeRenderer, GaussianRenderer, MeshRenderer +from ..representations import Octree, Gaussian, MeshExtractResult +from ..modules import sparse as sp +from .random_utils import sphere_hammersley_sequence + + +def yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitchs, rs, fovs): + is_list = isinstance(yaws, list) + if not is_list: + yaws = [yaws] + pitchs = [pitchs] + if not isinstance(rs, list): + rs = [rs] * len(yaws) + if not isinstance(fovs, list): + fovs = [fovs] * len(yaws) + extrinsics = [] + intrinsics = [] + for yaw, pitch, r, fov in zip(yaws, pitchs, rs, fovs): + fov = torch.deg2rad(torch.tensor(float(fov))).cuda() + yaw = torch.tensor(float(yaw)).cuda() + pitch = torch.tensor(float(pitch)).cuda() + orig = torch.tensor([ + torch.sin(yaw) * torch.cos(pitch), + torch.cos(yaw) * torch.cos(pitch), + torch.sin(pitch), + ]).cuda() * r + extr = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda()) + intr = utils3d.torch.intrinsics_from_fov_xy(fov, fov) + extrinsics.append(extr) + intrinsics.append(intr) + if not is_list: + extrinsics = extrinsics[0] + intrinsics = intrinsics[0] + return extrinsics, intrinsics + + +def render_frames(sample, extrinsics, intrinsics, options={}, colors_overwrite=None, verbose=True, **kwargs): + if isinstance(sample, Octree): + renderer = OctreeRenderer() + renderer.rendering_options.resolution = options.get('resolution', 512) + renderer.rendering_options.near = options.get('near', 0.8) + renderer.rendering_options.far = options.get('far', 1.6) + renderer.rendering_options.bg_color = options.get('bg_color', (0, 0, 0)) + renderer.rendering_options.ssaa = options.get('ssaa', 4) + renderer.pipe.primitive = sample.primitive + elif isinstance(sample, Gaussian): + renderer = GaussianRenderer() + renderer.rendering_options.resolution = options.get('resolution', 512) + renderer.rendering_options.near = options.get('near', 0.8) + renderer.rendering_options.far = options.get('far', 1.6) + renderer.rendering_options.bg_color = options.get('bg_color', (0, 0, 0)) + renderer.rendering_options.ssaa = options.get('ssaa', 1) + renderer.pipe.kernel_size = kwargs.get('kernel_size', 0.1) + renderer.pipe.use_mip_gaussian = True + elif isinstance(sample, MeshExtractResult): + renderer = MeshRenderer() + renderer.rendering_options.resolution = options.get('resolution', 512) + renderer.rendering_options.near = options.get('near', 1) + renderer.rendering_options.far = options.get('far', 100) + renderer.rendering_options.ssaa = options.get('ssaa', 4) + else: + raise ValueError(f'Unsupported sample type: {type(sample)}') + + rets = {} + for j, (extr, intr) in tqdm(enumerate(zip(extrinsics, intrinsics)), desc='Rendering', disable=not verbose): + if not isinstance(sample, MeshExtractResult): + res = renderer.render(sample, extr, intr, colors_overwrite=colors_overwrite) + if 'color' not in rets: rets['color'] = [] + if 'depth' not in rets: rets['depth'] = [] + rets['color'].append(np.clip(res['color'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8)) + if 'percent_depth' in res: + rets['depth'].append(res['percent_depth'].detach().cpu().numpy()) + elif 'depth' in res: + rets['depth'].append(res['depth'].detach().cpu().numpy()) + else: + rets['depth'].append(None) + else: + res = renderer.render(sample, extr, intr) + if 'normal' not in rets: rets['normal'] = [] + + rets['normal'].append(np.clip(res['normal'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8)) + return rets + +def render_video(sample, resolution=512, bg_color=(0, 0, 0), num_frames=300, r=2, fov=40, **kwargs): + yaws = torch.linspace(0, 2 * 3.1415, num_frames) + pitch = 0.25 + 0.5 * torch.sin(torch.linspace(0, 2 * 3.1415, num_frames)) + yaws = yaws.tolist() + pitch = pitch.tolist() + extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitch, r, fov) + return render_frames(sample, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': bg_color}, **kwargs) + + +def render_multiview(sample, resolution=512, nviews=30): + r = 2 + fov = 40 + cams = [sphere_hammersley_sequence(i, nviews) for i in range(nviews)] + yaws = [cam[0] for cam in cams] + pitchs = [cam[1] for cam in cams] + extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitchs, r, fov) + res = render_frames(sample, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': (0, 0, 0)}) + return res['color'], extrinsics, intrinsics + + +def render_snapshot(samples, resolution=512, bg_color=(0, 0, 0), offset=(-16 / 180 * np.pi, 20 / 180 * np.pi), r=10, fov=8, **kwargs): + yaw = [0, np.pi/2, np.pi, 3*np.pi/2] + yaw_offset = offset[0] + yaw = [y + yaw_offset for y in yaw] + pitch = [offset[1] for _ in range(4)] + extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaw, pitch, r, fov) + return render_frames(samples, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': bg_color}, **kwargs) diff --git a/utils.py b/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a1edc2fe81931d8fd60d11b785ce3353265dadd9 --- /dev/null +++ b/utils.py @@ -0,0 +1,88 @@ +import numpy as np +import open3d_pycg as o3d +import os +import torch +from sklearn.decomposition import PCA + + +def merge_meshes(mesh_list): + merged = o3d.geometry.TriangleMesh() + v_offset = 0 + + for mesh in mesh_list: + vertices = np.asarray(mesh.vertices) + triangles = np.asarray(mesh.triangles) + v_offset + + merged.vertices.extend(o3d.utility.Vector3dVector(vertices)) + merged.triangles.extend(o3d.utility.Vector3iVector(triangles)) + + v_offset += len(vertices) + + return merged + + +def voxelgrid_to_open3d(voxels: np.ndarray, threshold=0.5): + if len(voxels.shape) > 3: + C, D, H, W = voxels.shape + flat_feats = voxels.reshape(C, -1).transpose(1,0) + pca = PCA(n_components=3) + reduced = pca.fit_transform(flat_feats) + # Compute feature norm and PCA color std + + # Normalize for RGB + reduced -= reduced.min(0) + reduced /= reduced.max(0) + 1e-6 + + # Compute norms and color std + norms = np.linalg.norm(flat_feats, axis=1) + color_std = np.std(reduced, axis=1) + + # Filter: active voxels with non-trivial color + mask = (norms > threshold) & (color_std > 1e-3) + + # zz, yy, xx = np.meshgrid(np.arange(D), np.arange(H), np.arange(W), indexing='ij') + xx, yy, zz = np.meshgrid(np.arange(D), np.arange(H), np.arange(W), indexing='ij') + coords = np.stack([xx, yy, zz], axis=-1).reshape(-1, 3) + valid_coords = coords[mask] + valid_colors = reduced[mask] + + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(valid_coords.astype(np.float32)) + pcd.colors = o3d.utility.Vector3dVector(valid_colors.astype(np.float32)) + else: + coords = np.argwhere(voxels > threshold) + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(coords) + return pcd + + +def save_voxelgrid_as_ply(voxels: np.ndarray, filename: str, threshold=0.5): + pcd = voxelgrid_to_open3d(voxels, threshold) + o3d.io.write_point_cloud(filename, pcd) + + +def voxelize_sq_francis(file_name): + superquadric_mesh = o3d.io.read_triangle_mesh(file_name) + vertices = np.clip(np.asarray(superquadric_mesh.vertices), -0.5 + 1e-6, 0.5 - 1e-6) + superquadric_mesh.vertices = o3d.utility.Vector3dVector(vertices) + theta = np.pi / 2 + #superquadric_mesh.rotate(R_x, center=(0, 0, 0)) + voxel_grid = o3d.geometry.VoxelGrid.create_from_triangle_mesh_within_bounds( + superquadric_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()]) + unique_points = np.unique(vertices, axis=0) + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(unique_points) + o3d.io.write_point_cloud("merged_mesh_voxelized.ply", pcd) + + zeros = np.zeros((unique_points.shape[0], 1)) + unique_points_4d = np.hstack((zeros, unique_points)) # shape [N, 4] + unique_points_4d_torch = torch.from_numpy(unique_points_4d).to(dtype=torch.int32, device='cpu') + my_coords_orig = unique_points_4d_torch + + coords_dense = torch.ones(1, 1, 64, 64, 64).to(device='cpu', dtype=torch.float32) * 0.0 + for i in range(my_coords_orig.shape[0]): + x, y, z = my_coords_orig[i, 1], my_coords_orig[i, 2], my_coords_orig[i, 3] + coords_dense[0, 0, x, y, z] = 1.0 + return coords_dense \ No newline at end of file diff --git a/wheels/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl b/wheels/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl new file mode 100644 index 0000000000000000000000000000000000000000..eb06cf186a69a1a81ae9ac18ccc7e9425de667b1 --- /dev/null +++ b/wheels/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:77cc3f8fe1620242a502e5a1361d56dad083ad02468f10ef5fb9e81f54b65da0 +size 3213246