app.py

"""
SAM 3D Objects MCP Server
Image (+ text prompt) → 3D Object (GLB)

SAM3 concept segmentation + SAM 3D Objects reconstruction on ZeroGPU.

torch is provided by ZeroGPU. kaolin and pytorch3d are replaced by pure-python
stubs covering exactly the surface the pipeline touches — texture baking, mesh
postprocessing and layout postprocessing are disabled, so their compiled ops
are never called. Packages that need torch at install time are installed at
startup. The attention backend is pinned to sdpa via import order in
sam3d_inference.py (flash_attn is not available on ZeroGPU).

Models are loaded at module level per the official ZeroGPU pattern: a CUDA
emulation mode is active outside @spaces.GPU, and startup placements get
optimized transfers when the real GPU attaches.
"""
import os
import subprocess
import sys
import tempfile
from pathlib import Path

os.environ.setdefault("CUDA_HOME", "/usr/local/cuda")
os.environ.setdefault("CONDA_PREFIX", "/usr/local")
os.environ["LIDRA_SKIP_INIT"] = "true"
os.environ["ATTN_BACKEND"] = "sdpa"
os.environ["SPARSE_ATTN_BACKEND"] = "sdpa"
os.environ["SPARSE_BACKEND"] = "spconv"

# MUST import spaces before torch
import spaces
import gradio as gr
import numpy as np
from huggingface_hub import snapshot_download, login

if os.environ.get("HF_TOKEN"):
    login(token=os.environ["HF_TOKEN"])

APP_ROOT = Path(__file__).parent
for stub in ["kaolin_stub", "pytorch3d_stub"]:
    sys.path.insert(0, str(APP_ROOT / stub))


def _pip(*args):
    r = subprocess.run(
        [sys.executable, "-m", "pip", "install", "--no-cache-dir", *args],
        capture_output=True, text=True, timeout=1200,
    )
    print(f"  pip {'OK' if r.returncode == 0 else 'FAIL'}: {args[-1][:70]}")
    if r.returncode != 0:
        print(f"    {r.stderr[-500:]}")
    return r.returncode == 0


print("=== Runtime installs (need torch present) ===")
# utils3d pinned to the commit MoGe expects — newer commits dropped points_to_normals
_pip("--no-deps", "git+https://github.com/EasternJournalist/utils3d.git@3913c65d81e05e47b9f367250cf8c0f7462a0900")
_pip("--no-deps", "git+https://github.com/microsoft/MoGe.git@a8c37341bc0325ca99b9d57981cc3bb2bd3e255b")
# no prebuilt gsplat wheels beyond torch 2.4 — the PyPI sdist installs in JIT
# mode; kernels never compile here because rendering paths are disabled
_pip("--no-deps", "gsplat")

SAM3D_PATH = APP_ROOT / "sam-3d-objects"
if not SAM3D_PATH.exists():
    print("Cloning sam-3d-objects...")
    subprocess.run(["git", "clone", "--depth", "1",
                    "https://github.com/facebookresearch/sam-3d-objects.git",
                    str(SAM3D_PATH)], check=True)
sys.path.insert(0, str(SAM3D_PATH))

print("Downloading SAM3D checkpoints...")
CKPT_DIR = snapshot_download(repo_id="facebook/sam-3d-objects",
                             token=os.environ.get("HF_TOKEN"))
hf_ckpt = Path(CKPT_DIR) / "checkpoints"
local_ckpt = SAM3D_PATH / "checkpoints" / "hf"
if hf_ckpt.exists() and not local_ckpt.exists():
    local_ckpt.parent.mkdir(parents=True, exist_ok=True)
    local_ckpt.symlink_to(hf_ckpt)
CONFIG_PATH = str(local_ckpt / "pipeline.yaml")

print("=== Startup complete ===")

# Model construction needs a real GPU (the constructors run device-mixing
# tensor ops that ZeroGPU's startup CUDA emulation rejects), so models load
# inside @spaces.GPU and are cached for reuse.
SAM3 = None
SAM3D = None


