Andrea Pozzetti
e114ecefdf
Split Cadrille and CAD-Recode nodes from the CADabra monorepo. Multi-modal CAD reconstruction from point clouds, images, or text. Original source: https://github.com/PozzettiAndrea/ComfyUI-CADabra
288 lines
7.9 KiB
Python
288 lines
7.9 KiB
Python
"""
|
|
CAD-Recode Inference Pipeline
|
|
Handles point cloud preprocessing and code generation.
|
|
"""
|
|
|
|
import numpy as np
|
|
import torch
|
|
from typing import Optional, Tuple, Union
|
|
from transformers import AutoTokenizer
|
|
|
|
|
|
def normalize_point_cloud(
|
|
points: np.ndarray,
|
|
scale_range: float = 2.0,
|
|
center: bool = True
|
|
) -> np.ndarray:
|
|
"""
|
|
Normalize point cloud to fit within a centered cube.
|
|
|
|
Args:
|
|
points: (N, 3) point coordinates
|
|
scale_range: Size of the bounding cube (default 2.0 for [-1, 1])
|
|
center: Whether to center the point cloud at origin
|
|
|
|
Returns:
|
|
(N, 3) normalized point coordinates
|
|
"""
|
|
points = points.copy()
|
|
|
|
if center:
|
|
# Center at origin using bounding box center
|
|
bounds_min = points.min(axis=0)
|
|
bounds_max = points.max(axis=0)
|
|
center_point = (bounds_min + bounds_max) / 2.0
|
|
points = points - center_point
|
|
|
|
# Scale to fit in cube of size scale_range
|
|
extents = points.max(axis=0) - points.min(axis=0)
|
|
max_extent = max(extents)
|
|
if max_extent > 0:
|
|
points = points * (scale_range / max_extent)
|
|
|
|
return points
|
|
|
|
|
|
def sample_points_fps(
|
|
points: np.ndarray,
|
|
n_points: int = 256,
|
|
n_pre_points: int = 8192
|
|
) -> np.ndarray:
|
|
"""
|
|
Sample points using Farthest Point Sampling (v1.5 improvement).
|
|
Pure numpy implementation - no pytorch3d dependency.
|
|
|
|
This produces more uniform coverage than random sampling.
|
|
|
|
Args:
|
|
points: (N, 3) point coordinates
|
|
n_points: Number of points to sample (default 256 for CAD-Recode)
|
|
n_pre_points: Maximum points to consider before FPS
|
|
|
|
Returns:
|
|
(n_points, 3) sampled point coordinates
|
|
"""
|
|
n = len(points)
|
|
|
|
if n <= n_points:
|
|
# Pad if not enough points
|
|
if n < n_points:
|
|
indices = np.random.choice(n, n_points, replace=True)
|
|
return points[indices]
|
|
return points
|
|
|
|
# Limit input size for efficiency
|
|
if n > n_pre_points:
|
|
indices = np.random.choice(n, n_pre_points, replace=False)
|
|
points = points[indices]
|
|
n = n_pre_points
|
|
|
|
# Pure numpy FPS implementation
|
|
selected = np.zeros(n_points, dtype=np.int64)
|
|
distances = np.full(n, np.inf)
|
|
|
|
# Start with random point
|
|
selected[0] = np.random.randint(n)
|
|
|
|
for i in range(1, n_points):
|
|
# Distance from last selected point to all points
|
|
last_point = points[selected[i - 1]]
|
|
dist_to_last = np.linalg.norm(points - last_point, axis=1)
|
|
|
|
# Update minimum distances
|
|
distances = np.minimum(distances, dist_to_last)
|
|
|
|
# Select point with maximum distance to any selected point
|
|
selected[i] = np.argmax(distances)
|
|
|
|
return points[selected]
|
|
|
|
|
|
def sample_points_random(
|
|
points: np.ndarray,
|
|
n_points: int = 256
|
|
) -> np.ndarray:
|
|
"""
|
|
Sample points using random sampling.
|
|
|
|
Args:
|
|
points: (N, 3) point coordinates
|
|
n_points: Number of points to sample
|
|
|
|
Returns:
|
|
(n_points, 3) sampled point coordinates
|
|
"""
|
|
if len(points) <= n_points:
|
|
if len(points) < n_points:
|
|
indices = np.random.choice(len(points), n_points, replace=True)
|
|
return points[indices]
|
|
return points
|
|
|
|
indices = np.random.choice(len(points), n_points, replace=False)
|
|
return points[indices]
|
|
|
|
|
|
def mesh_to_point_cloud(
|
|
vertices: np.ndarray,
|
|
faces: np.ndarray,
|
|
n_points: int = 256,
|
|
use_fps: bool = True
|
|
) -> np.ndarray:
|
|
"""
|
|
Convert mesh to point cloud by sampling surface points.
|
|
|
|
Args:
|
|
vertices: (V, 3) mesh vertices
|
|
faces: (F, 3) mesh faces
|
|
n_points: Number of points to sample
|
|
use_fps: Use farthest point sampling (slower but better coverage)
|
|
|
|
Returns:
|
|
(n_points, 3) sampled point coordinates
|
|
"""
|
|
try:
|
|
import trimesh
|
|
mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
|
|
# Sample more points than needed, then use FPS/random to select
|
|
surface_points, _ = trimesh.sample.sample_surface(mesh, n_points * 32)
|
|
|
|
if use_fps:
|
|
return sample_points_fps(surface_points, n_points)
|
|
else:
|
|
return sample_points_random(surface_points, n_points)
|
|
except ImportError:
|
|
# Fallback: just sample from vertices
|
|
return sample_points_fps(vertices, n_points) if use_fps else sample_points_random(vertices, n_points)
|
|
|
|
|
|
def prepare_input(
|
|
point_cloud: np.ndarray,
|
|
tokenizer,
|
|
n_points: int = 256,
|
|
normalize: bool = True,
|
|
use_fps: bool = True,
|
|
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Prepare point cloud and tokenizer inputs for CAD-Recode inference.
|
|
|
|
Args:
|
|
point_cloud: (N, 3) point coordinates
|
|
tokenizer: Qwen2 tokenizer
|
|
n_points: Number of points (default 256)
|
|
normalize: Whether to normalize point cloud
|
|
use_fps: Use farthest point sampling
|
|
|
|
Returns:
|
|
Tuple of (input_ids, attention_mask, point_cloud_tensor)
|
|
"""
|
|
# Preprocess point cloud
|
|
if normalize:
|
|
point_cloud = normalize_point_cloud(point_cloud)
|
|
|
|
if use_fps:
|
|
point_cloud = sample_points_fps(point_cloud, n_points)
|
|
else:
|
|
point_cloud = sample_points_random(point_cloud, n_points)
|
|
|
|
# Create input sequence: [point_tokens...] [<|im_start|>]
|
|
# Point tokens use pad_token_id as placeholder
|
|
input_ids = [tokenizer.pad_token_id] * len(point_cloud) + \
|
|
[tokenizer('<|im_start|>')['input_ids'][0]]
|
|
|
|
# Attention mask: -1 for point positions (special marker), 1 for text
|
|
attention_mask = [-1] * len(point_cloud) + [1]
|
|
|
|
return (
|
|
torch.tensor(input_ids).unsqueeze(0),
|
|
torch.tensor(attention_mask).unsqueeze(0),
|
|
torch.tensor(point_cloud.astype(np.float32)).unsqueeze(0)
|
|
)
|
|
|
|
|
|
def generate_code(
|
|
model,
|
|
tokenizer,
|
|
point_cloud: np.ndarray,
|
|
n_points: int = 256,
|
|
max_tokens: int = 768,
|
|
normalize: bool = True,
|
|
use_fps: bool = True,
|
|
device: Optional[str] = None,
|
|
) -> str:
|
|
"""
|
|
Generate CadQuery code from point cloud using CAD-Recode model.
|
|
|
|
Args:
|
|
model: CADRecode model instance
|
|
tokenizer: Qwen2 tokenizer
|
|
point_cloud: (N, 3) point coordinates
|
|
n_points: Number of points to sample (default 256)
|
|
max_tokens: Maximum tokens to generate (default 768)
|
|
normalize: Whether to normalize point cloud
|
|
use_fps: Use farthest point sampling
|
|
device: Device to run on (None = auto-detect from model)
|
|
|
|
Returns:
|
|
Generated CadQuery Python code string
|
|
"""
|
|
if device is None:
|
|
device = next(model.parameters()).device
|
|
|
|
# Prepare inputs
|
|
input_ids, attention_mask, pc_tensor = prepare_input(
|
|
point_cloud, tokenizer, n_points, normalize, use_fps
|
|
)
|
|
|
|
# Move to device
|
|
input_ids = input_ids.to(device)
|
|
attention_mask = attention_mask.to(device)
|
|
pc_tensor = pc_tensor.to(device)
|
|
|
|
# Generate
|
|
with torch.no_grad():
|
|
output = model.generate(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
point_cloud=pc_tensor,
|
|
max_new_tokens=max_tokens,
|
|
pad_token_id=tokenizer.pad_token_id,
|
|
)
|
|
|
|
# Decode and extract code
|
|
text = tokenizer.batch_decode(output)[0]
|
|
|
|
# Extract code between markers
|
|
start_marker = '<|im_start|>'
|
|
end_marker = '<|endoftext|>'
|
|
|
|
start = text.find(start_marker)
|
|
if start != -1:
|
|
start += len(start_marker)
|
|
else:
|
|
start = 0
|
|
|
|
end = text.find(end_marker)
|
|
if end == -1:
|
|
end = len(text)
|
|
|
|
code = text[start:end].strip()
|
|
return code
|
|
|
|
|
|
def load_tokenizer(model_id: str = 'Qwen/Qwen2-1.5B'):
|
|
"""
|
|
Load the Qwen2 tokenizer configured for CAD-Recode.
|
|
|
|
Args:
|
|
model_id: HuggingFace model ID for tokenizer
|
|
|
|
Returns:
|
|
Configured tokenizer instance
|
|
"""
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
model_id,
|
|
pad_token='<|im_end|>',
|
|
padding_side='left'
|
|
)
|
|
return tokenizer
|