106 lines
3.9 KiB
Python
106 lines
3.9 KiB
Python
"""GPU validation for the TRELLIS.2 HiCache patch.
|
|
|
|
Loads a real stock TRELLIS.2 pipeline, applies the model-level forecast patch to
|
|
the sparse-structure + shape-SLaT flow DiTs, runs image-to-3D, and confirms it
|
|
runs, skips DiT forwards, is faster, and produces near-identical mesh geometry.
|
|
|
|
Run from a hermit-trellis2 checkout (which vendors the trellis2 package):
|
|
|
|
cd third_party/hermit-trellis2
|
|
PYTHONPATH=.:../comfyui-trellis2-hicache python ../comfyui-trellis2-hicache/tests/validate_gpu.py \
|
|
--ckpt /home/krishi/workspace/data/weights/TRELLIS.2-4B --image <some.png>
|
|
"""
|
|
import argparse
|
|
import sys
|
|
import time
|
|
import types
|
|
|
|
|
|
def _stub_render_deps():
|
|
for m in ("nvdiffrast", "nvdiffrast.torch", "nvdiffrec"):
|
|
sys.modules.setdefault(m, types.ModuleType(m))
|
|
sys.modules["nvdiffrast"].torch = sys.modules["nvdiffrast.torch"]
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--ckpt", required=True)
|
|
ap.add_argument("--image", required=True)
|
|
ap.add_argument("--method", default="hermite")
|
|
ap.add_argument("--interval", type=int, default=2)
|
|
ap.add_argument("--stages", default="both")
|
|
ap.add_argument("--pipeline_type", default="shape")
|
|
args = ap.parse_args()
|
|
|
|
_stub_render_deps()
|
|
import torch
|
|
from PIL import Image
|
|
from trellis2.pipelines import Trellis2ImageTo3DPipeline
|
|
from trellis_hicache_patch import apply_hicache, remove_hicache
|
|
|
|
pipe = Trellis2ImageTo3DPipeline.from_pretrained(args.ckpt)
|
|
pipe.to("cuda")
|
|
pipe.enable_faster("base") # stock samplers; we test the model-level patch
|
|
img = Image.open(args.image).convert("RGBA")
|
|
kw = dict(seed=0, preprocess_image=True, pipeline_type=args.pipeline_type)
|
|
|
|
def verts(mesh):
|
|
return mesh.vertices.detach().float()
|
|
|
|
torch.cuda.synchronize()
|
|
t0 = time.time()
|
|
m_stock = pipe.run(img, **kw)[0]
|
|
torch.cuda.synchronize()
|
|
t_stock = time.time() - t0
|
|
v_stock = verts(m_stock)
|
|
|
|
patched = apply_hicache(pipe, method=args.method, interval=args.interval, stages=args.stages)
|
|
torch.cuda.synchronize()
|
|
t0 = time.time()
|
|
m_fast = patched.run(img, **kw)[0]
|
|
torch.cuda.synchronize()
|
|
t_fast = time.time() - t0
|
|
v_fast = verts(m_fast)
|
|
|
|
def chamfer(a, b, k=20000, chunk_size=2048):
|
|
if a.shape[0] > k:
|
|
a = a[torch.randperm(a.shape[0], device=a.device)[:k]]
|
|
if b.shape[0] > k:
|
|
b = b[torch.randperm(b.shape[0], device=b.device)[:k]]
|
|
|
|
def directed_min_mean(src, dst):
|
|
mins = []
|
|
for start in range(0, src.shape[0], chunk_size):
|
|
mins.append(torch.cdist(src[start:start + chunk_size], dst)
|
|
.min(1).values)
|
|
return torch.cat(mins).mean()
|
|
|
|
# Keep the validator's peak distance matrix bounded instead of creating
|
|
# a 20,000 x 20,000 allocation (~1.6 GB in float32).
|
|
return float((directed_min_mean(a, b) + directed_min_mean(b, a)) / 2)
|
|
cham = chamfer(v_stock, v_fast)
|
|
|
|
total_skipped = 0
|
|
print("\n==================== TRELLIS.2-HiCache validation ====================")
|
|
print(f"method={args.method} interval={args.interval} stages={args.stages}")
|
|
for key, m in patched.models.items():
|
|
if getattr(m, "_hicache_is_patch", False):
|
|
c, s = m.computed_steps, m.skipped_steps
|
|
total_skipped += s
|
|
if c or s:
|
|
print(f" {key}: {c} computed + {s} skipped")
|
|
print(f"stock : {t_stock:.2f}s ({v_stock.shape[0]} verts)")
|
|
print(f"hicache: {t_fast:.2f}s ({v_fast.shape[0]} verts)")
|
|
print(f"speedup: {t_stock / max(t_fast, 1e-6):.2f}x")
|
|
print(f"chamfer(stock,hicache) = {cham:.5f}")
|
|
|
|
restored = remove_hicache(patched)
|
|
clean = not any(getattr(m, "_hicache_is_patch", False) for m in restored.models.values())
|
|
ok = total_skipped > 0 and t_fast < t_stock and cham < 0.02 and clean
|
|
print("RESULT:", "PASS" if ok else "FAIL")
|
|
sys.exit(0 if ok else 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|