def _get_sam3():
    global SAM3
    if SAM3 is None:
        import torch
        from transformers import Sam3Model, Sam3Processor
        model = Sam3Model.from_pretrained("facebook/sam3").to("cuda").eval()
        processor = Sam3Processor.from_pretrained("facebook/sam3")
        SAM3 = (model, processor)
    return SAM3


def _get_sam3d():
    global SAM3D
    if SAM3D is None:
        from sam3d_inference import SAM3DInference
        SAM3D = SAM3DInference(CONFIG_PATH)
    return SAM3D


def _segment(image_np, prompt):
    """SAM3 concept segmentation: returns the best-scoring mask for the
    prompt as a bool array, plus the number of instances found."""
    import torch
    model, processor = _get_sam3()
    inputs = processor(images=image_np, text=prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model(**inputs)
    results = processor.post_process_instance_segmentation(
        outputs, threshold=0.5, mask_threshold=0.5,
        target_sizes=inputs.get("original_sizes").tolist(),
    )[0]
    masks = results["masks"]
    if len(masks) == 0:
        return None, 0
    best = int(torch.as_tensor(results["scores"]).argmax())
    return masks[best].cpu().numpy().astype(bool), len(masks)


def _export_result(result, out_dir):
    """Export the pipeline result: prefer the ready-made GLB mesh, fall back
    to the raw gaussian splat PLY."""
    if not isinstance(result, dict):
        return None
    glb_mesh = result.get("glb")
    if glb_mesh is not None:
        path = f"{out_dir}/object.glb"
        glb_mesh.export(path)
        return path
    gs = result.get("gaussian")
    if gs is not None:
        gs = gs[0] if isinstance(gs, (list, tuple)) else gs
        if hasattr(gs, "save_ply"):
            path = f"{out_dir}/object.ply"
            gs.save_ply(path)
            return path
    return None


@spaces.GPU(duration=30)
def diagnose():
    """Check GPU and readiness of all models."""
    import torch
    lines = [f"torch={torch.__version__}", f"cuda={torch.cuda.is_available()}"]
    if torch.cuda.is_available():
        lines.append(f"gpu={torch.cuda.get_device_name()}")
    lines.append(f"SAM3: {'loaded' if SAM3 is not None else 'loads on first reconstruct'}")
    lines.append(f"SAM3D: {'loaded' if SAM3D is not None else 'loads on first reconstruct'}")
    lines.append(f"config: {Path(CONFIG_PATH).exists()}")
    return "\n".join(lines)


@spaces.GPU(duration=180)
def reconstruct_objects(image: np.ndarray, prompt: str = "object", progress=gr.Progress()):
    """
    Segment the object described by the text prompt with SAM3 and
    reconstruct it in 3D with SAM 3D Objects.

    Args:
        image: Input RGB image
        prompt: Short noun phrase describing what to reconstruct
            (e.g. "the chair"). Default "object" picks the most
            prominent object.

    Returns:
        tuple: (model_path, preview_image, status)
    """
    if image is None:
        return None, None, "❌ No image provided"
    prompt = (prompt or "object").strip()
    try:
        import time
        import torch
        t0 = time.time()
        print(f"GPU: {torch.cuda.get_device_name()}")

        progress(0.05, desc=f"Segmenting '{prompt}' (SAM3)...")
        image_np = np.asarray(image)
        best_mask, n_masks = _segment(image_np, prompt)
        if best_mask is None:
            return None, image_np, f"⚠️ No '{prompt}' found in image"

        preview = image_np.copy()
        preview[best_mask] = (preview[best_mask] * 0.5 + np.array([0, 255, 0]) * 0.5).astype(np.uint8)
        print(f"  {n_masks} instances of '{prompt}' ({time.time()-t0:.0f}s)")

        progress(0.25, desc="Reconstructing 3D (SAM 3D Objects, ~1-2 min)...")
        result = _get_sam3d()(image=image_np, mask=best_mask, seed=42)
        print(f"  Reconstructed ({time.time()-t0:.0f}s)")
        if result is None:
            return None, preview, "⚠️ Reconstruction returned None"

        progress(0.9, desc="Exporting GLB...")
        model_path = _export_result(result, tempfile.mkdtemp())
        print(f"  Exported ({time.time()-t0:.0f}s): {model_path}")
        if model_path is None:
            keys = list(result.keys()) if isinstance(result, dict) else dir(result)
            return None, preview, f"⚠️ Cannot extract 3D data. Keys: {keys}"

        import trimesh
        try:
            n_faces = len(trimesh.load(model_path, force="mesh").faces)
        except Exception:
            n_faces = 0
        return model_path, preview, f"✓ {n_masks} × '{prompt}' found, {n_faces:,} faces ({int(time.time()-t0)}s)"
    except Exception:
        import traceback
        tb = traceback.format_exc()
        print(tb)
        return None, None, f"❌ Error:\n{tb[-1500:]}"


with gr.Blocks(title="SAM 3D Objects MCP") as demo:
    gr.Markdown("""
    # 📦 SAM 3D Objects
    **Image (+ text prompt) → 3D Object (GLB)** · powered by Meta's SAM3 + SAM 3D Objects · usable as MCP server
    """)

    with gr.Tab("Reconstruct"):
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Input Image", type="numpy")
                input_prompt = gr.Textbox(
                    label="What to reconstruct", value="object",
                    placeholder="e.g. the chair, yellow car, laptop")
                btn = gr.Button("🚀 Segment & Reconstruct", variant="primary", size="lg")
            with gr.Column():
                preview = gr.Image(label="Segmented Object", type="numpy", interactive=False)
                status = gr.Textbox(label="Status")
        with gr.Row():
            output_model = gr.Model3D(label="3D Preview", height=420)
            with gr.Column():
                output_file = gr.File(label="Download")
        btn.click(reconstruct_objects, inputs=[input_image, input_prompt],
                  outputs=[output_model, preview, status])
        output_model.change(lambda x: x, inputs=[output_model], outputs=[output_file],
                            api_name=False)

    with gr.Tab("Guide"):
        gr.Markdown("""
        ## How it works

        1. **Upload an image** — photos and renders both work; one clearly visible object gives the best result.
        2. **Describe what to reconstruct** — a short noun phrase like `robot`, `the red chair`, `laptop`.
           The default `object` picks the most prominent thing in the image.
        3. **Segment & Reconstruct** — SAM3 finds the best-matching instance (green preview),
           SAM 3D Objects rebuilds it as a vertex-colored GLB mesh you can preview and download.

        The first request after a restart loads both models and takes ~2 minutes extra; after that
        a reconstruction takes roughly 1–2 minutes.

        ### Prompt tips
        - Short English noun phrases work best: `sofa`, `yellow school bus`, `coffee mug`
        - If several matching objects exist, the highest-confidence one is used
        - Getting the wrong object? Make the prompt more specific (`the left chair` won't help —
          SAM3 matches concepts, not positions — but `wooden chair` will)

        ## Use as MCP server

        Add to your MCP client (Claude Code, Claude Desktop, Cursor, ...):

        ```json
        {
          "mcpServers": {
            "sam3d-objects": {
              "url": "https://dev-bjoern-sam3d-objects-mcp.hf.space/gradio_api/mcp/"
            }
          }
        }
        ```

        Or with the Claude Code CLI:

        ```bash
        claude mcp add --transport http sam3d-objects https://dev-bjoern-sam3d-objects-mcp.hf.space/gradio_api/mcp/
        ```

        For stdio-only clients (e.g. Claude Desktop) use
        `npx mcp-remote https://dev-bjoern-sam3d-objects-mcp.hf.space/gradio_api/mcp/ --transport streamable-http`.

        The `reconstruct_objects` tool takes an image (URL or uploaded file) and a text
        prompt and returns the GLB, the segmentation preview and a status line.
        Full API details: the **"Use via API or MCP"** link in the page footer.
        """)

    with gr.Tab("Diagnose"):
        diag_btn = gr.Button("Diagnose GPU & Models")
        diag_out = gr.Textbox(lines=8, label="Diagnostics")
        diag_btn.click(diagnose, outputs=[diag_out])


if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Ocean())