321 lines
13 KiB
Python
321 lines
13 KiB
Python
"""Shared plumbing for the PyTorch reference dumps.
|
|
|
|
The reference environment is deliberately stock PyTorch (docker/Dockerfile.ref):
|
|
none of the custom CUDA extensions (flash-attn, FlexGEMM, cumesh, o-voxel) are
|
|
installed. This module makes the trellis2 package importable and usable in that
|
|
environment:
|
|
|
|
* stubs out `cumesh` (only needed by Mesh postprocess methods we don't call
|
|
while dumping activations),
|
|
* replaces the sparse attention dispatcher with a plain-SDPA implementation
|
|
(mathematically identical for the var-len batch-1 case we validate),
|
|
* registers a pure-PyTorch submanifold sparse-conv backend under the name
|
|
'none' (gather + GEMM per kernel offset — slow, but it runs anywhere and
|
|
doubles as the executable spec for the C++ implementation).
|
|
|
|
Import this before importing anything from `trellis2`.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import types
|
|
|
|
TRELLIS2_PY = os.environ.get("TRELLIS2_PY", "/trellis2")
|
|
if TRELLIS2_PY not in sys.path:
|
|
sys.path.insert(0, TRELLIS2_PY)
|
|
|
|
os.environ.setdefault("ATTN_BACKEND", "sdpa")
|
|
os.environ.setdefault("SPARSE_CONV_BACKEND", "none")
|
|
# The sparse attention config only accepts xformers/flash_attn/flash_attn_3;
|
|
# we leave it alone and monkeypatch the dispatcher below instead.
|
|
|
|
# --- cumesh stub (postprocess-only dependency of representations.mesh.base) ---
|
|
if "cumesh" not in sys.modules:
|
|
stub = types.ModuleType("cumesh")
|
|
|
|
class _CuMeshStub:
|
|
def __init__(self, *_a, **_k):
|
|
raise RuntimeError("cumesh is stubbed out in the reference container")
|
|
|
|
stub.CuMesh = _CuMeshStub
|
|
sys.modules["cumesh"] = stub
|
|
|
|
# --- o_voxel stub (CUDA hashmap mesher; we dump the decoder's raw 7-channel
|
|
# --- output and do mesh comparison with scripts/ref_dual_grid.py instead) ---
|
|
if "o_voxel" not in sys.modules:
|
|
ovx = types.ModuleType("o_voxel")
|
|
ovx_convert = types.ModuleType("o_voxel.convert")
|
|
|
|
def flexible_dual_grid_to_mesh(*_a, **_k):
|
|
raise RuntimeError("o_voxel is stubbed out in the reference container")
|
|
|
|
ovx_convert.flexible_dual_grid_to_mesh = flexible_dual_grid_to_mesh
|
|
ovx.convert = ovx_convert
|
|
sys.modules["o_voxel"] = ovx
|
|
sys.modules["o_voxel.convert"] = ovx_convert
|
|
|
|
# --- flex_gemm stub (CUDA kernels; representations.mesh.base imports
|
|
# --- grid_sample_3d for texture baking, unused on the geometry path) ---
|
|
if "flex_gemm" not in sys.modules:
|
|
fg = types.ModuleType("flex_gemm")
|
|
fg_ops = types.ModuleType("flex_gemm.ops")
|
|
fg_gs = types.ModuleType("flex_gemm.ops.grid_sample")
|
|
fg_sp = types.ModuleType("flex_gemm.ops.spconv")
|
|
|
|
def _fg_unavailable(*_a, **_k):
|
|
raise RuntimeError("flex_gemm is stubbed out in the reference container")
|
|
|
|
fg_gs.grid_sample_3d = _fg_unavailable
|
|
fg_sp.sparse_submanifold_conv3d = _fg_unavailable
|
|
fg_ops.grid_sample = fg_gs
|
|
fg_ops.spconv = fg_sp
|
|
fg.ops = fg_ops
|
|
sys.modules["flex_gemm"] = fg
|
|
sys.modules["flex_gemm.ops"] = fg_ops
|
|
sys.modules["flex_gemm.ops.grid_sample"] = fg_gs
|
|
sys.modules["flex_gemm.ops.spconv"] = fg_sp
|
|
|
|
|
|
def _install_sdpa_sparse_attention():
|
|
"""Replace trellis2's sparse attention with a dense-SDPA equivalent."""
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from trellis2.modules.sparse.attention import full_attn
|
|
from trellis2.modules.sparse import VarLenTensor
|
|
|
|
def sdpa_varlen(q, k, v, q_seqlen, kv_seqlen):
|
|
# q: [Tq, H, C], k/v: [Tkv, H, C] concatenated over batch. Query-chunked
|
|
# so the [H, chunk, L] score matrix stays bounded — the math SDPA backend
|
|
# (forced for true fp32) would otherwise OOM at the cascade's HR token
|
|
# counts. Chunking queries is mathematically exact (each query's softmax
|
|
# is independent), so the golden values are unchanged.
|
|
CHUNK = 2048
|
|
out = torch.empty_like(q)
|
|
qo = ko = 0
|
|
for ql, kl in zip(q_seqlen, kv_seqlen):
|
|
ks = k[ko:ko + kl].transpose(0, 1).unsqueeze(0) # [1,H,kl,C]
|
|
vs = v[ko:ko + kl].transpose(0, 1).unsqueeze(0)
|
|
for s in range(0, ql, CHUNK):
|
|
e = min(s + CHUNK, ql)
|
|
qs = q[qo + s:qo + e].transpose(0, 1).unsqueeze(0) # [1,H,chunk,C]
|
|
o = F.scaled_dot_product_attention(qs, ks, vs)
|
|
out[qo + s:qo + e] = o.squeeze(0).transpose(0, 1)
|
|
qo += ql
|
|
ko += kl
|
|
return out
|
|
|
|
def sparse_sdpa(*args, **kwargs):
|
|
num = len(args) + len(kwargs)
|
|
if num == 1:
|
|
qkv = args[0] if args else kwargs["qkv"]
|
|
assert isinstance(qkv, VarLenTensor)
|
|
q_seqlen = [qkv.layout[i].stop - qkv.layout[i].start for i in range(qkv.shape[0])]
|
|
q, k, v = qkv.feats.unbind(dim=1) # [T,3,H,C] -> 3x[T,H,C]
|
|
out = sdpa_varlen(q, k, v, q_seqlen, q_seqlen)
|
|
return qkv.replace(out)
|
|
if num == 2:
|
|
q = args[0] if len(args) > 0 else kwargs["q"]
|
|
kv = args[1] if len(args) > 1 else kwargs["kv"]
|
|
s = q if isinstance(q, VarLenTensor) else None
|
|
if isinstance(q, VarLenTensor):
|
|
q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])]
|
|
qf = q.feats
|
|
else:
|
|
N, L = q.shape[:2]
|
|
q_seqlen = [L] * N
|
|
qf = q.reshape(N * L, *q.shape[2:])
|
|
if isinstance(kv, VarLenTensor):
|
|
kv_seqlen = [kv.layout[i].stop - kv.layout[i].start for i in range(kv.shape[0])]
|
|
kvf = kv.feats
|
|
else:
|
|
N, L = kv.shape[:2]
|
|
kv_seqlen = [L] * N
|
|
kvf = kv.reshape(N * L, *kv.shape[2:])
|
|
k, v = kvf.unbind(dim=1)
|
|
out = sdpa_varlen(qf, k, v, q_seqlen, kv_seqlen)
|
|
if s is not None:
|
|
return s.replace(out)
|
|
N = len(q_seqlen)
|
|
return out.reshape(N, q_seqlen[0], *out.shape[1:])
|
|
if num == 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"]
|
|
s = q if isinstance(q, VarLenTensor) else None
|
|
if isinstance(q, VarLenTensor):
|
|
q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])]
|
|
qf = q.feats
|
|
else:
|
|
N, L = q.shape[:2]
|
|
q_seqlen = [L] * N
|
|
qf = q.reshape(N * L, *q.shape[2:])
|
|
if isinstance(k, VarLenTensor):
|
|
kv_seqlen = [k.layout[i].stop - k.layout[i].start for i in range(k.shape[0])]
|
|
kf, vf = k.feats, v.feats
|
|
else:
|
|
N, L = k.shape[:2]
|
|
kv_seqlen = [L] * N
|
|
kf = k.reshape(N * L, *k.shape[2:])
|
|
vf = v.reshape(N * L, *v.shape[2:])
|
|
out = sdpa_varlen(qf, kf, vf, q_seqlen, kv_seqlen)
|
|
if s is not None:
|
|
return s.replace(out)
|
|
N = len(q_seqlen)
|
|
return out.reshape(N, q_seqlen[0], *out.shape[1:])
|
|
raise AssertionError("bad arg count")
|
|
|
|
full_attn.sparse_scaled_dot_product_attention = sparse_sdpa
|
|
# modules.py imported the symbol by value; patch it there too.
|
|
from trellis2.modules.sparse.attention import modules as attn_modules
|
|
attn_modules.sparse_scaled_dot_product_attention = sparse_sdpa
|
|
|
|
|
|
def _install_torch_sparse_conv():
|
|
"""Register a pure-PyTorch submanifold conv backend as 'none'."""
|
|
import math
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from trellis2.modules.sparse.conv import conv as conv_dispatch
|
|
|
|
mod = types.ModuleType("trellis2.modules.sparse.conv.conv_none")
|
|
|
|
def sparse_conv3d_init(self, in_channels, out_channels, kernel_size,
|
|
stride=1, dilation=1, padding=None, bias=True,
|
|
indice_key=None):
|
|
assert stride == 1 and padding is None, "submanifold only"
|
|
self.in_channels = in_channels
|
|
self.out_channels = out_channels
|
|
ks = tuple(kernel_size) if isinstance(kernel_size, (list, tuple)) else (kernel_size,) * 3
|
|
self.kernel_size = ks
|
|
self.stride = (1, 1, 1)
|
|
self.dilation = tuple(dilation) if isinstance(dilation, (list, tuple)) else (dilation,) * 3
|
|
# flex_gemm weight layout: (Co, Kd, Kh, Kw, Ci)
|
|
self.weight = nn.Parameter(torch.empty(out_channels, *ks, in_channels))
|
|
if bias:
|
|
self.bias = nn.Parameter(torch.zeros(out_channels))
|
|
else:
|
|
self.register_parameter("bias", None)
|
|
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
|
|
|
def _coord_key(coords, spatial_shape):
|
|
# coords: [N,4] int (b,x,y,z) -> int64 linear key
|
|
b, x, y, z = coords.unbind(-1)
|
|
sx, sy, sz = spatial_shape
|
|
return ((b.long() * sx + x.long()) * sy + y.long()) * sz + z.long()
|
|
|
|
def sparse_conv3d_forward(self, x):
|
|
coords = x.coords
|
|
feats = x.feats
|
|
n = feats.shape[0]
|
|
spatial = tuple(x.spatial_shape)
|
|
keys = _coord_key(coords, spatial)
|
|
order = torch.argsort(keys)
|
|
keys_sorted = keys[order]
|
|
|
|
Co, Kd, Kh, Kw, Ci = self.weight.shape
|
|
w = self.weight
|
|
out = feats.new_zeros(n, Co)
|
|
if self.bias is not None:
|
|
out += self.bias.to(out.dtype)
|
|
rd, rh, rw = Kd // 2, Kh // 2, Kw // 2
|
|
dd, dh, dw = self.dilation
|
|
for kd in range(Kd):
|
|
for kh in range(Kh):
|
|
for kw in range(Kw):
|
|
off = coords.new_tensor([0, (kd - rd) * dd, (kh - rh) * dh, (kw - rw) * dw])
|
|
ncoords = coords + off
|
|
inb = ((ncoords[:, 1] >= 0) & (ncoords[:, 1] < spatial[0]) &
|
|
(ncoords[:, 2] >= 0) & (ncoords[:, 2] < spatial[1]) &
|
|
(ncoords[:, 3] >= 0) & (ncoords[:, 3] < spatial[2]))
|
|
nkeys = _coord_key(ncoords, spatial)
|
|
pos = torch.searchsorted(keys_sorted, nkeys)
|
|
pos_c = pos.clamp(max=n - 1)
|
|
hit = inb & (keys_sorted[pos_c] == nkeys)
|
|
src = order[pos_c[hit]]
|
|
# out[i] += feats[neighbor(i, offset)] @ w[:, kd, kh, kw, :]^T
|
|
contrib = feats[src] @ w[:, kd, kh, kw, :].to(feats.dtype).t()
|
|
out[hit] += contrib
|
|
return x.replace(out)
|
|
|
|
def sparse_inverse_conv3d_init(self, *a, **k):
|
|
raise NotImplementedError
|
|
|
|
def sparse_inverse_conv3d_forward(self, x):
|
|
raise NotImplementedError
|
|
|
|
mod.sparse_conv3d_init = sparse_conv3d_init
|
|
mod.sparse_conv3d_forward = sparse_conv3d_forward
|
|
mod.sparse_inverse_conv3d_init = sparse_inverse_conv3d_init
|
|
mod.sparse_inverse_conv3d_forward = sparse_inverse_conv3d_forward
|
|
conv_dispatch._backends["none"] = mod
|
|
from trellis2.modules import sparse as sp
|
|
sp.config.CONV = "none"
|
|
|
|
|
|
def _force_true_fp32():
|
|
"""Make CUDA math bit-comparable to a real fp32 reference.
|
|
|
|
PyTorch's default CUDA matmul/attention uses TF32 (≈10-bit mantissa) and
|
|
flash/mem-efficient SDPA kernels that accumulate in reduced precision —
|
|
that shows up as ~1e-3 relative error versus a true fp32 (or ggml-CPU)
|
|
forward, which would otherwise masquerade as a port bug. Force full-width
|
|
fp32 so the golden dumps are the real reference regardless of --device.
|
|
"""
|
|
import torch
|
|
torch.backends.cuda.matmul.allow_tf32 = False
|
|
torch.backends.cudnn.allow_tf32 = False
|
|
if hasattr(torch.backends.cuda, "enable_flash_sdp"):
|
|
torch.backends.cuda.enable_flash_sdp(False)
|
|
torch.backends.cuda.enable_mem_efficient_sdp(False)
|
|
torch.backends.cuda.enable_math_sdp(True)
|
|
|
|
|
|
def setup():
|
|
_force_true_fp32()
|
|
_install_sdpa_sparse_attention()
|
|
_install_torch_sparse_conv()
|
|
|
|
|
|
REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
MODELS = os.path.join(REPO, "models")
|
|
DUMPS = os.path.join(REPO, "dumps")
|
|
|
|
|
|
def write_dinodata(path, arr):
|
|
""".dinodata: DINOCOND | u32 version | u32 dtype(0=f32) | u32 ndim | dims | f32 payload"""
|
|
import struct
|
|
import numpy as np
|
|
arr = np.ascontiguousarray(arr, dtype="<f4")
|
|
with open(path, "wb") as f:
|
|
f.write(b"DINOCOND")
|
|
f.write(struct.pack("<III", 1, 0, arr.ndim))
|
|
f.write(struct.pack("<%dI" % arr.ndim, *arr.shape))
|
|
f.write(arr.tobytes())
|
|
|
|
|
|
def preprocess_rgba(img):
|
|
"""pipeline.preprocess_image for an RGBA input (no rembg): downscale to
|
|
<=1024, alpha bbox square crop, premultiply onto black. Returns RGB PIL."""
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
assert img.mode == "RGBA", "fixture must have an alpha channel"
|
|
max_size = max(img.size)
|
|
scale = min(1, 1024 / max_size)
|
|
if scale < 1:
|
|
img = img.resize((int(img.width * scale), int(img.height * scale)),
|
|
Image.Resampling.LANCZOS)
|
|
out = np.array(img)
|
|
alpha = out[:, :, 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)
|
|
bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2
|
|
img = img.crop(bbox)
|
|
out = np.array(img).astype(np.float32) / 255
|
|
out = out[:, :, :3] * out[:, :, 3:4]
|
|
return Image.fromarray((out * 255).astype(np.uint8))
|