diff --git a/.github/workflows/ci-cd.yml b/.github/workflows/ci-cd.yml new file mode 100644 index 00000000..c6609bf3 --- /dev/null +++ b/.github/workflows/ci-cd.yml @@ -0,0 +1,117 @@ +name: CI / CD + +on: + pull_request: + branches: [ "main" ] + types: [opened, synchronize, reopened, ready_for_review] + merge_group: + branches: [ "main" ] + push: + branches: [ "main" ] + workflow_dispatch: + +concurrency: + group: ci-cd-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + build: + name: Build checks (Python ${{ matrix.python-version }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: [ "3.13" ] + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + + - name: Upgrade pip + run: python -m pip install --upgrade pip + + - name: Install CI helper deps + run: python -m pip install packaging + + - name: Validate dependency spec files + run: | + python - <<'PY' + import tomllib + from pathlib import Path + from packaging.requirements import Requirement + + def check_req_line(path: Path): + bad = [] + for i, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): + s = line.strip() + if not s or s.startswith("#") or s.startswith("--") or s.startswith("./"): + continue + try: + Requirement(s) + except Exception as exc: + bad.append((i, s, str(exc))) + return bad + + errors = [] + req_errors = check_req_line(Path("requirements.txt")) + if req_errors: + for line_no, value, exc in req_errors: + errors.append(f"requirements.txt:{line_no}: {value} :: {exc}") + + pyproject = tomllib.loads(Path("o-voxel/pyproject.toml").read_text(encoding="utf-8")) + for dep in pyproject["project"]["dependencies"]: + try: + Requirement(dep) + except Exception as exc: + errors.append(f"o-voxel/pyproject.toml dependency {dep!r} :: {exc}") + + if errors: + raise SystemExit("\n".join(errors)) + print("Dependency specs are valid.") + PY + + - name: Compile Python sources + run: | + python -m compileall -q app.py app_texturing.py example.py example_texturing.py setup_windows.py train.py + python -m compileall -q trellis2 + python -m compileall -q o-voxel/o_voxel + + cd-artifacts: + name: CD artifact snapshot + if: github.event_name == 'push' && github.ref == 'refs/heads/main' + runs-on: ubuntu-latest + needs: build + permissions: + contents: read + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Create build manifest + run: | + mkdir -p out + cat > out/build-manifest.json <<'JSON' + { + "repo": "${{ github.repository }}", + "branch": "${{ github.ref_name }}", + "commit": "${{ github.sha }}", + "workflow": "${{ github.workflow }}", + "run_id": "${{ github.run_id }}", + "run_number": "${{ github.run_number }}", + "created_utc": "${{ github.event.head_commit.timestamp }}" + } + JSON + + - name: Upload build manifest + uses: actions/upload-artifact@v4 + with: + name: build-manifest-${{ github.run_number }} + path: out/build-manifest.json + if-no-files-found: error diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 00000000..6e676795 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,84 @@ +name: Release (tag) + +on: + push: + tags: + - "v*" + workflow_dispatch: + inputs: + tag: + description: "Existing tag to release (e.g. v1.0.0)" + required: true + type: string + +permissions: + contents: write + +concurrency: + group: release-${{ github.ref }} + cancel-in-progress: false + +jobs: + release: + name: Create GitHub release + runs-on: ubuntu-latest + + steps: + - name: Resolve release tag + id: tag + run: | + if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then + echo "value=${{ inputs.tag }}" >> "$GITHUB_OUTPUT" + else + echo "value=${GITHUB_REF#refs/tags/}" >> "$GITHUB_OUTPUT" + fi + + - name: Checkout repository at tag + uses: actions/checkout@v4 + with: + ref: ${{ steps.tag.outputs.value }} + fetch-depth: 0 + submodules: recursive + + - name: Build release archives + shell: bash + run: | + set -euo pipefail + TAG="${{ steps.tag.outputs.value }}" + STAGE="release-stage" + OUT="release-out" + PKG="TRELLIS.2-${TAG}" + mkdir -p "$STAGE/$PKG" "$OUT" + + rsync -a \ + --exclude '.git' \ + --exclude '.github' \ + --exclude '__pycache__' \ + --exclude '*.pyc' \ + ./ "$STAGE/$PKG/" + + cat > "$OUT/build-manifest.json" </dev/null) + (cd "$OUT" && sha256sum "${PKG}.tar.gz" "${PKG}.zip" > checksums.sha256) + + - name: Publish GitHub release + uses: softprops/action-gh-release@v2 + with: + tag_name: ${{ steps.tag.outputs.value }} + generate_release_notes: true + files: | + release-out/*.tar.gz + release-out/*.zip + release-out/checksums.sha256 + release-out/build-manifest.json diff --git a/.gitignore b/.gitignore index b7faf403..c4c6967f 100644 --- a/.gitignore +++ b/.gitignore @@ -205,3 +205,10 @@ cython_debug/ marimo/_static/ marimo/_lsp/ __marimo__/ + +# app.py runtime output +tmp/ + +# example.py generated output +sample.mp4 +sample.glb diff --git a/README.md b/README.md index 1c0ad06b..bc4b1b30 100644 --- a/README.md +++ b/README.md @@ -64,6 +64,74 @@ Data processing is streamlined for instant conversions that are fully **renderin - Python version 3.8 or higher is required. ### Installation Steps + +#### Windows (Python 3.13 + RTX 50-series / Blackwell) +This fork includes a pip-first Windows path for Python 3.13, CUDA Toolkit 13.0, PyTorch `2.13.0+cu130`, and NVIDIA RTX 50-series / Blackwell GPUs (`sm_120`). Conda is not required. + +1. Clone the repo and initialize submodules: + ```powershell + git clone -b windows-blackwell https://github.com/rwfsmith/TRELLIS.2.git --recursive + cd TRELLIS.2 + git submodule update --init --recursive + ``` + +2. Create and activate a Python 3.13 virtual environment: + ```powershell + py -3.13 -m venv venv + .\venv\Scripts\Activate.ps1 + python -m pip install --upgrade pip setuptools wheel + ``` + +3. Install Visual Studio 2022 C++ build tools and CUDA Toolkit 13.0. Then install dependencies: + ```powershell + powershell -ExecutionPolicy Bypass -File .\setup_windows.ps1 -Python .\venv\Scripts\python.exe + ``` + The setup script wraps: + ```powershell + python -m pip install -r requirements.txt --no-build-isolation + ``` + and sets the native build environment variables needed by PyTorch CUDA extensions: + `CUDA_HOME`, `TORCH_CUDA_ARCH_LIST=12.0`, `DISTUTILS_USE_SDK=1`, and `MSSdk=1`. + + If CUDA is installed somewhere other than `C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.0`, pass it explicitly: + ```powershell + powershell -ExecutionPolicy Bypass -File .\setup_windows.ps1 -Python .\venv\Scripts\python.exe -CudaHome "D:\CUDA\v13.0" + ``` + + If you are running from `cmd.exe` instead of PowerShell, use the bundled batch wrapper instead (it forwards all arguments to `setup_windows.ps1`): + ```bat + setup_windows.bat -Python .\venv\Scripts\python.exe + ``` + + Alternatively, a pure Python setup script is provided for those who prefer not to invoke PowerShell or a `.bat` file at all: + ``` + python setup_windows.py --python .\venv\Scripts\python.exe + ``` + It performs the same steps as `setup_windows.ps1` (locate VS2022 build tools, set `CUDA_HOME`/`TORCH_CUDA_ARCH_LIST`, then `pip install -r requirements.txt --no-build-isolation`). Use `--cuda-home` and `--torch-cuda-arch-list` to override the CUDA path or target GPU architecture. + + On Windows, TRELLIS.2 does not require `flash-attn`: at startup it probes whether `flash_attn` is installed and actually runs on your GPU, and automatically falls back to PyTorch SDPA if it's missing or fails. If you install a working `flash-attn` build, it will be used automatically; you can still force a specific backend with `ATTN_BACKEND` or `SPARSE_ATTN_BACKEND` (`xformers`, `flash_attn`, `flash_attn_3`, `sdpa`, `naive`). + + This fork's `requirements.txt` pins `cumesh` and `flex_gemm` to exact commits of [rwfsmith/CuMesh](https://github.com/rwfsmith/CuMesh) and [rwfsmith/FlexGEMM](https://github.com/rwfsmith/FlexGEMM). Besides the Windows/CUDA-13 build fixes, those commits carry two **cross-platform correctness fixes** to CuMesh that are not specific to this fork's platform support: + + - **Out-of-bounds UDF read in the narrow-band dual contouring remesh kernel.** Caused `o_voxel.postprocess.to_glb(..., remesh=True)` (used by `extract_glb()` in `app.py`) to non-deterministically produce shredded/spiky meshes — identical code and latents could yield a clean or corrupted mesh from run to run. Upstreamed as [CuMesh#39](https://github.com/JeffreyXiang/CuMesh/pull/39). + - **Uninitialized CUB reduction identity in `fill_holes`.** `fill_holes` averages hole-cap vertices with `cub::DeviceSegmentedReduce::Sum` over `Vec3f`. CUB builds the reduction identity on the *host* as `InitT{}`, but `Vec3f` has a user-provided default constructor (so `Vec3f{}` calls it instead of zero-filling) that is marked `__device__`-only (so it isn't callable from host code). The identity was therefore raw host stack memory, folded into every segment sum. Because it depends on stack residue, the first generation in a process was usually clean and **every generation after it was corrupt**, smearing the model into a cylindrical column of sliver triangles — so it only showed up from the second image onward in the long-lived Gradio app. Upstreamed as [CuMesh#41](https://github.com/JeffreyXiang/CuMesh/pull/41). + + If you see corrupted meshes, verify you are building against these pinned commits rather than upstream `main`. + + This fork also fixes the DINOv3 image conditioning for Transformers 5.x. Recent releases moved the DINOv3 transformer blocks from `DINOv3ViTModel.layer` to `DINOv3ViTModel.model.layer`, so the original lookup missed them. Falling back to `last_hidden_state` is *not* equivalent: `DINOv3ViTModel.forward` applies the backbone's trained final `LayerNorm`, whose learned affine shifts the feature distribution TRELLIS.2 was conditioned on. The resulting out-of-distribution conditioning made the sparse-structure flow sampler collapse to an all-empty voxel grid on roughly 1 seed in 6, surfacing 300 lines downstream as a cryptic `RuntimeError: max(): Expected reduction dim to be specified for input.numel() == 0`. `DinoV3FeatureExtractor.extract_features` now locates the block list across Transformers versions, runs the blocks directly, and applies a parameter-free `F.layer_norm`. On a fixed 40-seed sweep this took empty-grid failures from 7/40 to 0/40 and tightened the generated voxel count from 139–2169 to 3299–4338, so it improves output fidelity on every seed, not just the ones that crashed. + +4. Log in to Hugging Face before running the pretrained model. TRELLIS.2 loads gated dependencies, including `facebook/dinov3-vitl16-pretrain-lvd1689m`, so your account must have access: + ```powershell + huggingface-cli login + ``` + +5. Run the full image-to-3D example: + ```powershell + python example.py + ``` + A successful run writes `sample.mp4` and `sample.glb`. + +#### Linux 1. Clone the repo: ```sh git clone -b main https://github.com/microsoft/TRELLIS.2.git --recursive @@ -99,9 +167,22 @@ Data processing is streamlined for instant conversions that are fully **renderin --nvdiffrec Install nvdiffrec ``` +## 🔁 CI/CD (Fork) + +This fork includes a GitHub Actions workflow at [.github/workflows/ci-cd.yml](.github/workflows/ci-cd.yml): + +- **CI / PR build**: runs on every PR targeting `main` (opened/synchronize/reopened/ready_for_review), plus direct pushes to `main`. +- **Checks performed**: validates dependency spec syntax in `requirements.txt` and `o-voxel/pyproject.toml`, then compiles Python sources (`compileall`) for a fast build smoke test. +- **CD snapshot on `main`**: after a successful `main` push build, uploads a `build-manifest.json` artifact with commit/run metadata. + +This fork also includes a tag-triggered release workflow at [.github/workflows/release.yml](.github/workflows/release.yml): + +- **Release trigger**: pushes of tags matching `v*` (for example `v1.0.0`), plus optional manual `workflow_dispatch`. +- **Release artifacts**: publishes `TRELLIS.2-.tar.gz`, `TRELLIS.2-.zip`, `checksums.sha256`, and a `build-manifest.json` file to the GitHub Release. + ## 📦 Pretrained Weights -The pretrained model **TRELLIS.2-4B** is available on Hugging Face. Please refer to the model card there for more details. +The pretrained model **TRELLIS.2-4B** is available on Hugging Face. Please refer to the model card there for more details. The image encoder dependency `facebook/dinov3-vitl16-pretrain-lvd1689m` is gated on Hugging Face; authenticate with an account that has access before running the example or web demo. | Model | Parameters | Resolution | Link | | :--- | :--- | :--- | :--- | diff --git a/o-voxel/pyproject.toml b/o-voxel/pyproject.toml index 6b13d432..3222b856 100644 --- a/o-voxel/pyproject.toml +++ b/o-voxel/pyproject.toml @@ -29,6 +29,6 @@ dependencies = [ "tqdm", "zstandard", "easydict", - "cumesh @ git+https://github.com/JeffreyXiang/CuMesh.git", - "flex_gemm @ git+https://github.com/JeffreyXiang/FlexGEMM.git", + "cumesh @ git+https://github.com/rwfsmith/CuMesh.git@683440c8d3959653384647e7bb7218dbd99eecca", + "flex_gemm @ git+https://github.com/rwfsmith/FlexGEMM.git@7598fe133f5b1425d5624b4c00e1f381412b96ce", ] diff --git a/o-voxel/setup.py b/o-voxel/setup.py index 91cb5cec..fa43b265 100644 --- a/o-voxel/setup.py +++ b/o-voxel/setup.py @@ -1,9 +1,11 @@ from setuptools import setup from torch.utils.cpp_extension import CUDAExtension, BuildExtension, IS_HIP_EXTENSION import os +import platform ROOT = os.path.dirname(os.path.abspath(__file__)) BUILD_TARGET = os.environ.get("BUILD_TARGET", "auto") +IS_WINDOWS = platform.system() == "Windows" if BUILD_TARGET == "auto": if IS_HIP_EXTENSION: @@ -22,6 +24,13 @@ archs = os.getenv("GPU_ARCHS", "native").split(";") cc_flag = [f"--offload-arch={arch}" for arch in archs] +if IS_WINDOWS: + cxx_flags = ["/O2", "/std:c++17", "/EHsc", "/permissive-", "/Zc:__cplusplus"] + nvcc_flags = ["-O3", "-std=c++17", "--expt-relaxed-constexpr", "--extended-lambda", "-allow-unsupported-compiler", "-Xcompiler=/std:c++17", "-Xcompiler=/EHsc", "-Xcompiler=/permissive-", "-Xcompiler=/Zc:__cplusplus"] + cc_flag +else: + cxx_flags = ["-O3", "-std=c++17"] + nvcc_flags = ["-O3", "-std=c++17"] + cc_flag + setup( name="o_voxel", packages=[ @@ -56,8 +65,8 @@ os.path.join(ROOT, "third_party/eigen"), ], extra_compile_args={ - "cxx": ["-O3", "-std=c++17"], - "nvcc": ["-O3","-std=c++17"] + cc_flag, + "cxx": cxx_flags, + "nvcc": nvcc_flags, } ) ], diff --git a/o-voxel/src/convert/flexible_dual_grid.cpp b/o-voxel/src/convert/flexible_dual_grid.cpp index ad89edc0..2e48adc3 100644 --- a/o-voxel/src/convert/flexible_dual_grid.cpp +++ b/o-voxel/src/convert/flexible_dual_grid.cpp @@ -304,7 +304,7 @@ void boundry_qef( // Calculate the QEF for the edge (boundary) defined by v0 and v1 Eigen::Vector3d dir(v1.x() - v0.x(), v1.y() - v0.y(), v1.z() - v0.z()); double segment_length = dir.norm(); - if (segment_length < 1e-6d) continue; // Skip degenerate edges (zero-length) + if (segment_length < 1e-6) continue; // Skip degenerate edges (zero-length) dir.normalize(); // unit direction vector // Projection matrix orthogonal to the direction: I - d d^T @@ -334,7 +334,7 @@ void boundry_qef( Eigen::Vector3d tMax, tDelta; for (int axis = 0; axis < 3; ++axis) { - if (dir[axis] == 0.0d) { + if (dir[axis] == 0.0) { tMax[axis] = std::numeric_limits::infinity(); tDelta[axis] = std::numeric_limits::infinity(); } else { @@ -772,4 +772,3 @@ std::tuple mesh_to_flexible_dual_gr torch::from_blob(intersected.data(), {int(intersected.size()), 3}, torch::kBool).clone() ); } - diff --git a/o-voxel/src/io/filter_neighbor.cpp b/o-voxel/src/io/filter_neighbor.cpp index 3af406eb..2687deef 100644 --- a/o-voxel/src/io/filter_neighbor.cpp +++ b/o-voxel/src/io/filter_neighbor.cpp @@ -77,7 +77,7 @@ torch::Tensor encode_sparse_voxel_octree_attr_neighbor_cpu( } // Pack the deltas into a uint8 tensor - torch::Tensor delta = torch::zeros({N, C}, torch::dtype(torch::kUInt8)); + torch::Tensor delta = torch::zeros({static_cast(N), static_cast(C)}, torch::dtype(torch::kUInt8)); uint8_t* delta_data = delta.data_ptr(); for (int i = 0; i < N; i++) { int x = coord_data[i * 3 + 0]; @@ -163,7 +163,7 @@ torch::Tensor decode_sparse_voxel_octree_attr_neighbor_cpu( } // Pack the attribute into a uint8 tensor - torch::Tensor attr = torch::zeros({N, C}, torch::dtype(torch::kUInt8)); + torch::Tensor attr = torch::zeros({static_cast(N), static_cast(C)}, torch::dtype(torch::kUInt8)); uint8_t* attr_data = attr.data_ptr(); for (int i = 0; i < N; i++) { int x = coord_data[i * 3 + 0]; diff --git a/o-voxel/src/io/filter_parent.cpp b/o-voxel/src/io/filter_parent.cpp index 0beadaf4..4a4a0e58 100644 --- a/o-voxel/src/io/filter_parent.cpp +++ b/o-voxel/src/io/filter_parent.cpp @@ -80,7 +80,7 @@ torch::Tensor encode_sparse_voxel_octree_attr_parent_cpu( uint8_t* octree_data = octree.data_ptr(); uint8_t* attr_data = attr.data_ptr(); - torch::Tensor delta = torch::zeros({N_leaf, C}, torch::kUInt8); + torch::Tensor delta = torch::zeros({static_cast(N_leaf), static_cast(C)}, torch::kUInt8); uint32_t svo_ptr = 0; uint32_t attr_ptr = 0; uint32_t delta_ptr = C; @@ -151,7 +151,7 @@ torch::Tensor decode_sparse_voxel_octree_attr_parent_cpu( uint8_t* octree_data = octree.data_ptr(); uint8_t* delta_data = delta.data_ptr(); - torch::Tensor attr = torch::zeros({N_leaf, C}, torch::kUInt8); + torch::Tensor attr = torch::zeros({static_cast(N_leaf), static_cast(C)}, torch::kUInt8); uint32_t svo_ptr = 0; uint32_t attr_ptr = 0; uint32_t delta_ptr = C; diff --git a/o-voxel/src/io/svo.cpp b/o-voxel/src/io/svo.cpp index 6775284d..83b53282 100644 --- a/o-voxel/src/io/svo.cpp +++ b/o-voxel/src/io/svo.cpp @@ -71,7 +71,7 @@ torch::Tensor encode_sparse_voxel_octree_cpu( } // Convert SVO to tensor - torch::Tensor svo_tensor = torch::from_blob(svo.data(), {svo.size()}, torch::kUInt8).clone(); + torch::Tensor svo_tensor = torch::from_blob(svo.data(), {static_cast(svo.size())}, torch::kUInt8).clone(); return svo_tensor; } @@ -133,6 +133,6 @@ torch::Tensor decode_sparse_voxel_octree_cpu( // Decode SVO into list of codes decode_sparse_voxel_octree_cpu_recursive(octree_data, depth, ptr, stack, codes); // Convert codes to tensor - torch::Tensor codes_tensor = torch::from_blob(codes.data(), {codes.size()}, torch::kInt32).clone(); + torch::Tensor codes_tensor = torch::from_blob(codes.data(), {static_cast(codes.size())}, torch::kInt32).clone(); return codes_tensor; } diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 00000000..d810e102 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,35 @@ +--extra-index-url https://download.pytorch.org/whl/cu130 + +torch==2.13.0+cu130 +torchvision==0.28.0+cu130 +numpy<2.3.0,>=2 +pillow +imageio +imageio-ffmpeg +tqdm +easydict +opencv-python-headless==4.12.0.88 +ninja +trimesh +transformers +gradio==6.0.1 +tensorboard +pandas +lpips +zstandard +kornia +timm +safetensors +huggingface_hub +requests +plyfile +matplotlib +triton-windows +utils3d @ git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057d642626897ec8 +nvdiffrast @ git+https://github.com/NVlabs/nvdiffrast.git@v0.4.0 +nvdiffrec_render @ git+https://github.com/JeffreyXiang/nvdiffrec.git@renderutils +# Forks carrying Windows/CUDA-13 build fixes plus cross-platform correctness +# fixes. Pinned to commit SHAs so a rebuild reproduces the validated build. +cumesh @ git+https://github.com/rwfsmith/CuMesh.git@683440c8d3959653384647e7bb7218dbd99eecca +flex_gemm @ git+https://github.com/rwfsmith/FlexGEMM.git@7598fe133f5b1425d5624b4c00e1f381412b96ce +./o-voxel diff --git a/setup_windows.bat b/setup_windows.bat new file mode 100644 index 00000000..d06bf6d6 --- /dev/null +++ b/setup_windows.bat @@ -0,0 +1,6 @@ +@echo off +REM Convenience wrapper so setup_windows.ps1 can be run from cmd.exe. +REM Any arguments are forwarded as-is, e.g.: +REM setup_windows.bat -Python .\venv\Scripts\python.exe -CudaHome "D:\CUDA\v13.0" +powershell -NoProfile -ExecutionPolicy Bypass -File "%~dp0setup_windows.ps1" %* +exit /b %ERRORLEVEL% diff --git a/setup_windows.ps1 b/setup_windows.ps1 new file mode 100644 index 00000000..1595f23f --- /dev/null +++ b/setup_windows.ps1 @@ -0,0 +1,38 @@ +param( + [string]$Python = "python", + [string]$CudaHome = "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.0", + [string]$TorchCudaArchList = "12.0" +) + +$ErrorActionPreference = "Stop" + +$vcvars = @( + "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvars64.bat", + "C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvars64.bat", + "C:\Program Files\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvars64.bat" +) | Where-Object { Test-Path $_ } | Select-Object -First 1 + +if (-not $vcvars) { + throw "Visual Studio 2022 C++ build tools were not found. Install the Desktop development with C++ workload." +} + +if (-not (Test-Path $CudaHome)) { + throw "CUDA toolkit path not found: $CudaHome" +} + +$pythonExe = (Get-Command $Python -ErrorAction Stop).Source +$requirements = Join-Path $PSScriptRoot "requirements.txt" + +$command = @( + "call `"$vcvars`" >nul", + "set `"DISTUTILS_USE_SDK=1`"", + "set `"MSSdk=1`"", + "set `"CUDA_HOME=$CudaHome`"", + "set `"TORCH_CUDA_ARCH_LIST=$TorchCudaArchList`"", + "`"$pythonExe`" -m pip install -r `"$requirements`" --no-build-isolation" +) -join " && " + +cmd.exe /c $command +if ($LASTEXITCODE -ne 0) { + throw "Windows setup failed with exit code $LASTEXITCODE" +} diff --git a/setup_windows.py b/setup_windows.py new file mode 100644 index 00000000..f6f53388 --- /dev/null +++ b/setup_windows.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 +"""Windows setup helper for TRELLIS.2. + +Locates the Visual Studio 2022 C++ build tools and CUDA Toolkit, then runs +``pip install -r requirements.txt --no-build-isolation`` with the environment +variables required to compile the native CUDA extensions (CuMesh, FlexGEMM, +o-voxel) on Windows/MSVC. + +This is a Python equivalent of setup_windows.ps1 / setup_windows.bat for +users who prefer not to invoke PowerShell or a batch file directly, e.g.: + + python setup_windows.py + python setup_windows.py --python .\\venv\\Scripts\\python.exe --cuda-home "D:\\CUDA\\v13.0" +""" + +import argparse +import os +import shutil +import subprocess +import sys +from pathlib import Path + +DEFAULT_CUDA_HOME = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.0" +DEFAULT_TORCH_CUDA_ARCH_LIST = "12.0" + +VCVARS_CANDIDATES = [ + r"C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvars64.bat", + r"C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvars64.bat", + r"C:\Program Files\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvars64.bat", +] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--python", default=sys.executable, help="Path to the Python interpreter to install into (default: current interpreter)") + parser.add_argument("--cuda-home", default=DEFAULT_CUDA_HOME, help=f"Path to the CUDA Toolkit install (default: {DEFAULT_CUDA_HOME})") + parser.add_argument("--torch-cuda-arch-list", default=DEFAULT_TORCH_CUDA_ARCH_LIST, help=f"Value for TORCH_CUDA_ARCH_LIST (default: {DEFAULT_TORCH_CUDA_ARCH_LIST})") + return parser.parse_args() + + +def find_vcvars() -> str: + for candidate in VCVARS_CANDIDATES: + if os.path.isfile(candidate): + return candidate + raise SystemExit( + "Visual Studio 2022 C++ build tools were not found. " + "Install the 'Desktop development with C++' workload." + ) + + +def resolve_python(python_arg: str) -> str: + resolved = shutil.which(python_arg) + if resolved is None: + raise SystemExit(f"Could not find python executable: {python_arg}") + return resolved + + +def main() -> int: + if os.name != "nt": + raise SystemExit("setup_windows.py is intended to be run on Windows.") + + args = parse_args() + + vcvars = find_vcvars() + + cuda_home = args.cuda_home + if not os.path.isdir(cuda_home): + raise SystemExit(f"CUDA toolkit path not found: {cuda_home}") + + python_exe = resolve_python(args.python) + requirements = Path(__file__).resolve().parent / "requirements.txt" + if not requirements.is_file(): + raise SystemExit(f"requirements.txt not found: {requirements}") + + # vcvars64.bat only sets environment variables for the cmd.exe process it + # runs in, so vcvars + pip install must happen in the same cmd.exe call. + command = " && ".join( + [ + f'call "{vcvars}" >nul', + 'set "DISTUTILS_USE_SDK=1"', + 'set "MSSdk=1"', + f'set "CUDA_HOME={cuda_home}"', + f'set "TORCH_CUDA_ARCH_LIST={args.torch_cuda_arch_list}"', + f'"{python_exe}" -m pip install -r "{requirements}" --no-build-isolation', + ] + ) + + print(f"Using Python: {python_exe}") + print(f"Using CUDA_HOME: {cuda_home}") + print(f"Using TORCH_CUDA_ARCH_LIST: {args.torch_cuda_arch_list}") + print(f"Using vcvars64.bat: {vcvars}") + + result = subprocess.run(["cmd.exe", "/c", command]) + if result.returncode != 0: + raise SystemExit(f"Windows setup failed with exit code {result.returncode}") + + print("Windows setup completed successfully.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/trellis2/modules/_attn_probe.py b/trellis2/modules/_attn_probe.py new file mode 100644 index 00000000..831f5439 --- /dev/null +++ b/trellis2/modules/_attn_probe.py @@ -0,0 +1,24 @@ +""" +Shared helper to detect whether flash_attn is actually usable on the current +GPU/CUDA build. Import succeeding is not sufficient: flash_attn can be +installed but fail at runtime on unsupported architectures or CUDA/torch +version mismatches. The result is cached so the (relatively expensive) CUDA +probe only runs once per process even though both the dense and sparse +attention config modules need it. +""" + +from functools import lru_cache + + +@lru_cache(maxsize=1) +def flash_attn_usable() -> bool: + try: + import torch + import flash_attn + if not torch.cuda.is_available(): + return False + q = torch.zeros(1, 1, 1, 8, dtype=torch.float16, device='cuda') + flash_attn.flash_attn_func(q, q, q) + return True + except Exception: + return False diff --git a/trellis2/modules/attention/config.py b/trellis2/modules/attention/config.py index a6d5180c..655c9ba5 100644 --- a/trellis2/modules/attention/config.py +++ b/trellis2/modules/attention/config.py @@ -1,6 +1,15 @@ from typing import * +import platform +from .._attn_probe import flash_attn_usable -BACKEND = 'flash_attn' + +def _default_backend() -> str: + if platform.system() == 'Windows': + return 'flash_attn' if flash_attn_usable() else 'sdpa' + return 'flash_attn' + + +BACKEND = _default_backend() DEBUG = False def __from_env(): diff --git a/trellis2/modules/image_feature_extractor.py b/trellis2/modules/image_feature_extractor.py index c3cb515a..06992185 100644 --- a/trellis2/modules/image_feature_extractor.py +++ b/trellis2/modules/image_feature_extractor.py @@ -78,17 +78,36 @@ def cuda(self): def cpu(self): self.model.cpu() + def _encoder_layers(self): + """ + Locate the transformer block list across transformers versions. + + Older releases exposed the blocks directly on ``DINOv3ViTModel.layer``; current + releases nest them inside the encoder submodule (``DINOv3ViTModel.model.layer``). + """ + for holder in (self.model, getattr(self.model, 'model', None), getattr(self.model, 'encoder', None)): + if holder is not None and hasattr(holder, 'layer'): + return holder.layer + return None + def extract_features(self, image: torch.Tensor) -> torch.Tensor: image = image.to(self.model.embeddings.patch_embeddings.weight.dtype) + layers = self._encoder_layers() + if layers is None: + return self.model(pixel_values=image, return_dict=True).last_hidden_state + hidden_states = self.model.embeddings(image, bool_masked_pos=None) position_embeddings = self.model.rope_embeddings(image) - - for i, layer_module in enumerate(self.model.layer): + for layer_module in layers: hidden_states = layer_module( hidden_states, position_embeddings=position_embeddings, ) + # Parameter-free normalization of the pre-norm states. Do not substitute the + # backbone's trained final LayerNorm (``self.model.norm``): its learned affine + # shifts the feature distribution the flow models were conditioned on, which + # makes the sparse-structure sampler collapse to an empty grid on some seeds. return F.layer_norm(hidden_states, hidden_states.shape[-1:]) @torch.no_grad() diff --git a/trellis2/modules/sparse/attention/full_attn.py b/trellis2/modules/sparse/attention/full_attn.py index 4eb74d24..834a86a4 100644 --- a/trellis2/modules/sparse/attention/full_attn.py +++ b/trellis2/modules/sparse/attention/full_attn.py @@ -9,6 +9,20 @@ ] +def _sdpa_varlen(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_seqlen: List[int], kv_seqlen: List[int]) -> torch.Tensor: + out = [] + q_start = 0 + kv_start = 0 + for q_len, kv_len in zip(q_seqlen, kv_seqlen): + q_slice = q[q_start:q_start + q_len].transpose(0, 1).unsqueeze(0) + k_slice = k[kv_start:kv_start + kv_len].transpose(0, 1).unsqueeze(0) + v_slice = v[kv_start:kv_start + kv_len].transpose(0, 1).unsqueeze(0) + out.append(torch.nn.functional.scaled_dot_product_attention(q_slice, k_slice, v_slice).squeeze(0).transpose(0, 1)) + q_start += q_len + kv_start += kv_len + return torch.cat(out, dim=0) + + @overload def sparse_scaled_dot_product_attention(qkv: VarLenTensor) -> VarLenTensor: """ @@ -211,6 +225,12 @@ def sparse_scaled_dot_product_attention(*args, **kwargs): max_q_seqlen = max(q_seqlen) max_kv_seqlen = max(kv_seqlen) out = flash_attn_3.flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_q_seqlen, max_kv_seqlen) + elif config.ATTN == 'sdpa': + if num_all_args == 1: + q, k, v = qkv.unbind(dim=1) + elif num_all_args == 2: + k, v = kv.unbind(dim=1) + out = _sdpa_varlen(q, k, v, q_seqlen, kv_seqlen) else: raise ValueError(f"Unknown attention module: {config.ATTN}") diff --git a/trellis2/modules/sparse/attention/windowed_attn.py b/trellis2/modules/sparse/attention/windowed_attn.py index 04307889..dc28832a 100644 --- a/trellis2/modules/sparse/attention/windowed_attn.py +++ b/trellis2/modules/sparse/attention/windowed_attn.py @@ -11,6 +11,22 @@ ] +def _sdpa_varlen(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_seq_lens: torch.Tensor, kv_seq_lens: torch.Tensor) -> torch.Tensor: + out = [] + q_start = 0 + kv_start = 0 + for q_len_tensor, kv_len_tensor in zip(q_seq_lens, kv_seq_lens): + q_len = int(q_len_tensor.item()) + kv_len = int(kv_len_tensor.item()) + q_slice = q[q_start:q_start + q_len].transpose(0, 1).unsqueeze(0) + k_slice = k[kv_start:kv_start + kv_len].transpose(0, 1).unsqueeze(0) + v_slice = v[kv_start:kv_start + kv_len].transpose(0, 1).unsqueeze(0) + out.append(torch.nn.functional.scaled_dot_product_attention(q_slice, k_slice, v_slice).squeeze(0).transpose(0, 1)) + q_start += q_len + kv_start += kv_len + return torch.cat(out, dim=0) + + def calc_window_partition( tensor: SparseTensor, window_size: Union[int, Tuple[int, ...]], @@ -60,6 +76,8 @@ def calc_window_partition( 'cu_seqlens': torch.cat([torch.tensor([0], device=tensor.device), torch.cumsum(seq_lens, dim=0)], dim=0).int(), 'max_seqlen': torch.max(seq_lens) } + elif config.ATTN == 'sdpa': + attn_func_args = {} return fwd_indices, bwd_indices, seq_lens, attn_func_args @@ -113,6 +131,9 @@ def sparse_windowed_scaled_dot_product_self_attention( if 'flash_attn' not in globals(): import flash_attn out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, **attn_func_args) # [M, H, C] + elif config.ATTN == 'sdpa': + q, k, v = qkv_feats.unbind(dim=1) + out = _sdpa_varlen(q, k, v, seq_lens, seq_lens) out = out[bwd_indices] # [T, H, C] @@ -184,6 +205,9 @@ def sparse_windowed_scaled_dot_product_cross_attention( cu_seqlens_q=q_attn_func_args['cu_seqlens'], cu_seqlens_k=kv_attn_func_args['cu_seqlens'], max_seqlen_q=q_attn_func_args['max_seqlen'], max_seqlen_k=kv_attn_func_args['max_seqlen'], ) # [M, H, C] + elif config.ATTN == 'sdpa': + k, v = kv_feats.unbind(dim=1) + out = _sdpa_varlen(q_feats, k, v, q_seq_lens, kv_seq_lens) out = out[q_bwd_indices] # [T, H, C] diff --git a/trellis2/modules/sparse/config.py b/trellis2/modules/sparse/config.py index a5f4d532..c11b4620 100644 --- a/trellis2/modules/sparse/config.py +++ b/trellis2/modules/sparse/config.py @@ -1,8 +1,18 @@ from typing import * +import platform +from .._attn_probe import flash_attn_usable CONV = 'flex_gemm' DEBUG = False -ATTN = 'flash_attn' + + +def _default_attn_backend() -> str: + if platform.system() == 'Windows': + return 'flash_attn' if flash_attn_usable() else 'sdpa' + return 'flash_attn' + + +ATTN = _default_attn_backend() def __from_env(): import os @@ -21,7 +31,7 @@ def __from_env(): CONV = env_sparse_conv_backend if env_sparse_debug is not None: DEBUG = env_sparse_debug == '1' - if env_sparse_attn_backend is not None and env_sparse_attn_backend in ['xformers', 'flash_attn', 'flash_attn_3']: + if env_sparse_attn_backend is not None and env_sparse_attn_backend in ['xformers', 'flash_attn', 'flash_attn_3', 'sdpa']: ATTN = env_sparse_attn_backend print(f"[SPARSE] Conv backend: {CONV}; Attention backend: {ATTN}") @@ -38,6 +48,6 @@ def set_debug(debug: bool): global DEBUG DEBUG = debug -def set_attn_backend(backend: Literal['xformers', 'flash_attn']): +def set_attn_backend(backend: Literal['xformers', 'flash_attn', 'flash_attn_3', 'sdpa']): global ATTN ATTN = backend diff --git a/trellis2/pipelines/trellis2_image_to_3d.py b/trellis2/pipelines/trellis2_image_to_3d.py index a7b84b98..a318ad29 100644 --- a/trellis2/pipelines/trellis2_image_to_3d.py +++ b/trellis2/pipelines/trellis2_image_to_3d.py @@ -231,6 +231,12 @@ def sample_sparse_structure( ratio = decoded.shape[2] // resolution decoded = torch.nn.functional.max_pool3d(decoded.float(), ratio, ratio, 0) > 0.5 coords = torch.argwhere(decoded)[:, [0, 2, 3, 4]].int() + if coords.shape[0] == 0: + raise RuntimeError( + "The sparse structure sampler produced an empty voxel grid, so there is no " + "geometry to generate. Try a different seed, or an input image with a clearly " + "segmented foreground object." + ) return coords diff --git a/trellis2/trainers/flow_matching/mixins/image_conditioned.py b/trellis2/trainers/flow_matching/mixins/image_conditioned.py index 13b932eb..5e4195db 100644 --- a/trellis2/trainers/flow_matching/mixins/image_conditioned.py +++ b/trellis2/trainers/flow_matching/mixins/image_conditioned.py @@ -80,17 +80,36 @@ def cuda(self): def cpu(self): self.model.cpu() + def _encoder_layers(self): + """ + Locate the transformer block list across transformers versions. + + Older releases exposed the blocks directly on ``DINOv3ViTModel.layer``; current + releases nest them inside the encoder submodule (``DINOv3ViTModel.model.layer``). + """ + for holder in (self.model, getattr(self.model, 'model', None), getattr(self.model, 'encoder', None)): + if holder is not None and hasattr(holder, 'layer'): + return holder.layer + return None + def extract_features(self, image: torch.Tensor) -> torch.Tensor: image = image.to(self.model.embeddings.patch_embeddings.weight.dtype) + layers = self._encoder_layers() + if layers is None: + return self.model(pixel_values=image, return_dict=True).last_hidden_state + hidden_states = self.model.embeddings(image, bool_masked_pos=None) position_embeddings = self.model.rope_embeddings(image) - - for i, layer_module in enumerate(self.model.layer): + for layer_module in layers: hidden_states = layer_module( hidden_states, position_embeddings=position_embeddings, ) + # Parameter-free normalization of the pre-norm states. Do not substitute the + # backbone's trained final LayerNorm (``self.model.norm``): its learned affine + # shifts the feature distribution the flow models were conditioned on, which + # makes the sparse-structure sampler collapse to an empty grid on some seeds. return F.layer_norm(hidden_states, hidden_states.shape[-1:]) @torch.no_grad()