from __future__ import annotations import os import re import subprocess import sys import time from pathlib import Path import gradio as gr import spaces from huggingface_hub import hf_hub_download ROOT = Path(__file__).resolve().parent SRC_ROOT = ROOT / "src" if str(SRC_ROOT) not in sys.path: sys.path.insert(0, str(SRC_ROOT)) CKPT_REPO_ID = os.environ.get("PHYSFORMER_CKPT_REPO_ID", "yslan/physformer") CKPT_FILENAME = os.environ.get("PHYSFORMER_CKPT_FILENAME", "checkpoint-best.pt") CKPT_PATH = ROOT / "checkpoints" / "checkpoint-best.pt" PREVIEW_VERSION = "v4" RENDER_COLORS = [ (0.86, 0.24, 0.20, 1.0), (0.20, 0.64, 0.42, 1.0), (0.20, 0.44, 0.86, 1.0), (0.92, 0.67, 0.22, 1.0), (0.62, 0.32, 0.76, 1.0), ] NAMED_COLORS = { "cow": (0.00, 0.62, 0.66, 1.0), "horse": (0.88, 0.30, 0.24, 1.0), } MESH_EDGE_COLOR = (0.05, 0.06, 0.07, 0.62) RIGID_RENDER_ALPHA = 0.96 ELASTIC_RENDER_ALPHA = 0.38 def _natural_path_key(path: Path) -> tuple[object, ...]: import re parts: list[object] = [] for part in path.parts: parts.extend(int(text) if text.isdigit() else text.lower() for text in re.split(r"(\d+)", part) if text) return tuple(parts) def _is_input_sample_dir(path: Path) -> bool: return ( path.name.isdigit() and (path / "metadata.json").is_file() and (path / "meshes" / "combined_frame_000.obj").is_file() and (path / "vertex_velocities" / "combined_frame_000.npy").is_file() ) def _ood_sample_paths(group: str, fallback_count: int = 10) -> list[str]: root = ROOT / "ood_examples" / group samples = [ path.relative_to(ROOT).as_posix() for path in sorted(root.iterdir(), key=_natural_path_key) if path.is_dir() and _is_input_sample_dir(path) ] if root.is_dir() else [] if samples: return samples return [f"ood_examples/{group}/{i}" for i in range(fallback_count)] OOD_SAMPLE_CHOICES = { "2 objects": _ood_sample_paths("2obj_cow_horse"), "3 objects": _ood_sample_paths("3obj_teapot_fish_bunny"), } SAMPLE_CHOICES = { "OOD mixed materials": OOD_SAMPLE_CHOICES["2 objects"] + OOD_SAMPLE_CHOICES["3 objects"], "In-distribution rigid": [ "indistri_examples/rigid/sample_000007", "indistri_examples/rigid/sample_000114", "indistri_examples/rigid/sample_000969", "indistri_examples/rigid/sample_001151", "indistri_examples/rigid/sample_001557", "indistri_examples/rigid/sample_001874", "indistri_examples/rigid/sample_002143", ], "In-distribution elastic": [ "indistri_examples/elastic/sample_000047", "indistri_examples/elastic/sample_000105", "indistri_examples/elastic/sample_000121", "indistri_examples/elastic/sample_000203", ], } DEFAULT_MATERIALS = { "cow": "rigid", "horse": "elastic", "teapot": "rigid", "fish": "elastic", "bunny": "elastic", } def example_ids_for_setting(setting: str, object_count: str = "2 objects") -> list[str]: if setting == "OOD mixed materials": samples = OOD_SAMPLE_CHOICES.get(str(object_count), OOD_SAMPLE_CHOICES["2 objects"]) return [str(i) for i in range(len(samples))] samples = SAMPLE_CHOICES.get(setting, SAMPLE_CHOICES["OOD mixed materials"]) return [str(i) for i in range(len(samples))] def sample_path_for_example_id(setting: str, object_count: str, example_id: str) -> str: if setting == "OOD mixed materials": samples = OOD_SAMPLE_CHOICES.get(str(object_count), OOD_SAMPLE_CHOICES["2 objects"]) else: samples = SAMPLE_CHOICES.get(setting, SAMPLE_CHOICES["OOD mixed materials"]) try: idx = int(str(example_id).strip()) except ValueError as exc: raise ValueError(f"Example must be an integer index, got {example_id!r}") from exc if idx < 0 or idx >= len(samples): raise ValueError(f"Example index {idx} is out of range for {setting!r}; valid range is 0..{len(samples) - 1}") return samples[idx] def _tail(text: str, max_chars: int = 18000) -> str: if len(text) <= max_chars: return text return "[log truncated]\n" + text[-max_chars:] def _timing_summary(log: str) -> str: values = dict(re.findall(r"\[timing\]\s+([A-Za-z0-9_\[\]\.]+)=([0-9.]+)", log)) hardware = re.findall(r"\[hardware\]\s+(.+)", log) attention = re.findall(r"\[attention\]\[rank=\d+\]\s+(.+)", log) lines = [] if hardware: lines.append("Hardware: " + hardware[-1]) if attention: lines.append("Attention: " + attention[-1]) if "sample[0].gen[0].inference_model_generate_s" in values: lines.append(f"Model inference: {values['sample[0].gen[0].inference_model_generate_s']} s") fields = [ ("checkpoint_load_s", "Checkpoint load"), ("model_setup_s", "Model setup"), ("sample[0].setup_s", "Input setup"), ("sample[0].gen[0].postprocess_save_npz_s", "Postprocess/save"), ("sample[0].gen[0].render_encode_s", "Render/encode"), ("engine_total_wall_s", "Engine total"), ("gradio_subprocess_wall_s", "Gradio subprocess wall"), ] for key, label in fields: if key in values: lines.append(f"{label}: {values[key]} s") if lines: return "\n".join(lines) return ( "No timing markers were found in the inference output.\n" "The Space may still be running an older build, or the inference process exited before timing was emitted." ) def _safe_preview_name(*parts: object) -> str: text = "_".join([PREVIEW_VERSION, *(str(part) for part in parts)]).lower() return "".join(ch if ch.isalnum() else "_" for ch in text).strip("_") def _object_names_from_metadata(metadata_path: Path) -> list[str]: import json with metadata_path.open("r", encoding="utf-8") as f: metadata = json.load(f) out: list[str] = [] for obj in metadata.get("objects", []): if isinstance(obj, dict): name = obj.get("name") or obj.get("mesh_used") or obj.get("mesh_source") or "" out.append(Path(str(name)).stem) return out def _color_for_object(index: int, object_name: str | None) -> tuple[float, float, float, float]: name = str(object_name or "").lower() for pattern, color in NAMED_COLORS.items(): if pattern in name: return color return RENDER_COLORS[int(index) % len(RENDER_COLORS)] def _with_alpha(color: tuple[float, float, float, float], alpha: float) -> tuple[float, float, float, float]: return (float(color[0]), float(color[1]), float(color[2]), float(alpha)) def _is_elastic_material(material: object) -> bool: if isinstance(material, str): return material.strip().lower() in {"elastic", "soft"} if isinstance(material, dict): kind = str(material.get("kind", "")).strip().lower() if kind in {"elastic", "soft"}: return True if kind in {"rigid", "hard"}: return False for key in ("effective_softness", "softness"): value = material.get(key) if isinstance(value, (int, float)): return float(value) >= 0.5 return False def _default_material_values(sample: str) -> list[str]: values = [DEFAULT_MATERIALS.get(obj, "elastic") for obj in objects_for_sample(sample)] while len(values) < 3: values.append("elastic") return values[:3] def _preview_object_alphas( setting: str, sample: str, metadata_path: Path, material_0: str = "", material_1: str = "", material_2: str = "", ) -> list[float]: import json if setting == "OOD mixed materials": defaults = _default_material_values(sample) materials = [ str(material_0 or defaults[0]), str(material_1 or defaults[1]), str(material_2 or defaults[2]), ] return [ELASTIC_RENDER_ALPHA if _is_elastic_material(material) else RIGID_RENDER_ALPHA for material in materials] with metadata_path.open("r", encoding="utf-8") as f: metadata = json.load(f) alphas: list[float] = [] for obj in metadata.get("objects", []): material = obj.get("material") if isinstance(obj, dict) else None alphas.append(ELASTIC_RENDER_ALPHA if _is_elastic_material(material) else RIGID_RENDER_ALPHA) return alphas def _shaded_facecolors(vertices, faces, base_color): import numpy as np light_direction = np.asarray([0.45, -0.65, 0.75], dtype=np.float32) tris = vertices[faces] normals = np.cross(tris[:, 1] - tris[:, 0], tris[:, 2] - tris[:, 0]) normals /= np.maximum(np.linalg.norm(normals, axis=1, keepdims=True), 1e-8) light = light_direction / np.linalg.norm(light_direction) intensity = 0.42 + 0.58 * np.clip(normals @ light, 0.0, 1.0) base = np.asarray(base_color, dtype=np.float32) facecolors = np.empty((faces.shape[0], 4), dtype=np.float32) facecolors[:, :3] = np.clip(base[:3][None, :] * intensity[:, None] + 0.10 * (1.0 - intensity[:, None]), 0.0, 1.0) facecolors[:, 3] = base[3] return facecolors def _faces_for_vertex_slice(faces, start: int, end: int): import numpy as np in_range = (faces >= int(start)) & (faces < int(end)) keep = np.all(in_range, axis=1) return faces[keep] - int(start) def _velocity_indices_for_object(speed, start: int, end: int, max_arrows: int): import numpy as np local = np.arange(int(start), int(end), dtype=np.int64) active = local[speed[local] > 1e-9] if active.size <= int(max_arrows): return active # Deterministic subsample across the object vertices so the preview does not become an arrow cloud. positions = np.linspace(0, active.size - 1, int(max_arrows)).round().astype(np.int64) return active[positions] def _draw_unit_bounds(ax) -> None: corners = [ (-1.0, -1.0, -1.0), (-1.0, -1.0, 1.0), (-1.0, 1.0, -1.0), (-1.0, 1.0, 1.0), (1.0, -1.0, -1.0), (1.0, -1.0, 1.0), (1.0, 1.0, -1.0), (1.0, 1.0, 1.0), ] edges = [ (0, 1), (0, 2), (0, 4), (3, 1), (3, 2), (3, 7), (5, 1), (5, 4), (5, 7), (6, 2), (6, 4), (6, 7), ] for start, end in edges: xs = [corners[start][0], corners[end][0]] ys = [corners[start][1], corners[end][1]] zs = [corners[start][2], corners[end][2]] ax.plot(xs, ys, zs, color=(0.18, 0.22, 0.28, 0.52), linewidth=0.9) def render_initial_preview( setting: str, object_count: str, example_id: str, material_0: str = "", material_1: str = "", material_2: str = "", ) -> str | None: sample = sample_path_for_example_id(setting, object_count, example_id) sample_dir = ROOT / sample obj_path = sample_dir / "meshes" / "combined_frame_000.obj" vel_path = sample_dir / "vertex_velocities" / "combined_frame_000.npy" metadata_path = sample_dir / "metadata.json" if not obj_path.is_file() or not vel_path.is_file(): return None out_dir = ROOT / ".inference_work" / "previews" out_dir.mkdir(parents=True, exist_ok=True) material_tag = "_".join(str(value or "default") for value in (material_0, material_1, material_2)) out_path = out_dir / f"{_safe_preview_name(setting, object_count, example_id, material_tag)}.png" if out_path.is_file() and out_path.stat().st_mtime >= max(obj_path.stat().st_mtime, vel_path.stat().st_mtime): return str(out_path) os.environ.setdefault("MPLCONFIGDIR", str(ROOT / ".inference_work" / "matplotlib")) import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np from mpl_toolkits.mplot3d.art3d import Poly3DCollection from physformer.data.multiobj_utils_multiobj import ( default_vertex_count_json_path, load_mesh_vertex_counts, scene_info_from_metadata, ) from physformer.data.obj_io import load_obj_vertices_faces vertices, faces = load_obj_vertices_faces(str(obj_path)) velocities = np.load(vel_path).astype(np.float32, copy=False) if velocities.shape != vertices.shape: raise ValueError(f"Velocity shape mismatch for {sample}: {velocities.shape} != {vertices.shape}") scene = scene_info_from_metadata( str(metadata_path), vertex_counts=load_mesh_vertex_counts(default_vertex_count_json_path()), max_num_objects=10, ) object_names = _object_names_from_metadata(metadata_path) object_alphas = _preview_object_alphas(setting, sample, metadata_path, material_0, material_1, material_2) fig = plt.figure(figsize=(6.0, 5.0), dpi=150, facecolor="#f7f8fb") ax = fig.add_subplot(111, projection="3d") ax.set_facecolor("#f7f8fb") for obj_idx, (start, end) in enumerate(scene.vertex_slices): obj_vertices = vertices[int(start) : int(end)] obj_faces = _faces_for_vertex_slice(faces, int(start), int(end)) if obj_vertices.size == 0 or obj_faces.size == 0: continue object_name = object_names[obj_idx] if obj_idx < len(object_names) else None alpha = object_alphas[obj_idx] if obj_idx < len(object_alphas) else RIGID_RENDER_ALPHA color = _with_alpha(_color_for_object(obj_idx, object_name), alpha) poly = Poly3DCollection( obj_vertices[obj_faces], facecolors=_shaded_facecolors(obj_vertices, obj_faces, color), edgecolors=MESH_EDGE_COLOR, linewidths=0.28, alpha=alpha, antialiased=True, ) ax.add_collection3d(poly) speed = np.linalg.norm(velocities, axis=1) arrow_cap = 20 if setting == "OOD mixed materials" else 10 arrow_indices = [ _velocity_indices_for_object(speed, int(start), int(end), arrow_cap) for start, end in scene.vertex_slices ] active = np.concatenate([idx for idx in arrow_indices if idx.size]) if any(idx.size for idx in arrow_indices) else np.empty((0,), dtype=np.int64) bbox_diag = float(np.linalg.norm(np.asarray([2.0, 2.0, 2.0], dtype=np.float32))) speed_ref = float(np.percentile(speed[active], 95)) if active.size else 0.0 scale = (0.20 * bbox_diag / speed_ref) if speed_ref > 0 and bbox_diag > 0 else 1.0 for obj_idx, active_obj in enumerate(arrow_indices): if not active_obj.size: continue object_name = object_names[obj_idx] if obj_idx < len(object_names) else None color = _color_for_object(obj_idx, object_name) v = velocities[active_obj] * scale ax.quiver( vertices[active_obj, 0], vertices[active_obj, 1], vertices[active_obj, 2], v[:, 0], v[:, 1], v[:, 2], color=color[:3], linewidth=0.85, arrow_length_ratio=0.22, normalize=False, ) _draw_unit_bounds(ax) ax.set_xlim(-1.0, 1.0) ax.set_ylim(-1.0, 1.0) ax.set_zlim(-1.0, 1.0) ax.set_box_aspect([1, 1, 1]) ax.view_init(elev=24, azim=-56) ax.set_title("Initial mesh per-vertex position and velocity", fontsize=10) ax.set_xlabel("x") ax.set_ylabel("y") ax.set_zlabel("z") ax.grid(True, linewidth=0.35, alpha=0.35) for axis in (ax.xaxis, ax.yaxis, ax.zaxis): axis.pane.set_facecolor((0.95, 0.96, 0.98, 0.72)) axis.pane.set_edgecolor((0.72, 0.75, 0.80, 0.50)) fig.tight_layout() fig.savefig(out_path, bbox_inches="tight") plt.close(fig) return str(out_path) def ensure_checkpoint() -> str: if CKPT_PATH.is_file(): return f"Checkpoint found: {CKPT_PATH}" CKPT_PATH.parent.mkdir(parents=True, exist_ok=True) token = os.environ.get("HF_TOKEN") or None downloaded = hf_hub_download( repo_id=CKPT_REPO_ID, filename=CKPT_FILENAME, local_dir=str(CKPT_PATH.parent), token=token, ) downloaded_path = Path(downloaded) if downloaded_path.resolve() != CKPT_PATH.resolve(): downloaded_path.replace(CKPT_PATH) return f"Downloaded checkpoint from {CKPT_REPO_ID}/{CKPT_FILENAME}" def objects_for_sample(sample: str) -> list[str]: sample = str(sample) if "/2obj_cow_horse/" in sample: return ["cow", "horse"] if "/3obj_teapot_fish_bunny/" in sample: return ["teapot", "fish", "bunny"] return [] def material_controls_for_sample(setting: str, object_count: str, example_id: str) -> tuple[dict, dict, dict]: sample = sample_path_for_example_id(setting, object_count, example_id) objects = objects_for_sample(sample) if setting == "OOD mixed materials" else [] updates: list[dict] = [] for idx in range(3): if idx < len(objects): obj = objects[idx] updates.append( gr.update( label=f"{obj} material", value=DEFAULT_MATERIALS[obj], visible=True, ) ) else: updates.append(gr.update(visible=False)) return tuple(updates) # type: ignore[return-value] def material_controls_and_preview(setting: str, object_count: str, example_id: str) -> tuple[dict, dict, dict, str | None]: sample = sample_path_for_example_id(setting, object_count, example_id) material_0, material_1, material_2 = _default_material_values(sample) return ( *material_controls_for_sample(setting, object_count, example_id), render_initial_preview(setting, object_count, example_id, material_0, material_1, material_2), ) def preview_for_materials( setting: str, object_count: str, example_id: str, material_0: str, material_1: str, material_2: str, ) -> str | None: return render_initial_preview(setting, object_count, example_id, material_0, material_1, material_2) def update_setting_controls(setting: str) -> tuple[dict, dict, dict, dict, dict, str | None]: if setting == "OOD mixed materials": object_count = "2 objects" choices = example_ids_for_setting(setting, object_count) example_id = choices[0] return ( gr.update(visible=True, value=object_count), gr.update(choices=choices, value=example_id), *material_controls_and_preview(setting, object_count, example_id), ) choices = example_ids_for_setting(setting) example_id = choices[0] return ( gr.update(visible=False, value="2 objects"), gr.update(choices=choices, value=example_id), *material_controls_and_preview(setting, "2 objects", example_id), ) def update_ood_object_count_controls(setting: str, object_count: str) -> tuple[dict, dict, dict, dict, str | None]: choices = example_ids_for_setting(setting, object_count) example_id = choices[0] return (gr.update(choices=choices, value=example_id), *material_controls_and_preview(setting, object_count, example_id)) def _ood_material_args(sample: str, material_0: str, material_1: str, material_2: str) -> list[str]: objects = objects_for_sample(sample) materials = [material_0, material_1, material_2] args: list[str] = [] for obj, material in zip(objects, materials): material = str(material).strip().lower() if material not in {"elastic", "rigid"}: raise ValueError(f"Invalid material for {obj}: {material!r}") args.extend([f"--{material}", obj]) return args def _command_for_example( setting: str, object_count: str, example_id: str, sampling_steps: int, material_0: str, material_1: str, material_2: str, ) -> list[str]: sample = sample_path_for_example_id(setting, object_count, example_id) common = [ sys.executable, "run_official_demo_inference.py", "--demo-root", sample, "--include", "all", "--generations", "1", "--num-sampling-steps", str(int(sampling_steps)), "--checkpoint", str(CKPT_PATH), "--device", "cuda", "--amp", os.environ.get("PHYSFORMER_AMP", "bf16"), "--overwrite", "--save-mp4", "--verbose", "--attention-debug", ] if setting == "OOD mixed materials": return common + _ood_material_args(sample, material_0, material_1, material_2) if setting == "In-distribution rigid": return common + ["--rigid", "all"] if setting == "In-distribution elastic": return common + ["--elastic", "all"] raise ValueError(f"Unknown setting: {setting}") def _latest_mp4_since(start_time: float) -> Path | None: candidates = sorted( [ path for path in ROOT.glob("**/inference.mp4") if ".inference_work" not in path.parts and path.stat().st_mtime >= start_time - 1.0 ], key=lambda path: path.stat().st_mtime, reverse=True, ) return candidates[0] if candidates else None DEMO_CSS = """ #generated-rollout { width: min(100%, 840px) !important; max-width: 840px !important; } #generated-rollout video { width: 100% !important; max-height: 480px !important; object-fit: contain !important; } """ @spaces.GPU(duration=120) def run_inference( setting: str, object_count: str, example_id: str, sampling_steps: int, material_0: str, material_1: str, material_2: str, setup_log: str, ) -> tuple[str | None, str, str]: if not CKPT_PATH.is_file(): log = setup_log + "\nCheckpoint is missing; click Run again after the download finishes." return None, "Checkpoint missing.", log start_time = time.time() subprocess_t0 = time.perf_counter() cmd = _command_for_example(setting, str(object_count), str(example_id), int(sampling_steps), material_0, material_1, material_2) env = os.environ.copy() env.setdefault("PYTHONUNBUFFERED", "1") env.setdefault("MPLCONFIGDIR", str(ROOT / ".inference_work" / "matplotlib")) proc = subprocess.run( cmd, cwd=ROOT, env=env, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, check=False, timeout=900, ) subprocess_s = time.perf_counter() - subprocess_t0 log = setup_log + "\n\n$ " + " ".join(cmd) + "\n" + proc.stdout + f"\n[timing] gradio_subprocess_wall_s={subprocess_s:.3f}" mp4 = _latest_mp4_since(start_time) if proc.returncode != 0: fail_log = log + f"\nInference failed with exit code {proc.returncode}." return None, _timing_summary(fail_log), _tail(fail_log) if mp4 is None: missing_log = log + "\nInference finished, but no inference.mp4 was found." return None, _timing_summary(missing_log), _tail(missing_log) final_log = log + f"\nGenerated video: {mp4.relative_to(ROOT)}" return str(mp4), _timing_summary(final_log), _tail(final_log) with gr.Blocks(title="PhysFormer", css=DEMO_CSS) as demo: gr.Markdown( """ # PhysiFormer Minimal ZeroGPU Demo Select the example, material conditions, and denoising step numbers to run PhysiFormer inference. This demo runs on Hugging Face ZeroGPU, dynamically allocating a 48GB NVIDIA RTX Pro 6000 Blackwell GPU for each generation. """ ) with gr.Row(): setting = gr.Dropdown( choices=["OOD mixed materials", "In-distribution rigid", "In-distribution elastic"], value="OOD mixed materials", label="Setting", ) object_count = gr.Dropdown( choices=["2 objects", "3 objects"], value="2 objects", label="Object Count", visible=True, ) example_id = gr.Dropdown( choices=example_ids_for_setting("OOD mixed materials"), value="0", label="Example", ) sampling_steps = gr.Slider(5, 50, value=10, step=1, label="Denoising steps") with gr.Row(): material_0 = gr.Dropdown( choices=["elastic", "rigid"], value="rigid", label="cow material", visible=True, ) material_1 = gr.Dropdown( choices=["elastic", "rigid"], value="elastic", label="horse material", visible=True, ) material_2 = gr.Dropdown( choices=["elastic", "rigid"], value="elastic", label="material", visible=False, ) preview = gr.Image( value=render_initial_preview("OOD mixed materials", "2 objects", "0", "rigid", "elastic", "elastic"), label="Initial mesh per-vertex position and velocity", type="filepath", height=420, ) run_button = gr.Button("Generate", variant="primary") video = gr.Video(label="Generated rollout", height=480, width=840, elem_id="generated-rollout") timing = gr.Textbox(label="Timing summary", lines=8, value="Run a rollout to see timing.") log = gr.Textbox(label="Log", lines=18) setting.change( update_setting_controls, inputs=setting, outputs=[object_count, example_id, material_0, material_1, material_2, preview], ) object_count.change( update_ood_object_count_controls, inputs=[setting, object_count], outputs=[example_id, material_0, material_1, material_2, preview], ) example_id.change( material_controls_and_preview, inputs=[setting, object_count, example_id], outputs=[material_0, material_1, material_2, preview], ) for material_control in (material_0, material_1, material_2): material_control.change( preview_for_materials, inputs=[setting, object_count, example_id, material_0, material_1, material_2], outputs=preview, ) run_button.click(ensure_checkpoint, outputs=log).then( run_inference, inputs=[setting, object_count, example_id, sampling_steps, material_0, material_1, material_2, log], outputs=[video, timing, log], ) if __name__ == "__main__": demo.queue(default_concurrency_limit=1, max_size=8).launch()