import modal import textwrap # Volume that already contains: # sam-3d-objects/checkpoints/pipeline.yaml # AND will now cache DINOv2 / other model weights volume = modal.Volume.from_name("sam3d-weights", create_if_missing=False) # --------------------------------------------------------------------------- # Image build: CUDA base + PyTorch + PyTorch3D + SAM-3D repo + deps # --------------------------------------------------------------------------- sam3d_image = ( modal.Image.from_registry( "nvidia/cuda:12.4.1-devel-ubuntu22.04", add_python="3.11", # Python 3.11 ) .apt_install( "git", "g++", "gcc", "clang", "build-essential", "libgl1-mesa-glx", "libglib2.0-0", "libopenexr-dev", "wget", ) # STEP 1: Install PyTorch CUDA 12.4 stack (hard fail if broken) .pip_install( "torch==2.5.1", "torchvision", "torchaudio", index_url="https://download.pytorch.org/whl/cu124", ) # STEP 1.5: Build deps (needed for PyTorch3D / SAM-3D) .pip_install( "fvcore", "iopath", "numpy", "ninja", "setuptools", "wheel", ) # STEP 2: Clone the SAM-3D Objects repo .run_commands( "echo '[STEP 2] Cloning facebookresearch/sam-3d-objects' && " "git clone https://github.com/facebookresearch/sam-3d-objects.git /sam3d" ) # STEP 2.1: Remove nvidia-pyindex from pyproject so pip doesn't try to build it .run_commands( "echo '[STEP 2.1] Removing nvidia-pyindex from pyproject.toml (if present)' && " "cd /sam3d && " "if [ -f pyproject.toml ]; then " " sed -i '/nvidia-pyindex/d' pyproject.toml; " "fi" ) # STEP 3: Install [p3d] extras (PyTorch3D-related deps), fail-soft .run_commands( "echo '[STEP 3] Installing sam-3d-objects extra [p3d]' && " "cd /sam3d && " "PIP_EXTRA_INDEX_URL='https://pypi.ngc.nvidia.com https://download.pytorch.org/whl/cu124' " "pip install -e '.[p3d]' " "|| echo '[WARN] [p3d] extras failed to install, continuing without them.'" ) # STEP 4: Install [inference] extras (Kaolin etc.), fail-soft .run_commands( "echo '[STEP 4] Installing sam-3d-objects extra [inference] (includes Kaolin etc.)' && " "cd /sam3d && " "PIP_FIND_LINKS='https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.5.1_cu121.html' " "pip install -e '.[inference]' " "|| echo '[WARN] [inference] extras failed to install, continuing without them.'" ) # STEP 5: Helper libs (open3d, trimesh, seaborn) – fail-soft .run_commands( "echo '[STEP 5] Installing helper libraries: open3d, trimesh, seaborn' && " "pip install open3d trimesh seaborn " "|| echo '[WARN] Helper libs (open3d/trimesh/seaborn) failed to install, continuing.'" ) # STEP 5.5: Config libs required by inference.py (omegaconf, hydra-core) .run_commands( "echo '[STEP 5.5] Installing config libraries: omegaconf, hydra-core' && " "pip install omegaconf hydra-core " "|| echo '[WARN] omegaconf/hydra-core failed to install, continuing.'" ) # STEP 5.6: Install utils3d explicitly (inference.py imports this) .run_commands( "echo '[STEP 5.6] Installing utils3d' && " "pip install " "'git+https://github.com/EasternJournalist/utils3d.git@3913c65d81e05e47b9f367250cf8c0f7462a0900' " "|| echo '[WARN] utils3d failed to install, continuing.'" ) # STEP 5.7: Installing gradio (inference.py imports this) .run_commands( "echo '[STEP 5.7] Installing gradio' && " "pip install gradio " "|| echo '[WARN] gradio failed to install, continuing.'" ) .run_commands( "echo '[STEP 5.8] Installing kaolin from NVIDIA index' && " "pip install kaolin " "-f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.5.1_cu121.html " "|| echo '[WARN] kaolin install failed, continuing.'" ) # STEP 5.9: Install loguru (needed by inference_pipeline_pointmap) .run_commands( "echo '[STEP 5.9] Installing loguru' && " "pip install loguru " "|| echo '[WARN] loguru failed to install, continuing.'" ) # STEP 5.91: Install timm (vision transformer lib) .run_commands( "echo '[STEP 5.91] Installing timm' && " "pip install timm " "|| echo '[WARN] timm failed to install, continuing.'" ) # STEP 5.8: Install PyTorch3D from GitHub @stable, using the pattern that worked for you .run_commands( "echo '[STEP 5.92] Installing PyTorch3D from GitHub @stable (no build isolation, no deps)' && " "python -c 'import pytorch3d' 2>/dev/null && " "echo 'PyTorch3D already installed, skipping...' || ( " "export FORCE_CUDA=1 && " "export TORCH_CUDA_ARCH_LIST='8.0;8.6;8.9;9.0' && " "pip install --no-build-isolation --no-deps " "\"git+https://github.com/facebookresearch/pytorch3d.git@stable\" " ")" ) .run_commands( "cd /sam3d && pip install '.[dev]' --no-deps" ) .run_commands("pip install optree") .run_commands("pip install astor==0.8.1") .run_commands("pip install opencv-python") .run_commands("pip install lightning") .run_commands("pip install spconv-cu121==2.3.8") .run_commands("pip install psutil && pip install --no-build-isolation flash_attn==2.8.3 || echo '[WARN] flash_attn failed'") .run_commands("pip install xatlas==0.0.9") .run_commands("pip install pyvista") .run_commands("pip install pymeshfix==0.17.0") .run_commands("pip install igraph") .run_commands("pip install easydict") .run_commands("pip install igraph") .run_commands( "export TORCH_CUDA_ARCH_LIST='8.0;8.6;8.9;9.0' && " "pip install --no-build-isolation 'git+https://github.com/nerfstudio-project/gsplat.git@2323de5905d5e90e035f792fe65bad0fedd413e7'" ) .run_commands("pip install igraph") .run_commands("pip install 'git+https://github.com/microsoft/MoGe.git@a8c37341bc0325ca99b9d57981cc3bb2bd3e255b'") .run_commands("pip install imageio") # STEP 6: Patch hydra – skip if it fails .run_commands( "echo '[STEP 6] Patching hydra' && " "cd /sam3d && " "./patching/hydra " "|| echo '[WARN] Hydra patch failed, continuing without patch.'" ) ) app = modal.App("sam3d-objects-inference", image=sam3d_image) # --------------------------------------------------------------------------- # Runtime helper: minimal pytorch3d stub so SAM-3D imports work (fallback) @app.cls( image=sam3d_image, gpu="A10G", timeout=600, volumes={"/weights": volume}, scaledown_window=300, # renamed from container_idle_timeout enable_memory_snapshot=True, # required for snap=True ) class SAM3DModel: @modal.enter(snap=True) def setup(self): """Model loads once when container starts. snap=True caches the loaded state.""" import os import sys import math import types import torch # Cache setup CACHE_DIR = "/weights/model_cache" os.makedirs(CACHE_DIR, exist_ok=True) os.environ["TORCH_HOME"] = CACHE_DIR os.environ["TORCH_HUB"] = os.path.join(CACHE_DIR, "hub") os.environ["HF_HOME"] = os.path.join(CACHE_DIR, "huggingface") os.environ["TRANSFORMERS_CACHE"] = os.path.join(CACHE_DIR, "huggingface") os.environ["XDG_CACHE_HOME"] = CACHE_DIR os.environ["TIMM_CACHE"] = os.path.join(CACHE_DIR, "timm") os.environ.setdefault("CUDA_HOME", "/usr/local/cuda") os.environ.setdefault("CONDA_PREFIX", "/usr/local/cuda") # pytorch3d stub try: import pytorch3d except Exception: pkg = types.ModuleType("pytorch3d") transforms_mod = types.ModuleType("pytorch3d.transforms") renderer_mod = types.ModuleType("pytorch3d.renderer") def _quat_conj(q): w, x, y, z = q.unbind(-1) return torch.stack((w, -x, -y, -z), dim=-1) def quaternion_multiply(q1, q2): w1, x1, y1, z1 = q1.unbind(-1) w2, x2, y2, z2 = q2.unbind(-1) return torch.stack([w1*w2-x1*x2-y1*y2-z1*z2, w1*x2+x1*w2+y1*z2-z1*y2, w1*y2-x1*z2+y1*w2+z1*x2, w1*z2+x1*y2-y1*x2+z1*w2], dim=-1) def quaternion_invert(q): return _quat_conj(q) / (q.norm(dim=-1, keepdim=True) ** 2 + 1e-8) transforms_mod.quaternion_multiply = quaternion_multiply transforms_mod.quaternion_invert = quaternion_invert class Transform3d: def __init__(self, matrix=None, device=None): self.matrix = torch.eye(4, device=device).unsqueeze(0) if matrix is None else matrix def compose(self, other): return Transform3d(other.matrix @ self.matrix) def transform_points(self, points): if points.dim() == 2: pts = torch.cat([points, torch.ones(points.shape[0], 1, device=points.device)], dim=-1) return (self.matrix[0] @ pts.T).T[..., :3] elif points.dim() == 3: B, N, _ = points.shape pts = torch.cat([points, torch.ones(B, N, 1, device=points.device)], dim=-1) mat = self.matrix.expand(B, -1, -1) if self.matrix.shape[0] == 1 and B > 1 else self.matrix return torch.bmm(mat, pts.transpose(1, 2)).transpose(1, 2)[..., :3] transforms_mod.Transform3d = Transform3d def look_at_view_transform(dist=1.0, elev=0.0, azim=0.0, device=None): dist_t = torch.tensor([dist], device=device, dtype=torch.float32) elev_rad = torch.tensor([elev], device=device) * math.pi / 180.0 azim_rad = torch.tensor([azim], device=device) * math.pi / 180.0 x = dist_t * torch.cos(elev_rad) * torch.sin(azim_rad) y = dist_t * torch.sin(elev_rad) z = dist_t * torch.cos(elev_rad) * torch.cos(azim_rad) cam_pos = torch.stack([x, y, z], dim=-1) up = torch.tensor([[0.0, 1.0, 0.0]], device=device) z_axis = torch.nn.functional.normalize(cam_pos, dim=-1) x_axis = torch.nn.functional.normalize(torch.cross(up, z_axis, dim=-1), dim=-1) y_axis = torch.cross(z_axis, x_axis, dim=-1) R = torch.stack([x_axis, y_axis, z_axis], dim=-1) T = -torch.bmm(R, cam_pos.unsqueeze(-1)).squeeze(-1) return R, T renderer_mod.look_at_view_transform = look_at_view_transform pkg.transforms = transforms_mod pkg.renderer = renderer_mod sys.modules["pytorch3d"] = pkg sys.modules["pytorch3d.transforms"] = transforms_mod sys.modules["pytorch3d.renderer"] = renderer_mod sys.path.insert(0, "/sam3d") sys.path.insert(0, "/sam3d/notebook") from inference import Inference, load_image self.load_image = load_image self.model = Inference("/weights/sam-3d-objects/checkpoints/pipeline.yaml", compile=False) print("[SETUP] Model loaded!") @modal.method() def reconstruct(self, image_bytes: bytes, mask_bytes: bytes = None) -> tuple[bytes, bytes]: import os, io, tempfile, shutil import numpy as np from PIL import Image import torch temp_dir = tempfile.mkdtemp() image_path = os.path.join(temp_dir, "image.png") mask_path = os.path.join(temp_dir, "mask.png") with open(image_path, 'wb') as f: f.write(image_bytes) pil_image = Image.open(image_path) if mask_bytes is not None: with open(mask_path, 'wb') as f: f.write(mask_bytes) mask = np.array(Image.open(mask_path).convert('L')) elif pil_image.mode == 'RGBA': alpha = np.array(pil_image)[:, :, 3] mask = (alpha > 128).astype(np.uint8) * 255 pil_image = pil_image.convert('RGB') pil_image.save(image_path) else: raise ValueError("Provide either: 1) separate mask_bytes, or 2) RGBA image with alpha mask") if np.sum(mask > 0) < 100: raise ValueError("Mask too small!") image = self.load_image(image_path) if mask.shape[0] != image.shape[0] or mask.shape[1] != image.shape[1]: mask = np.array(Image.fromarray(mask).resize((image.shape[1], image.shape[0]), Image.NEAREST)) with torch.inference_mode(): output = self.model(image, mask, seed=42) shutil.rmtree(temp_dir, ignore_errors=True) ply_buffer = io.BytesIO() output["gs"].save_ply(ply_buffer) glb_bytes = None if "mesh" in output and output["mesh"]: import trimesh mesh = output["mesh"][0] if isinstance(output["mesh"], list) else output["mesh"] glb_bytes = trimesh.Trimesh( vertices=mesh.vertices.cpu().numpy(), faces=mesh.faces.cpu().numpy() ).export(file_type="glb") return ply_buffer.getvalue(), glb_bytes @app.local_entrypoint() def main( input_path: str = "sam3d_1.png", mask_path: str = "sam3d_1gray.png", output_path: str = "output_model.ply", ): """ Local test: # With RGBA image (mask in alpha): modal run modal_sam3d.py --input-path image_rgba.png # With separate mask file (official pattern): modal run modal_sam3d.py --input-path image.png --mask-path mask.png """ from pathlib import Path input_file = Path(input_path) if not input_file.exists(): print(f"[LOCAL] ERROR: Input image not found: {input_file.resolve()}") return mask_bytes = None if mask_path: mask_file = Path(mask_path) if mask_file.exists(): mask_bytes = mask_file.read_bytes() print(f"[LOCAL] Using separate mask file: {mask_file}") else: print(f"[LOCAL] WARNING: Mask file not found: {mask_file}") print(f"[LOCAL] Sending {input_file} to SAM-3D on Modal...") model = SAM3DModel() ply_bytes, glb_bytes = model.reconstruct.remote(input_file.read_bytes(), mask_bytes) output_file = Path(output_path) output_file.write_bytes(ply_bytes) if glb_bytes: glb_file = Path(output_path).with_suffix(".glb") glb_file.write_bytes(glb_bytes) print(f"[LOCAL] Saved mesh to: {glb_file}") print(f"[LOCAL] Saved 3D model to: {output_file.resolve()} ({len(ply_bytes)} bytes)")