Kelly Yang
51f4675c88
Integrates TRI-ML RecGen into ComfyUI with 2D RGB-D to mesh/splat nodes, macOS spconv fallback, smoke-test helpers, and project metadata for kaili-yang/ComfyUI-Recgen.
116 lines
3.6 KiB
Python
116 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Local end-to-end smoke demo for ComfyUI-Recgen (no ComfyUI server required)."""
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Redirect HuggingFace / torch caches to external drive if available
|
|
EXTERNAL_DRIVE = Path("/Volumes/HD-SL2")
|
|
if EXTERNAL_DRIVE.exists():
|
|
external_cache = EXTERNAL_DRIVE / ".cache"
|
|
(external_cache / "huggingface").mkdir(parents=True, exist_ok=True)
|
|
(external_cache / "torch").mkdir(parents=True, exist_ok=True)
|
|
os.environ["HF_HOME"] = str(external_cache / "huggingface")
|
|
os.environ["TORCH_HOME"] = str(external_cache / "torch")
|
|
print(f"[ComfyUI-Recgen Test] Redirected caches to: {external_cache}")
|
|
|
|
import torch
|
|
|
|
WORKSPACE_DIR = Path(__file__).resolve().parent.parent.parent
|
|
sys.path.insert(0, str(WORKSPACE_DIR))
|
|
sys.path.insert(0, str(WORKSPACE_DIR / "ComfyUI-Recgen" / "src"))
|
|
|
|
os.environ.setdefault("SPCONV_ALGO", "native")
|
|
os.environ.setdefault("SPARSE_ATTN_BACKEND", "sdpa")
|
|
os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.0")
|
|
|
|
try:
|
|
import certifi
|
|
|
|
os.environ.setdefault("SSL_CERT_FILE", certifi.where())
|
|
os.environ.setdefault("REQUESTS_CA_BUNDLE", certifi.where())
|
|
except ImportError:
|
|
pass
|
|
|
|
from comfyui_recgen.nodes import RecGen2DTo3DNode
|
|
from comfyui_recgen.smoke_inputs import load_example_smoke_tensors
|
|
|
|
SMOKE_SIZE = 256
|
|
|
|
|
|
def _pick_device() -> str:
|
|
if torch.cuda.is_available():
|
|
return "cuda"
|
|
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
|
return "mps"
|
|
return "cpu"
|
|
|
|
|
|
def _print_mps_status() -> None:
|
|
mps_mod = getattr(torch.backends, "mps", None)
|
|
if mps_mod is None:
|
|
print("PyTorch MPS backend: not available in this build")
|
|
return
|
|
print(f"PyTorch MPS built: {mps_mod.is_built()}")
|
|
print(f"PyTorch MPS available: {mps_mod.is_available()}")
|
|
|
|
|
|
def main():
|
|
print("=== ComfyUI-Recgen Local Smoke Demo ===")
|
|
_print_mps_status()
|
|
|
|
recgen_dir = WORKSPACE_DIR / "recgen" / "examples"
|
|
if not recgen_dir.is_dir():
|
|
raise SystemExit(f"Missing example data: {recgen_dir}")
|
|
|
|
image_tensor, depth_tensor, mask_tensor, (fx, fy, cx, cy) = load_example_smoke_tensors(
|
|
recgen_dir, size=SMOKE_SIZE
|
|
)
|
|
|
|
print(f"Smoke RGB tensor: {tuple(image_tensor.shape)}")
|
|
print(f"Smoke depth tensor: {tuple(depth_tensor.shape)}")
|
|
print(f"Smoke mask tensor: {tuple(mask_tensor.shape)}")
|
|
print(f"Intrinsics (256²): fx={fx:.1f}, fy={fy:.1f}, cx={cx:.1f}, cy={cy:.1f}")
|
|
|
|
device = _pick_device()
|
|
print(f"Device: {device}")
|
|
|
|
node = RecGen2DTo3DNode()
|
|
print("Running 3D generation...")
|
|
out_dir, overlay_tensor = node.generate_3d(
|
|
image=image_tensor,
|
|
depth=depth_tensor,
|
|
mask=mask_tensor,
|
|
checkpoint="recgen_base.multiview_stereo",
|
|
device=device,
|
|
fx=fx,
|
|
fy=fy,
|
|
cx=cx,
|
|
cy=cy,
|
|
seed=42,
|
|
save_splat=False,
|
|
save_glb=False,
|
|
)
|
|
|
|
out_path = Path(out_dir)
|
|
mesh_path = out_path / "posed_mesh.obj"
|
|
overlay_path = out_path / "overlay.png"
|
|
|
|
print("\n=== Success ===")
|
|
print(f"Output directory: {out_path}")
|
|
print(f" posed_mesh.obj: {mesh_path.exists()}")
|
|
print(f" overlay.png: {overlay_path.exists()}")
|
|
print(f" overlay tensor shape: {tuple(overlay_tensor.shape)}")
|
|
|
|
if not mesh_path.exists():
|
|
raise SystemExit("Demo failed: posed_mesh.obj was not created")
|
|
if not overlay_path.exists():
|
|
raise SystemExit("Demo failed: overlay.png was not created")
|
|
|
|
print("\nLocal smoke demo completed successfully.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|