modal_sam3d.py

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)")