import spaces # must be imported before torch / anything touching CUDA import os import sys import time import tempfile from functools import lru_cache import numpy as np import torch import gradio as gr import trimesh import imageio.v2 as imageio from PIL import Image, ImageDraw, ImageFont from huggingface_hub import hf_hub_download ROOT = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, ROOT) sys.path.insert(0, os.path.join(ROOT, "src")) from src.train import CUT3RDynamicsModel, _adapt_state_dict_for_height_input # noqa: E402 from src.data.multi_dataloader import MultiObjectDynamics # noqa: E402 MODEL_REPO = "frankzydou/RigidFormer" DATA_REPO = "frankzydou/RigidFormer-MOVi" DATASETS = { # UI name: (parquet config, checkpoint) "MOVi-B": ("movi_b", "rigidformer_movib.pth"), "MOVi-A": ("movi_a", "rigidformer_movia.pth"), "MOVi-Sphere": ("movi_sphere", "rigidformer_movisphere.pth"), } SHAPES = ("cone", "cube", "cylinder", "gear", "sphere", "sponge", "spot", "suzanne", "teapot", "torus", "torus_knot") FPS_PHYS = 240 # physical frames per second of the MOVi simulations WARMUP = 4 # ground-truth context frames (evaluation protocol of the paper) ANCHOR_K = 4 VIDEO_FPS = 48 # one rendered frame every 5 physical frames (real time) GRAVITY = np.array([0.0, 0.0, -10.0]) # -------------------------------------------------------------------------------------- # Data: test scenes of RigidFormer-MOVi (CC BY 4.0) and KuBasic collision meshes (CC BY-SA 4.0) # -------------------------------------------------------------------------------------- COLLISION = {} for s in SHAPES: p = hf_hub_download(DATA_REPO, f"assets/kubasic/{s}/collision_geometry.obj", repo_type="dataset") # file order of the vertices == the order of the object's vertices in the training data m = trimesh.load(p, process=False) COLLISION[s] = (np.asarray(m.vertices, np.float64), np.asarray(m.faces, np.int64)) def _render_mesh(shape): V, F = COLLISION[shape] if shape == "sphere": # nicer sphere for display only (same radius as the collision sphere) ico = trimesh.creation.icosphere(subdivisions=3, radius=0.5) V, F = np.asarray(ico.vertices), np.asarray(ico.faces) return V, F RENDER_MESH = {s: _render_mesh(s) for s in SHAPES} TRAJ_COLS = ("positions", "quaternions", "velocities") def _read_rows(path): import pyarrow.parquet as pq cols = ["scene", "num_objects", "shape", "size", "material", "mass", "friction", "restitution", "color", "positions", "quaternions", "velocities"] table = pq.read_table(path, columns=cols) frame_cols = [(k, 3 if k != "quaternions" else 4) for k in TRAJ_COLS] plain = table.drop_columns([k for k, _ in frame_cols]).to_pylist() for key, d in frame_cols: col = table.column(key).combine_chunks() frame_off = col.offsets.to_numpy() frames = col.flatten() obj_off = frames.offsets.to_numpy() values = frames.flatten().flatten().to_numpy(zero_copy_only=False).reshape(-1, d) for r, rec in enumerate(plain): f0, f1 = frame_off[r], frame_off[r + 1] rec[key] = values[obj_off[f0]:obj_off[f1]].reshape(f1 - f0, rec["num_objects"], d).astype(np.float64) return plain SCENES = {} for name, (cfg, _) in DATASETS.items(): p = hf_hub_download(DATA_REPO, f"data/{cfg}/test-00000-of-00001.parquet", repo_type="dataset") SCENES[name] = _read_rows(p) print(f"[data] {name}: {len(SCENES[name])} test scenes", flush=True) # -------------------------------------------------------------------------------------- # Models (one per dataset), loaded at module scope and moved to CUDA eagerly (ZeroGPU) # -------------------------------------------------------------------------------------- MODELS = {} for name, (_, ckpt_name) in DATASETS.items(): ckpt = torch.load(hf_hub_download(MODEL_REPO, ckpt_name), map_location="cpu", weights_only=False) config = ckpt["config"] model = CUT3RDynamicsModel(config) sd = {k[len("module."):] if k.startswith("module.") else k: v for k, v in ckpt["model_state_dict"].items()} sd = _adapt_state_dict_for_height_input(model, sd) model.load_state_dict(sd, strict=True) model.eval() MODELS[name] = model.to("cuda") del ckpt, sd print(f"[model] {name}: {ckpt_name} loaded", flush=True) # -------------------------------------------------------------------------------------- # Geometry helpers # -------------------------------------------------------------------------------------- def quat_wxyz_to_rotmat(q): q = np.asarray(q, np.float64) q = q / np.linalg.norm(q, axis=-1, keepdims=True) w, x, y, z = q[..., 0], q[..., 1], q[..., 2], q[..., 3] R = np.stack([ 1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w), 2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w), 2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)], axis=-1) return R.reshape(q.shape[:-1] + (3, 3)) def kabsch_poses(reference, frames): """Rigid poses (R (T,3,3), t (T,3)) taking reference (V,3) to frames (T,V,3).""" X0, X = np.asarray(reference, np.float64), np.asarray(frames, np.float64) c0, c = X0.mean(0), X.mean(1) U, _, Vt = np.linalg.svd(np.einsum("vi,tvj->tij", X0 - c0, X - c[:, None])) V, Ut = np.transpose(Vt, (0, 2, 1)), np.transpose(U, (0, 2, 1)) D = np.tile(np.eye(3), (len(X), 1, 1)) D[:, 2, 2] = np.sign(np.linalg.det(V @ Ut)) R = V @ D @ Ut return R, c - np.einsum("tij,j->ti", R, c0) @lru_cache(maxsize=None) def object_pool(dataset): """Every object of the dataset's test scenes as it appears at frame 0 (for random new scenes).""" pool = [] for rec in SCENES[dataset]: for i in range(rec["num_objects"]): pool.append({k: rec[k][i] for k in ("shape", "size", "material", "mass", "friction", "restitution", "color")} | {"quat": rec["quaternions"][0, i]}) return pool def random_scene(dataset, num_objects, seed, speed=1.0, horizon=480): """A new MOVi-style scene: objects drawn from the dataset's object pool, dropped from 1.5-4.5 m with a horizontal throw toward the centre (as in the MOVi generator). Only the first context frames are known (free flight before any contact), so the result has no ground truth.""" rng = np.random.default_rng(int(seed)) pool = object_pool(dataset) objs, placed = [], [] tries = 0 while len(objs) < num_objects and tries < 5000: tries += 1 o = pool[rng.integers(len(pool))] radius = float(o["size"]) * np.linalg.norm(COLLISION[o["shape"]][0], axis=1).max() pos = np.array([rng.uniform(-4.0, 4.0), rng.uniform(-4.0, 4.0), rng.uniform(1.5, 4.5)]) if any(np.linalg.norm(pos - q) < radius + r + 0.05 for q, r in placed): continue placed.append((pos, radius)) target = rng.uniform(-4.0, 4.0, size=2) vel = np.array([*(speed * (target - pos[:2])), 0.0]) objs.append(dict(o, pos=pos, vel=vel)) n = len(objs) T = horizon t = np.arange(T)[:, None] / FPS_PHYS positions = np.stack([o["pos"][None] + o["vel"][None] * t + 0.5 * GRAVITY[None] * t ** 2 for o in objs], 1) quats = np.stack([np.tile(o["quat"], (T, 1)) for o in objs], 1) return { "scene": f"random (seed {int(seed)})", "num_objects": n, "shape": [o["shape"] for o in objs], "size": [float(o["size"]) for o in objs], "material": [o["material"] for o in objs], "mass": [float(o["mass"]) for o in objs], "friction": [float(o["friction"]) for o in objs], "restitution": [float(o["restitution"]) for o in objs], "color": [o["color"] for o in objs], "positions": positions, "quaternions": quats, "ground_truth": False, } @lru_cache(maxsize=64) def _fps_cached(key, V_bytes, n): V = np.frombuffer(V_bytes, np.float32).reshape(n, 3) anchors = MultiObjectDynamics._farthest_point_sample_np(V, ANCHOR_K, rule="centroid_farthest") order = MultiObjectDynamics._farthest_point_sample_np(V, n, rule="centroid_farthest") return anchors, order def build_batch(rec, step_size): """The model input exactly as the authors' data loader builds it for evaluation (step size s: context frames 0, s, 2s, 3s; reference pose = global frame 0; FPS anchors and vertex ranks from the stored, unscaled frame-0 vertices).""" n = rec["num_objects"] nvs = [len(COLLISION[s][0]) for s in rec["shape"]] max_nv = max(nvs) frames = [k * step_size for k in range(WARMUP)] R_all = quat_wxyz_to_rotmat(rec["quaternions"][: frames[-1] + 1]) # (T, n, 3, 3) P_all = rec["positions"][: frames[-1] + 1] def verts(t, i, scaled=True): v = COLLISION[rec["shape"][i]][0] s = rec["size"][i] if scaled else 1.0 return P_all[t, i] + s * v @ R_all[t, i].T obj = np.zeros((1, WARMUP, n, max_nv, 3), np.float32) vmask = np.zeros((1, WARMUP, n, max_nv), bool) ref = np.zeros((1, n, max_nv, 3), np.float32) anchor_idx = np.zeros((1, n, ANCHOR_K), np.int64) vtx_rank = np.full((1, n, max_nv), max_nv, np.int64) for i in range(n): nv = nvs[i] for k, t in enumerate(frames): obj[0, k, i, :nv] = verts(t, i) vmask[0, k, i, :nv] = True ref[0, i, :nv] = verts(0, i) V0 = verts(0, i, scaled=False).astype(np.float32) anchors, order = _fps_cached(rec["shape"][i], V0.tobytes(), nv) anchor_idx[0, i] = anchors vtx_rank[0, i, order] = np.arange(nv) phys = np.array([[rec["mass"][i], rec["friction"][i], rec["restitution"][i]] for i in range(n)], np.float32)[None] return { "obj_positions": obj, "valid_obj_mask": np.ones((1, n), bool), "valid_vtx_mask": vmask, "ref_obj_positions": ref, "anchor_idx": anchor_idx, "vtx_rank": vtx_rank, "phys_params": phys, "skip_step_size": np.array([float(step_size)], np.float32), } @spaces.GPU(duration=6) def run_rollout(dataset, batch, num_steps): model = MODELS[dataset] tb = {k: torch.from_numpy(v).to("cuda") for k, v in batch.items()} torch.cuda.synchronize() t0 = time.time() try: # the authors' --cuda-graph fast path: one step is captured and replayed (bit-identical trajectories) os.environ["RIGIDFORMER_CUDA_GRAPH"] = "1" out = model.rollout(initial_batch=tb, num_steps=num_steps, dt=1.0, return_extras=False, warmup_frames=WARMUP) except Exception as e: # fall back to the default eager path print(f"[rollout] CUDA-graph path failed ({type(e).__name__}: {e}); using the eager path", flush=True) os.environ["RIGIDFORMER_CUDA_GRAPH"] = "0" tb = {k: torch.from_numpy(v).to("cuda") for k, v in batch.items()} out = model.rollout(initial_batch=tb, num_steps=num_steps, dt=1.0, return_extras=False, warmup_frames=WARMUP) torch.cuda.synchronize() elapsed = time.time() - t0 print(f"[rollout] {dataset}: {num_steps} frames, step size {float(batch['skip_step_size'][0]):g}, " f"{batch['obj_positions'].shape[2]} objects: {elapsed:.2f} s", flush=True) return out["obj_positions"][0].float().cpu().numpy(), elapsed # -------------------------------------------------------------------------------------- # Rendering (CPU, painter's algorithm with PIL) # -------------------------------------------------------------------------------------- W, H, SS = 640, 480, 2 LIGHT = np.array([0.35, -0.45, 0.82]); LIGHT /= np.linalg.norm(LIGHT) def _camera(): target = np.array([0.0, 0.0, 0.6]) eye = np.array([0.0, -11.5, 7.8]) f = target - eye; f /= np.linalg.norm(f) r = np.cross(f, [0, 0, 1.0]); r /= np.linalg.norm(r) u = np.cross(r, f) focal = (H * SS / 2) / np.tan(np.radians(42) / 2) return eye, f, r, u, focal CAM = _camera() def _project(X): eye, f, r, u, focal = CAM d = X - eye z = d @ f x = d @ r y = d @ u z = np.maximum(z, 1e-3) return np.stack([W * SS / 2 + focal * x / z, H * SS / 2 - focal * y / z], -1), z @lru_cache(maxsize=1) def _floor_image(): img = Image.new("RGB", (W * SS, H * SS), (236, 238, 242)) draw = ImageDraw.Draw(img) # horizon gradient for yy in range(0, H * SS, 4): c = int(246 - 18 * yy / (H * SS)) draw.rectangle([0, yy, W * SS, yy + 4], fill=(c, c, c + 4)) quads = [] for i in range(-10, 10): for j in range(-10, 10): Q = np.array([[i, j, 0], [i + 1, j, 0], [i + 1, j + 1, 0], [i, j + 1, 0]], float) xy, z = _project(Q) if (z < 0.5).any(): continue shade = 222 if (i + j) % 2 == 0 else 200 quads.append((z.mean(), xy, shade)) for _, xy, shade in sorted(quads, key=lambda q: -q[0]): draw.polygon([tuple(p) for p in xy], fill=(shade, shade, shade + 6)) return img def render_frame(objects, title, subtitle=""): """objects: list of (vertices (V,3) world, faces (F,3), rgb 0-1).""" img = _floor_image().copy() draw = ImageDraw.Draw(img, "RGBA") eye = CAM[0] # blob shadows for V, _, _ in objects: c = V.mean(0) rad = 0.5 * (V.max(0) - V.min(0))[:2].max() h = max(V[:, 2].min(), 0.0) a = int(110 * np.exp(-h / 2.0)) ring = np.array([[c[0] + rad * np.cos(t), c[1] + rad * np.sin(t), 0.001] for t in np.linspace(0, 2 * np.pi, 24)]) xy, z = _project(ring) if (z > 0.5).all(): draw.polygon([tuple(p) for p in xy], fill=(40, 40, 50, a)) tris = [] for V, F, rgb in objects: A, B, C = V[F[:, 0]], V[F[:, 1]], V[F[:, 2]] N = np.cross(B - A, C - A) N /= np.linalg.norm(N, axis=1, keepdims=True) + 1e-12 centers = (A + B + C) / 3 vis = ((eye - centers) * N).sum(1) > 0 if not vis.any(): continue xy, z = _project(V) lam = np.clip(N[vis] @ LIGHT, 0, 1) view = eye - centers[vis]; view /= np.linalg.norm(view, axis=1, keepdims=True) hvec = view + LIGHT; hvec /= np.linalg.norm(hvec, axis=1, keepdims=True) spec = np.clip((N[vis] * hvec).sum(1), 0, 1) ** 24 base = np.asarray(rgb, float) col = np.clip(base[None] * (0.32 + 0.68 * lam[:, None]) + 0.35 * spec[:, None], 0, 1) col = (col * 255).astype(int) depth = z[F[vis]].mean(1) for d, tri, c in zip(depth, F[vis], col): tris.append((d, xy[tri], tuple(c))) tris.sort(key=lambda q: -q[0]) for _, xy, c in tris: draw.polygon([tuple(p) for p in xy], fill=c) img = img.resize((W, H), Image.LANCZOS) d2 = ImageDraw.Draw(img) try: font = ImageFont.load_default(size=22) small = ImageFont.load_default(size=15) except TypeError: font = small = ImageFont.load_default() d2.text((14, 10), title, fill=(30, 30, 40), font=font) if subtitle: d2.text((14, 38), subtitle, fill=(70, 70, 85), font=small) return np.asarray(img) def _color(c): c = np.clip(np.asarray(c, float), 0, 1) return 0.15 + 0.8 * c # keep colours vivid but never pure black/white def poses_to_objects(rec, R, t): """R (n,3,3), t (n,3) -> renderable objects.""" objs = [] for i in range(rec["num_objects"]): V, F = RENDER_MESH[rec["shape"][i]] Vw = t[i] + rec["size"][i] * V @ R[i].T objs.append((Vw, F, _color(rec["color"][i]))) return objs def interp_pose(R_k, t_k, times_k, t): """Pose at time t from key poses (linear translation, slerp rotation).""" from scipy.spatial.transform import Rotation, Slerp j = np.searchsorted(times_k, t, side="right") - 1 j = int(np.clip(j, 0, len(times_k) - 2)) if len(times_k) > 1 else 0 if len(times_k) == 1 or t <= times_k[0]: return R_k[0], t_k[0] if t >= times_k[-1]: return R_k[-1], t_k[-1] a = (t - times_k[j]) / (times_k[j + 1] - times_k[j]) rot = Slerp([0, 1], Rotation.from_matrix(np.stack([R_k[j], R_k[j + 1]])))([a]).as_matrix()[0] return rot, (1 - a) * t_k[j] + a * t_k[j + 1] # -------------------------------------------------------------------------------------- # Main entry point # -------------------------------------------------------------------------------------- def simulate(dataset: str = "MOVi-B", scene_source: str = "MOVi test scene", scene_index: int = 0, step_size: int = 1, num_frames: int = 480, num_objects: int = 6, seed: int = 0, progress=gr.Progress()): """Roll out RigidFormer on a multi-object rigid-body scene and render the prediction. Args: dataset: Which pretrained model / object set to use: "MOVi-B", "MOVi-A" or "MOVi-Sphere". scene_source: "MOVi test scene" (a held-out scene with simulator ground truth) or "Random new scene". scene_index: Index (0-119) of the held-out test scene, used when scene_source is "MOVi test scene". step_size: Integration step size in physical frames (1, 5 or 10; the simulation runs at 240 frames/s). num_frames: Rollout horizon in physical frames (480 = 2 seconds). num_objects: Number of objects for a random new scene (3-10). seed: Random seed for a random new scene. Returns: Rendered rollout video, per-frame error plot, a markdown summary and an .npz file with the predicted poses. """ step_size = int(step_size) if step_size not in (1, 5, 10): raise gr.Error("Step size must be 1, 5 or 10 (the step sizes RigidFormer was trained with).") num_frames = int(num_frames) num_frames = max(num_frames - num_frames % step_size, (WARMUP + 1) * step_size) if scene_source == "Random new scene": rec = random_scene(dataset, int(num_objects), int(seed), horizon=num_frames) has_gt = False else: recs = SCENES[dataset] rec = dict(recs[int(scene_index) % len(recs)]) rec["ground_truth"] = True has_gt = True n = rec["num_objects"] progress(0.05, desc="Building the point-cloud input") batch = build_batch(rec, step_size) num_steps = num_frames // step_size # network frames 0, s, ..., num_frames - s (as the evaluator) progress(0.15, desc=f"RigidFormer rollout ({num_steps - WARMUP} steps)") pred, elapsed = run_rollout(dataset, batch, num_steps) # (K, n, Nv, 3) K = pred.shape[0] times_k = np.arange(K) * step_size # Rigid poses of every object at every predicted frame (Kabsch fit to the predicted vertices) R_pred = np.zeros((K, n, 3, 3)); t_pred = np.zeros((K, n, 3)) for i in range(n): local = rec["size"][i] * COLLISION[rec["shape"][i]][0] nv = len(local) R_pred[:, i], t_pred[:, i] = kabsch_poses(local, pred[:, i, :nv]) # Errors vs. the simulator metrics_md = "" plot = None if has_gt: P_gt = rec["positions"][times_k] # (K, n, 3) R_gt = quat_wxyz_to_rotmat(rec["quaternions"][times_k]) pos_rmse = np.sqrt(((t_pred - P_gt) ** 2).sum(-1).mean(-1)) rel = np.einsum("tnij,tnlj->tnil", R_pred, R_gt) # R_pred @ R_gt^T cosang = (np.trace(rel, axis1=-2, axis2=-1) - 1) / 2 rot_err = np.degrees(np.sqrt((np.arccos(np.clip(cosang, -1, 1)) ** 2).mean(-1))) rows = [] for cp in (25, 50, 100, 200, 300, 400): if cp <= times_k[-1] and cp % step_size == 0 and cp // step_size >= WARMUP: k = cp // step_size rows.append(f"| @{cp} | {pos_rmse[k]:.4f} | {rot_err[k]:.2f} |") metrics_md = ("\n\n| physical frame | position RMSE (m) | rotation error (deg) |\n|---|---|---|\n" + "\n".join(rows) if rows else "") import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt fig, ax1 = plt.subplots(figsize=(7, 3.2)) ax1.plot(times_k / FPS_PHYS, pos_rmse, color="#d97706", lw=2, label="position RMSE (m)") ax1.set_xlabel("time (s)"); ax1.set_ylabel("position RMSE (m)", color="#d97706") ax2 = ax1.twinx() ax2.plot(times_k / FPS_PHYS, rot_err, color="#475569", lw=2, ls="--", label="rotation error (deg)") ax2.set_ylabel("rotation error (deg)", color="#475569") ax1.set_title("RigidFormer vs. simulator ground truth") ax1.grid(alpha=0.3) fig.tight_layout() plot = fig # Video: ground truth (if any) | RigidFormer, one frame every 5 physical frames (real time at 48 fps) progress(0.4, desc="Rendering") out_path = os.path.join(tempfile.mkdtemp(), "rigidformer_rollout.mp4") writer = imageio.get_writer(out_path, fps=VIDEO_FPS, codec="libx264", quality=7, macro_block_size=16, ffmpeg_log_level="error") show_t = np.arange(0, times_k[-1] + 1, FPS_PHYS // VIDEO_FPS) R_q_all = quat_wxyz_to_rotmat(rec["quaternions"][: times_k[-1] + 1]) if has_gt else None for fi, t in enumerate(show_t): if fi % 8 == 0: progress(0.4 + 0.55 * fi / len(show_t), desc="Rendering") Rt = [interp_pose(R_pred[:, i], t_pred[:, i], times_k, t) for i in range(n)] R_now = np.stack([a for a, _ in Rt]); t_now = np.stack([b for _, b in Rt]) ctx = "context (ground truth)" if t <= times_k[WARMUP - 1] else "prediction" sub = f"t = {t / FPS_PHYS:.2f} s step size {step_size} [{ctx}]" right = render_frame(poses_to_objects(rec, R_now, t_now), "RigidFormer", sub) if has_gt: left = render_frame(poses_to_objects(rec, R_q_all[t], rec["positions"][t]), "Simulator (ground truth)", f"t = {t / FPS_PHYS:.2f} s") frame = np.concatenate([left, right], 1) else: frame = right writer.append_data(frame) writer.close() npz_path = os.path.join(tempfile.mkdtemp(), "rigidformer_poses.npz") np.savez(npz_path, rotation=R_pred.astype(np.float32), translation=t_pred.astype(np.float32), frame=times_k.astype(np.int32), shape=np.array(rec["shape"]), size=np.array(rec["size"], np.float32), material=np.array(rec["material"])) objs_desc = ", ".join(f"{m} {s}" for s, m in zip(rec["shape"], rec["material"])) summary = (f"**{dataset} — {rec['scene']}** · {n} objects ({objs_desc})\n\n" f"Rolled out {K - WARMUP} steps of size {step_size} ({times_k[-1] / FPS_PHYS:.2f} s of physics) " f"from {WARMUP} context frames in **{elapsed:.2f} s** on the GPU " f"({1000 * elapsed / max(K - WARMUP, 1):.1f} ms/step).") if not has_gt: summary += ("\n\nThis is a new scene with no simulator ground truth: only the first frames (free flight) are " "given, everything after that is predicted by RigidFormer.") summary += metrics_md return out_path, plot, summary, npz_path # -------------------------------------------------------------------------------------- # UI # -------------------------------------------------------------------------------------- DESCRIPTION = """ # RigidFormer: Learning Rigid Dynamics using Transformers **RigidFormer** ([paper](https://huggingface.co/papers/2605.09196) · [code](https://github.com/Frank-ZY-Dou/Dynamics-Modeling/tree/main/RigidFormer) · [weights](https://huggingface.co/frankzydou/RigidFormer)) is an object-centric Transformer that simulates multi-object rigid-body dynamics from **point sets** — no mesh connectivity. Given four context frames of each object's points, it predicts per-anchor accelerations, integrates them, and projects the update onto rigid motions (differentiable Kabsch), autoregressively for hundreds of steps at step size 1, 5 or 10. Pick a held-out **MOVi test scene** to compare against the physics simulator, or drop a **random new scene** of objects and watch the learned simulator take over. """ with gr.Blocks(title="RigidFormer") as demo: gr.Markdown(DESCRIPTION) with gr.Row(): with gr.Column(scale=1): dataset = gr.Dropdown(list(DATASETS), value="MOVi-B", label="Model / object set", info="MOVi-B: 11 KuBasic shapes · MOVi-A: cubes, cylinders, spheres · MOVi-Sphere: spheres") scene_source = gr.Radio(["MOVi test scene", "Random new scene"], value="MOVi test scene", label="Scene") scene_index = gr.Slider(0, 119, value=0, step=1, label="Test scene #", info="held-out scenes with simulator ground truth") with gr.Row(): num_objects = gr.Slider(3, 10, value=6, step=1, label="Objects (random scene)") seed = gr.Number(value=0, precision=0, label="Seed (random scene)") step_size = gr.Radio([1, 5, 10], value=1, label="Integration step size (physical frames)") with gr.Accordion("Advanced", open=False): num_frames = gr.Slider(60, 480, value=480, step=10, label="Horizon (physical frames @ 240 fps)") run = gr.Button("Simulate", variant="primary") with gr.Column(scale=2): video = gr.Video(label="Rollout", autoplay=True, loop=True) summary = gr.Markdown() plot = gr.Plot(label="Error vs. time") poses = gr.File(label="Predicted object poses (.npz)") inputs = [dataset, scene_source, scene_index, step_size, num_frames, num_objects, seed] outputs = [video, plot, summary, poses] run.click(simulate, inputs=inputs, outputs=outputs) gr.Examples( examples=[ ["MOVi-B", "MOVi test scene", 0, 1], ["MOVi-B", "MOVi test scene", 7, 5], ["MOVi-A", "MOVi test scene", 3, 1], ["MOVi-Sphere", "MOVi test scene", 11, 10], ["MOVi-B", "Random new scene", 0, 1], ], inputs=[dataset, scene_source, scene_index, step_size], outputs=outputs, fn=simulate, cache_examples=True, cache_mode="lazy", ) gr.Markdown("Scenes: [RigidFormer-MOVi](https://huggingface.co/datasets/frankzydou/RigidFormer-MOVi) test split " "(CC BY 4.0, derived from HopNet's MOVi datasets); meshes: Kubric KuBasic assets (CC BY-SA 4.0). " "Both are downloaded at startup.") if __name__ == "__main__": demo.launch(theme=gr.themes.Citrus(), mcp_server=True)