Spaces:

HaiderCG
/

sam3d-lab

Running on Zero

Lower ZeroGPU durations; account for 1.5x duration factor

ee28d62 23 days ago

24 kB

"""Prompt-guided SAM 3 → SAM 3D Objects demo for Hugging Face ZeroGPU."""

from __future__ import annotations

import os

# Configure backends before importing the model packages. ZeroGPU requires this
# import to precede torch so it can install its CUDA compatibility layer.
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"

import spaces

import subprocess
import sys
import tempfile
import time
import traceback
import uuid
from pathlib import Path
from typing import Any

import gradio as gr
import numpy as np
from huggingface_hub import snapshot_download
from PIL import Image


ROOT = Path(__file__).resolve().parent
SAM3D_REVISION = "f91db41"
SAM3D_CHECKPOINT_ID = "facebook/sam-3d-objects"
SAM3_CHECKPOINT_ID = "facebook/sam3"
HF_TOKEN = os.environ.get("HF_TOKEN")

# The official SAM 3D repo imports these optional libraries at module import
# time, although the inference path below disables their post-processing.
for _shim in ("kaolin_stub", "pytorch3d_stub"):
    sys.path.insert(0, str(ROOT / "compat" / _shim))


def _run(
    command: list[str],
    *,
    timeout: int = 1200,
    env: dict[str, str] | None = None,
) -> None:
    """Run setup work outside ZeroGPU allocation and show useful errors."""
    result = subprocess.run(
        command,
        check=False,
        capture_output=True,
        text=True,
        timeout=timeout,
        env=env,
    )
    if result.returncode:
        tail = (result.stderr or result.stdout)[-5000:]
        raise RuntimeError(f"Command failed ({result.returncode}): {tail}")


def _install_runtime_packages() -> None:
    """Install only packages whose setup needs the platform-provided torch."""
    marker = Path(tempfile.gettempdir()) / ".sam3d-runtime-deps-ready"
    if marker.exists():
        return

    requirements = ROOT / "runtime-requirements.txt"
    if not requirements.exists():
        raise RuntimeError("runtime-requirements.txt is missing")
    print("Installing SAM 3D packages that require ZeroGPU's torch runtime...")
    _run(
        [
            sys.executable,
            "-m",
            "pip",
            "install",
            "--no-cache-dir",
            "--no-deps",
            "-r",
            str(requirements),
        ],
        timeout=1800,
    )
    # gsplat's setup.py imports torch while preparing its wheel. PEP 517's
    # isolated build environment hides the platform's installed torch, so
    # disable isolation and use gsplat's documented no-CUDA-build/JIT mode.
    gsplat_env = os.environ.copy()
    gsplat_env["BUILD_NO_CUDA"] = "1"
    _run(
        [
            sys.executable,
            "-m",
            "pip",
            "install",
            "--no-cache-dir",
            "--no-deps",
            "--no-build-isolation",
            "git+https://github.com/nerfstudio-project/gsplat.git@2323de5905d5e90e035f792fe65bad0fedd413e7",
        ],
        timeout=1800,
        env=gsplat_env,
    )
    marker.write_text("ready\n", encoding="utf-8")


SETUP_ERROR: str | None = None
CONFIG_PATH: Path | None = None
SAM3_CACHE_DIR: Path | None = None


def _prepare_runtime() -> None:
    global SETUP_ERROR, CONFIG_PATH, SAM3_CACHE_DIR
    try:
        _install_runtime_packages()
        from sam3d_objects import __file__ as _sam3d_package_file  # noqa: F401

        if not HF_TOKEN:
            SETUP_ERROR = "Add a read token as the HF_TOKEN Space secret."
            return

        # Download once at Space startup. Checkpoint transfer is CPU/network
        # work and therefore does not occupy a ZeroGPU allocation.
        SAM3_CACHE_DIR = Path(
            snapshot_download(
                repo_id=SAM3_CHECKPOINT_ID,
                token=HF_TOKEN,
            )
        )
        cache_dir = Path(
            snapshot_download(
                repo_id=SAM3D_CHECKPOINT_ID,
                token=HF_TOKEN,
            )
        )
        checkpoint_dir = cache_dir / "checkpoints"
        config = checkpoint_dir / "pipeline.yaml"
        if not config.is_file():
            raise RuntimeError(
                f"Expected pipeline config at {config}; verify model access and files."
            )
        # Keep the config in its original Hub snapshot directory. snapshot_download
        # represents large checkpoint files with relative symlinks to the cache's
        # blobs directory; symlinking only this subdirectory elsewhere breaks those
        # relative links and causes FileNotFoundError while loading weights.
        CONFIG_PATH = config
        from omegaconf import OmegaConf

        checkpoint_config = OmegaConf.load(CONFIG_PATH)
        for key in (
            "ss_generator_config_path",
            "ss_generator_ckpt_path",
            "slat_generator_config_path",
            "slat_generator_ckpt_path",
            "ss_decoder_config_path",
            "ss_decoder_ckpt_path",
            "slat_decoder_gs_config_path",
            "slat_decoder_gs_ckpt_path",
            "slat_decoder_gs_4_config_path",
            "slat_decoder_gs_4_ckpt_path",
            "slat_decoder_mesh_config_path",
            "slat_decoder_mesh_ckpt_path",
        ):
            relative_path = checkpoint_config.get(key)
            if relative_path:
                checkpoint_file = CONFIG_PATH.parent / str(relative_path)
                if not checkpoint_file.is_file():
                    raise RuntimeError(
                        f"SAM 3D checkpoint is missing or unreadable: "
                        f"{checkpoint_file} (from {key})."
                    )
        print(f"Using SAM 3D pipeline config: {CONFIG_PATH}")
    except Exception as exc:  # Keep the UI available for clear setup errors.
        SETUP_ERROR = str(exc)
        traceback.print_exc()


_prepare_runtime()

SAM3: tuple[Any, Any] | None = None
SAM3D: Any | None = None


def _load_sam3() -> tuple[Any, Any]:
    global SAM3
    if SAM3 is None:
        if not HF_TOKEN:
            raise RuntimeError("Add a read token as the HF_TOKEN Space secret.")
        from transformers import Sam3Model, Sam3Processor

        print("Loading gated SAM 3 model...")
        model = Sam3Model.from_pretrained(str(SAM3_CACHE_DIR)).to("cuda").eval()
        processor = Sam3Processor.from_pretrained(
            str(SAM3_CACHE_DIR),
        )
        SAM3 = (model, processor)
    return SAM3


SAM3_TRACKER: tuple[Any, Any] | None = None


def _load_sam3_tracker() -> tuple[Any, Any]:
    """SAM 3's interactive (click/box) model, shipped in the same facebook/sam3 repo."""
    global SAM3_TRACKER
    if SAM3_TRACKER is None:
        if not HF_TOKEN:
            raise RuntimeError("Add a read token as the HF_TOKEN Space secret.")
        from transformers import Sam3TrackerModel, Sam3TrackerProcessor

        print("Loading SAM 3 point-prompt (tracker) model...")
        model = Sam3TrackerModel.from_pretrained(str(SAM3_CACHE_DIR)).to("cuda").eval()
        processor = Sam3TrackerProcessor.from_pretrained(str(SAM3_CACHE_DIR))
        SAM3_TRACKER = (model, processor)
    return SAM3_TRACKER


def _load_sam3d() -> Any:
    global SAM3D
    if SAM3D is None:
        if CONFIG_PATH is None:
            raise RuntimeError(SETUP_ERROR or "SAM 3D checkpoints are unavailable.")
        from sam3d_inference import SAM3DInference

        print("Loading SAM 3D Objects model...")
        SAM3D = SAM3DInference(str(CONFIG_PATH))
    return SAM3D


# Models are loaded lazily inside the @spaces.GPU handlers: ZeroGPU only exposes a
# real GPU there, and loading at startup made ZeroGPU pack/ship ~17 GB per call.


def _segment(
    image: Image.Image, prompt: str, request_id: str
) -> tuple[np.ndarray | None, int, float | None]:
    import torch

    model, processor = _load_sam3()
    _request_log(
        request_id,
        "SAM 3 inference started",
        f"prompt={prompt!r}; image={image.width}x{image.height}",
    )
    inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)
    with torch.inference_mode():
        output = model(**inputs)
    result = processor.post_process_instance_segmentation(
        output,
        threshold=0.5,
        mask_threshold=0.5,
        target_sizes=inputs["original_sizes"].tolist(),
    )[0]
    masks = result["masks"]
    if len(masks) == 0:
        _request_log(request_id, "SAM 3 inference completed", "instances=0")
        return None, 0, None
    scores = torch.as_tensor(result["scores"])
    best = int(scores.argmax().item())
    best_score = float(scores[best].item())
    _request_log(
        request_id,
        "SAM 3 inference completed",
        f"instances={len(masks)}; selected_mask={best}; score={best_score:.4f}",
    )
    return masks[best].detach().cpu().numpy().astype(bool), len(masks), best_score


GRANULARITY = ("Best score", "Smallest", "Medium", "Largest")


def _segment_points(
    image: Image.Image,
    points: list[list[float]],
    granularity: str,
    request_id: str,
) -> tuple[np.ndarray | None, int, float | None]:
    """Segment the object under keep clicks, excluding the exclude clicks."""
    import torch

    model, processor = _load_sam3_tracker()
    _request_log(
        request_id,
        "SAM 3 point inference started",
        f"points={[(int(x), int(y), int(l)) for x, y, l in points]}; granularity={granularity}",
    )
    input_points = [[[[float(x), float(y)] for x, y, _ in points]]]
    input_labels = [[[int(label) for _, _, label in points]]]
    inputs = processor(
        images=image,
        input_points=input_points,
        input_labels=input_labels,
        return_tensors="pt",
    ).to(model.device)
    # One click is ambiguous (part vs whole), so keep SAM's 3 candidates; with
    # several clicks the extra clicks already disambiguate, so take one mask.
    with torch.inference_mode():
        output = model(**inputs, multimask_output=len(points) == 1)
    masks = processor.post_process_masks(
        output.pred_masks.cpu(), inputs["original_sizes"]
    )[0][0]
    scores = output.iou_scores.detach().cpu().reshape(-1)
    count = masks.shape[0]
    areas = masks.flatten(1).sum(1)
    by_area = [int(i) for i in areas.argsort()]
    if count == 1 or granularity == GRANULARITY[0]:
        pick = int(scores.argmax().item())
    elif granularity == GRANULARITY[1]:
        pick = by_area[0]
    elif granularity == GRANULARITY[3]:
        pick = by_area[-1]
    else:
        pick = by_area[count // 2]
    mask = masks[pick].numpy().astype(bool)
    _request_log(
        request_id,
        "SAM 3 point inference completed",
        f"candidates={count}; selected={pick}; score={float(scores[pick]):.4f}; area_px={int(mask.sum())}",
    )
    if not mask.any():
        return None, 0, None
    return mask, count, float(scores[pick])


def _pick_mask(
    image: Image.Image,
    prompt: str,
    points: list[list[float]],
    granularity: str,
    request_id: str,
) -> tuple[np.ndarray | None, int, float | None]:
    """Clicks take over from the text prompt whenever at least one exists."""
    if points:
        return _segment_points(image, points, granularity, request_id)
    return _segment(image, prompt, request_id)


def _draw_points(image: Image.Image, points: list[list[float]]) -> Image.Image:
    """Return a copy with green (keep) and red (exclude) click markers."""
    from PIL import ImageDraw

    out = image.convert("RGB").copy()
    draw = ImageDraw.Draw(out)
    radius = max(5, round(min(out.size) / 70))
    for x, y, label in points:
        color = (0, 220, 70) if label else (235, 40, 40)
        draw.ellipse(
            (x - radius, y - radius, x + radius, y + radius),
            fill=color,
            outline=(255, 255, 255),
            width=2,
        )
    return out


def _overlay(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
    preview = image.copy()
    green = np.array([0, 220, 100], dtype=np.float32)
    preview[mask] = (preview[mask].astype(np.float32) * 0.45 + green * 0.55).astype(
        np.uint8
    )
    return preview


def _export_glb(result: dict[str, Any]) -> str:
    mesh = result.get("glb")
    if mesh is None:
        raise RuntimeError("SAM 3D finished without producing a GLB mesh.")
    output_dir = Path(tempfile.mkdtemp(prefix="sam3d-output-"))
    output_path = output_dir / "sam3d-object.glb"
    mesh.export(str(output_path))
    if not output_path.is_file() or output_path.stat().st_size == 0:
        raise RuntimeError("GLB export produced an empty file.")
    return str(output_path)


def _request_log(request_id: str, event: str, detail: str = "") -> None:
    timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    suffix = f" | {detail}" if detail else ""
    print(f"[{timestamp}] [reconstruct:{request_id}] {event}{suffix}", flush=True)


@spaces.GPU(duration=30)
def gpu_check() -> str:
    """Cheap check (seconds of quota) that a ZeroGPU worker starts and sees CUDA."""
    import torch

    _request_log("gpucheck", "handler entered")
    if not torch.cuda.is_available():
        return "CUDA is NOT available inside the ZeroGPU worker."
    x = torch.randn(2048, 2048, device="cuda")
    torch.cuda.synchronize()
    name = torch.cuda.get_device_name(0)
    _request_log("gpucheck", "ok", f"{name}; sum={float((x @ x).sum()):.1f}")
    return f"ZeroGPU worker OK: {name}"


@spaces.GPU(duration=60)
def sam3_only(
    image: Image.Image | None,
    prompt: str,
    points: list[list[float]],
    original: Image.Image | None,
    granularity: str,
) -> tuple[Any, str]:
    """Run only SAM 3 so the mask can be verified before spending SAM 3D time."""
    image = original if original is not None else image
    if image is None:
        return None, "Upload an image first."
    prompt = (prompt or "").strip()
    if not points and not prompt:
        return None, "Enter an object prompt, or click the object in the image."
    if SETUP_ERROR:
        return None, f"Space setup is incomplete: {SETUP_ERROR}"
    request_id = uuid.uuid4().hex[:8]
    _request_log(request_id, "SAM 3 test handler entered")
    started = time.monotonic()
    try:
        rgb = image.convert("RGB")
        mask, count, score = _pick_mask(rgb, prompt, points, granularity, request_id)
        if mask is None:
            return None, "SAM 3 ran but found no matching object."
        preview = _overlay(np.asarray(rgb, dtype=np.uint8), mask)
        if points:
            preview = np.asarray(_draw_points(Image.fromarray(preview), points))
        return preview, (
            f"SAM 3 OK ({'clicks' if points else 'text'}): {count} candidate(s), "
            f"score {score:.3f}, "
            f"{time.monotonic() - started:.1f}s (includes model load)."
        )
    except Exception as exc:
        traceback.print_exc()
        return None, f"SAM 3 failed: {type(exc).__name__}: {exc}"


@spaces.GPU(duration=100)
def reconstruct(
    image: Image.Image | None,
    prompt: str,
    points: list[list[float]],
    original: Image.Image | None,
    granularity: str,
) -> tuple[Any, Any, Any, str]:
    """Gradio's direct ZeroGPU event handler; keep it directly decorated."""
    image = original if original is not None else image
    if image is None:
        return None, None, None, "Upload an image first."
    prompt = (prompt or "").strip()
    if not points and not prompt:
        return (
            None,
            None,
            None,
            "Enter a short object description such as 'chair', or click the object.",
        )
    if SETUP_ERROR:
        return None, None, None, f"Space setup is incomplete: {SETUP_ERROR}"

    request_id = uuid.uuid4().hex[:8]
    _request_log(
        request_id,
        "GPU handler entered",
        "maximum_gpu_seconds=150",
    )
    started = time.monotonic()
    rgb = image.convert("RGB")
    image_array = np.asarray(rgb, dtype=np.uint8)
    stage = "SAM 3 segmentation"
    _request_log(
        request_id,
        "GPU work started",
        f"image={rgb.width}x{rgb.height}; max_gpu_seconds=150",
    )
    try:
        _request_log(request_id, "stage started", stage)
        stage_started = time.monotonic()
        mask, count, score = _pick_mask(rgb, prompt, points, granularity, request_id)
        _request_log(
            request_id,
            "SAM 3 stage completed",
            f"matching_instances={count}; selected_score={score}; stage_seconds={time.monotonic() - stage_started:.1f}; total_seconds={time.monotonic() - started:.1f}",
        )
        if mask is None:
            message = (
                "SAM 3 found no object at those clicks."
                if points
                else f"SAM 3 did not find '{prompt}' in this image."
            )
            _request_log(request_id, "request completed without a match", message)
            return None, None, image_array, message

        preview = _overlay(image_array, mask)
        if points:
            preview = np.asarray(_draw_points(Image.fromarray(preview), points))
        stage = "SAM 3D reconstruction"
        _request_log(request_id, "stage started", stage)
        stage_started = time.monotonic()
        result = _load_sam3d()(
            image=image_array,
            mask=mask,
            seed=42,
        )
        _request_log(
            request_id,
            "stage completed",
            f"{stage}; stage_seconds={time.monotonic() - stage_started:.1f}; total_seconds={time.monotonic() - started:.1f}",
        )
        stage = "GLB export"
        _request_log(request_id, "stage started", stage)
        stage_started = time.monotonic()
        model_file = _export_glb(result)
        _request_log(
            request_id,
            "stage completed",
            f"{stage}; stage_seconds={time.monotonic() - stage_started:.1f}; total_seconds={time.monotonic() - started:.1f}",
        )
        elapsed = int(time.monotonic() - started)
        status_text = f"Complete: {count} candidate mask(s); GLB ready in {elapsed}s."
        _request_log(request_id, "request completed", f"elapsed={elapsed}s")
        return model_file, model_file, preview, status_text
    except Exception as exc:
        elapsed = int(time.monotonic() - started)
        _request_log(
            request_id,
            "request failed",
            f"stage={stage}; elapsed={elapsed}s; error={type(exc).__name__}: {exc}",
        )
        traceback.print_exc()
        error_lower = str(exc).lower()
        timeout_hint = " ZeroGPU may have reached its runtime limit; check the run logs." if (
            isinstance(exc, TimeoutError)
            or any(word in error_lower for word in ("timeout", "timed out", "deadline", "time limit"))
        ) else ""
        return (
            None,
            None,
            image_array,
            f"Failed during {stage} after {elapsed}s: {type(exc).__name__}: {exc}.{timeout_hint}",
        )


def _readiness() -> str:
    if SETUP_ERROR:
        return f"Setup: {SETUP_ERROR}"
    if CONFIG_PATH is None:
        return "Setup: SAM 3D checkpoint is unavailable."
    return "Checkpoints downloaded. Models load inside the ZeroGPU call."


def _on_upload(image: Image.Image | None):
    """A new upload resets the click state and re-enables the text prompt."""
    return image, [], gr.update(interactive=True), _point_note([])


def _on_clear():
    return None, [], gr.update(interactive=True), _point_note([])


def _point_note(points: list[list[float]]) -> str:
    if not points:
        return "No clicks: the object prompt is used."
    keep = sum(1 for _, _, label in points if label)
    return (
        f"{keep} keep (green) / {len(points) - keep} exclude (red) click(s). "
        "The text prompt is ignored while clicks exist."
    )


def _redraw(image, original, points):
    base = original if original is not None else image
    shown = _draw_points(base, points) if (base is not None and points) else base
    return (
        shown,
        base,
        points,
        gr.update(interactive=not points),
        _point_note(points),
    )


def _add_point(image, original, points, kind, evt: gr.SelectData):
    if image is None and original is None:
        return _redraw(image, original, points)
    x, y = evt.index[0], evt.index[1]
    points = [*points, [float(x), float(y), 1 if kind.startswith("Keep") else 0]]
    return _redraw(image, original, points)


def _undo_point(image, original, points):
    return _redraw(image, original, points[:-1])


def _clear_points(image, original, points):
    return _redraw(image, original, [])


with gr.Blocks(title="SAM 3D Lab") as demo:
    gr.Markdown(
        """
# SAM 3D Lab
Upload a photo and pick one object, then reconstruct it with **SAM 3** and
**SAM 3D Objects**. Pick it either by **typing a prompt** or by **clicking it**
(green = keep this, red = exclude this, so you can isolate a single part).
Checkpoints download when the Space starts; models load inside each request,
which asks for up to ~150 seconds of ZeroGPU time (typically ~75-130 s used).
"""
    )
    readiness = gr.Markdown(_readiness())
    original_state = gr.State(None)
    points_state = gr.State([])
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input image (click to add points)", type="pil")
            point_note = gr.Markdown(_point_note([]))
            click_kind = gr.Radio(
                ["Keep (green +)", "Exclude (red -)"],
                value="Keep (green +)",
                label="Next click adds",
            )
            with gr.Row():
                undo_button = gr.Button("Undo last click")
                clear_button = gr.Button("Clear clicks")
            granularity = gr.Radio(
                list(GRANULARITY),
                value=GRANULARITY[0],
                label="Single-click mask size (part vs whole)",
            )
            prompt = gr.Textbox(
                label="Object prompt (used only when there are no clicks)",
                placeholder="chair, coffee mug, shoe...",
                value="chair",
            )
            with gr.Row():
                check_button = gr.Button("1. GPU check (cheap)")
                sam_button = gr.Button("2. SAM 3 only")
            run_button = gr.Button("3. Segment and reconstruct", variant="primary")
        with gr.Column():
            preview = gr.Image(label="SAM 3 mask preview", interactive=False)
            status = gr.Textbox(label="Status", lines=3, interactive=False)
    with gr.Row():
        model = gr.Model3D(label="3D preview", height=460)
        download = gr.File(label="Download GLB")

    point_outputs = [input_image, original_state, points_state, prompt, point_note]
    input_image.upload(
        _on_upload,
        inputs=[input_image],
        outputs=[original_state, points_state, prompt, point_note],
    )
    input_image.clear(
        _on_clear,
        inputs=None,
        outputs=[original_state, points_state, prompt, point_note],
    )
    input_image.select(
        _add_point,
        inputs=[input_image, original_state, points_state, click_kind],
        outputs=point_outputs,
    )
    undo_button.click(
        _undo_point,
        inputs=[input_image, original_state, points_state],
        outputs=point_outputs,
    )
    clear_button.click(
        _clear_points,
        inputs=[input_image, original_state, points_state],
        outputs=point_outputs,
    )

    job_inputs = [input_image, prompt, points_state, original_state, granularity]
    check_button.click(gpu_check, inputs=None, outputs=[status])
    sam_button.click(sam3_only, inputs=job_inputs, outputs=[preview, status])
    run_button.click(
        reconstruct,
        inputs=job_inputs,
        outputs=[model, download, preview, status],
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch()

	"""Prompt-guided SAM 3 → SAM 3D Objects demo for Hugging Face ZeroGPU."""

	from __future__ import annotations

	import os

	# Configure backends before importing the model packages. ZeroGPU requires this
	# import to precede torch so it can install its CUDA compatibility layer.
	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"

	import spaces

	import subprocess
	import sys
	import tempfile
	import time
	import traceback
	import uuid
	from pathlib import Path
	from typing import Any

	import gradio as gr
	import numpy as np
	from huggingface_hub import snapshot_download
	from PIL import Image


	ROOT = Path(__file__).resolve().parent
	SAM3D_REVISION = "f91db41"
	SAM3D_CHECKPOINT_ID = "facebook/sam-3d-objects"
	SAM3_CHECKPOINT_ID = "facebook/sam3"
	HF_TOKEN = os.environ.get("HF_TOKEN")

	# The official SAM 3D repo imports these optional libraries at module import
	# time, although the inference path below disables their post-processing.
	for _shim in ("kaolin_stub", "pytorch3d_stub"):
	sys.path.insert(0, str(ROOT / "compat" / _shim))


	def _run(
	command: list[str],
	*,
	timeout: int = 1200,
	env: dict[str, str] \| None = None,
	) -> None:
	"""Run setup work outside ZeroGPU allocation and show useful errors."""
	result = subprocess.run(
	command,
	check=False,
	capture_output=True,
	text=True,
	timeout=timeout,
	env=env,
	)
	if result.returncode:
	tail = (result.stderr or result.stdout)[-5000:]
	raise RuntimeError(f"Command failed ({result.returncode}): {tail}")


	def _install_runtime_packages() -> None:
	"""Install only packages whose setup needs the platform-provided torch."""
	marker = Path(tempfile.gettempdir()) / ".sam3d-runtime-deps-ready"
	if marker.exists():
	return

	requirements = ROOT / "runtime-requirements.txt"
	if not requirements.exists():
	raise RuntimeError("runtime-requirements.txt is missing")
	print("Installing SAM 3D packages that require ZeroGPU's torch runtime...")
	_run(
	[
	sys.executable,
	"-m",
	"pip",
	"install",
	"--no-cache-dir",
	"--no-deps",
	"-r",
	str(requirements),
	],
	timeout=1800,
	)
	# gsplat's setup.py imports torch while preparing its wheel. PEP 517's
	# isolated build environment hides the platform's installed torch, so
	# disable isolation and use gsplat's documented no-CUDA-build/JIT mode.
	gsplat_env = os.environ.copy()
	gsplat_env["BUILD_NO_CUDA"] = "1"
	_run(
	[
	sys.executable,
	"-m",
	"pip",
	"install",
	"--no-cache-dir",
	"--no-deps",
	"--no-build-isolation",
	"git+https://github.com/nerfstudio-project/gsplat.git@2323de5905d5e90e035f792fe65bad0fedd413e7",
	],
	timeout=1800,
	env=gsplat_env,
	)
	marker.write_text("ready\n", encoding="utf-8")


	SETUP_ERROR: str \| None = None
	CONFIG_PATH: Path \| None = None
	SAM3_CACHE_DIR: Path \| None = None


	def _prepare_runtime() -> None:
	global SETUP_ERROR, CONFIG_PATH, SAM3_CACHE_DIR
	try:
	_install_runtime_packages()
	from sam3d_objects import __file__ as _sam3d_package_file # noqa: F401

	if not HF_TOKEN:
	SETUP_ERROR = "Add a read token as the HF_TOKEN Space secret."
	return

	# Download once at Space startup. Checkpoint transfer is CPU/network
	# work and therefore does not occupy a ZeroGPU allocation.
	SAM3_CACHE_DIR = Path(
	snapshot_download(
	repo_id=SAM3_CHECKPOINT_ID,
	token=HF_TOKEN,
	)
	)
	cache_dir = Path(
	snapshot_download(
	repo_id=SAM3D_CHECKPOINT_ID,
	token=HF_TOKEN,
	)
	)
	checkpoint_dir = cache_dir / "checkpoints"
	config = checkpoint_dir / "pipeline.yaml"
	if not config.is_file():
	raise RuntimeError(
	f"Expected pipeline config at {config}; verify model access and files."
	)
	# Keep the config in its original Hub snapshot directory. snapshot_download
	# represents large checkpoint files with relative symlinks to the cache's
	# blobs directory; symlinking only this subdirectory elsewhere breaks those
	# relative links and causes FileNotFoundError while loading weights.
	CONFIG_PATH = config
	from omegaconf import OmegaConf

	checkpoint_config = OmegaConf.load(CONFIG_PATH)
	for key in (
	"ss_generator_config_path",
	"ss_generator_ckpt_path",
	"slat_generator_config_path",
	"slat_generator_ckpt_path",
	"ss_decoder_config_path",
	"ss_decoder_ckpt_path",
	"slat_decoder_gs_config_path",
	"slat_decoder_gs_ckpt_path",
	"slat_decoder_gs_4_config_path",
	"slat_decoder_gs_4_ckpt_path",
	"slat_decoder_mesh_config_path",
	"slat_decoder_mesh_ckpt_path",
	):
	relative_path = checkpoint_config.get(key)
	if relative_path:
	checkpoint_file = CONFIG_PATH.parent / str(relative_path)
	if not checkpoint_file.is_file():
	raise RuntimeError(
	f"SAM 3D checkpoint is missing or unreadable: "
	f"{checkpoint_file} (from {key})."
	)
	print(f"Using SAM 3D pipeline config: {CONFIG_PATH}")
	except Exception as exc: # Keep the UI available for clear setup errors.
	SETUP_ERROR = str(exc)
	traceback.print_exc()


	_prepare_runtime()

	SAM3: tuple[Any, Any] \| None = None
	SAM3D: Any \| None = None


	def _load_sam3() -> tuple[Any, Any]:
	global SAM3
	if SAM3 is None:
	if not HF_TOKEN:
	raise RuntimeError("Add a read token as the HF_TOKEN Space secret.")
	from transformers import Sam3Model, Sam3Processor

	print("Loading gated SAM 3 model...")
	model = Sam3Model.from_pretrained(str(SAM3_CACHE_DIR)).to("cuda").eval()
	processor = Sam3Processor.from_pretrained(
	str(SAM3_CACHE_DIR),
	)
	SAM3 = (model, processor)
	return SAM3


	SAM3_TRACKER: tuple[Any, Any] \| None = None


	def _load_sam3_tracker() -> tuple[Any, Any]:
	"""SAM 3's interactive (click/box) model, shipped in the same facebook/sam3 repo."""
	global SAM3_TRACKER
	if SAM3_TRACKER is None:
	if not HF_TOKEN:
	raise RuntimeError("Add a read token as the HF_TOKEN Space secret.")
	from transformers import Sam3TrackerModel, Sam3TrackerProcessor

	print("Loading SAM 3 point-prompt (tracker) model...")
	model = Sam3TrackerModel.from_pretrained(str(SAM3_CACHE_DIR)).to("cuda").eval()
	processor = Sam3TrackerProcessor.from_pretrained(str(SAM3_CACHE_DIR))
	SAM3_TRACKER = (model, processor)
	return SAM3_TRACKER


	def _load_sam3d() -> Any:
	global SAM3D
	if SAM3D is None:
	if CONFIG_PATH is None:
	raise RuntimeError(SETUP_ERROR or "SAM 3D checkpoints are unavailable.")
	from sam3d_inference import SAM3DInference

	print("Loading SAM 3D Objects model...")
	SAM3D = SAM3DInference(str(CONFIG_PATH))
	return SAM3D


	# Models are loaded lazily inside the @spaces.GPU handlers: ZeroGPU only exposes a
	# real GPU there, and loading at startup made ZeroGPU pack/ship ~17 GB per call.


	def _segment(
	image: Image.Image, prompt: str, request_id: str
	) -> tuple[np.ndarray \| None, int, float \| None]:
	import torch

	model, processor = _load_sam3()
	_request_log(
	request_id,
	"SAM 3 inference started",
	f"prompt={prompt!r}; image={image.width}x{image.height}",
	)
	inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)
	with torch.inference_mode():
	output = model(**inputs)
	result = processor.post_process_instance_segmentation(
	output,
	threshold=0.5,
	mask_threshold=0.5,
	target_sizes=inputs["original_sizes"].tolist(),
	)[0]
	masks = result["masks"]
	if len(masks) == 0:
	_request_log(request_id, "SAM 3 inference completed", "instances=0")
	return None, 0, None
	scores = torch.as_tensor(result["scores"])
	best = int(scores.argmax().item())
	best_score = float(scores[best].item())
	_request_log(
	request_id,
	"SAM 3 inference completed",
	f"instances={len(masks)}; selected_mask={best}; score={best_score:.4f}",
	)
	return masks[best].detach().cpu().numpy().astype(bool), len(masks), best_score


	GRANULARITY = ("Best score", "Smallest", "Medium", "Largest")


	def _segment_points(
	image: Image.Image,
	points: list[list[float]],
	granularity: str,
	request_id: str,
	) -> tuple[np.ndarray \| None, int, float \| None]:
	"""Segment the object under keep clicks, excluding the exclude clicks."""
	import torch

	model, processor = _load_sam3_tracker()
	_request_log(
	request_id,
	"SAM 3 point inference started",
	f"points={[(int(x), int(y), int(l)) for x, y, l in points]}; granularity={granularity}",
	)
	input_points = [[[[float(x), float(y)] for x, y, _ in points]]]
	input_labels = [[[int(label) for _, _, label in points]]]
	inputs = processor(
	images=image,
	input_points=input_points,
	input_labels=input_labels,
	return_tensors="pt",
	).to(model.device)
	# One click is ambiguous (part vs whole), so keep SAM's 3 candidates; with
	# several clicks the extra clicks already disambiguate, so take one mask.
	with torch.inference_mode():
	output = model(**inputs, multimask_output=len(points) == 1)
	masks = processor.post_process_masks(
	output.pred_masks.cpu(), inputs["original_sizes"]
	)[0][0]
	scores = output.iou_scores.detach().cpu().reshape(-1)
	count = masks.shape[0]
	areas = masks.flatten(1).sum(1)
	by_area = [int(i) for i in areas.argsort()]
	if count == 1 or granularity == GRANULARITY[0]:
	pick = int(scores.argmax().item())
	elif granularity == GRANULARITY[1]:
	pick = by_area[0]
	elif granularity == GRANULARITY[3]:
	pick = by_area[-1]
	else:
	pick = by_area[count // 2]
	mask = masks[pick].numpy().astype(bool)
	_request_log(
	request_id,
	"SAM 3 point inference completed",
	f"candidates={count}; selected={pick}; score={float(scores[pick]):.4f}; area_px={int(mask.sum())}",
	)
	if not mask.any():
	return None, 0, None
	return mask, count, float(scores[pick])


	def _pick_mask(
	image: Image.Image,
	prompt: str,
	points: list[list[float]],
	granularity: str,
	request_id: str,
	) -> tuple[np.ndarray \| None, int, float \| None]:
	"""Clicks take over from the text prompt whenever at least one exists."""
	if points:
	return _segment_points(image, points, granularity, request_id)
	return _segment(image, prompt, request_id)


	def _draw_points(image: Image.Image, points: list[list[float]]) -> Image.Image:
	"""Return a copy with green (keep) and red (exclude) click markers."""
	from PIL import ImageDraw

	out = image.convert("RGB").copy()
	draw = ImageDraw.Draw(out)
	radius = max(5, round(min(out.size) / 70))
	for x, y, label in points:
	color = (0, 220, 70) if label else (235, 40, 40)
	draw.ellipse(
	(x - radius, y - radius, x + radius, y + radius),
	fill=color,
	outline=(255, 255, 255),
	width=2,
	)
	return out


	def _overlay(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
	preview = image.copy()
	green = np.array([0, 220, 100], dtype=np.float32)
	preview[mask] = (preview[mask].astype(np.float32) * 0.45 + green * 0.55).astype(
	np.uint8
	)
	return preview


	def _export_glb(result: dict[str, Any]) -> str:
	mesh = result.get("glb")
	if mesh is None:
	raise RuntimeError("SAM 3D finished without producing a GLB mesh.")
	output_dir = Path(tempfile.mkdtemp(prefix="sam3d-output-"))
	output_path = output_dir / "sam3d-object.glb"
	mesh.export(str(output_path))
	if not output_path.is_file() or output_path.stat().st_size == 0:
	raise RuntimeError("GLB export produced an empty file.")
	return str(output_path)


	def _request_log(request_id: str, event: str, detail: str = "") -> None:
	timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
	suffix = f" \| {detail}" if detail else ""
	print(f"[{timestamp}] [reconstruct:{request_id}] {event}{suffix}", flush=True)


	@spaces.GPU(duration=30)
	def gpu_check() -> str:
	"""Cheap check (seconds of quota) that a ZeroGPU worker starts and sees CUDA."""
	import torch

	_request_log("gpucheck", "handler entered")
	if not torch.cuda.is_available():
	return "CUDA is NOT available inside the ZeroGPU worker."
	x = torch.randn(2048, 2048, device="cuda")
	torch.cuda.synchronize()
	name = torch.cuda.get_device_name(0)
	_request_log("gpucheck", "ok", f"{name}; sum={float((x @ x).sum()):.1f}")
	return f"ZeroGPU worker OK: {name}"


	@spaces.GPU(duration=60)
	def sam3_only(
	image: Image.Image \| None,
	prompt: str,
	points: list[list[float]],
	original: Image.Image \| None,
	granularity: str,
	) -> tuple[Any, str]:
	"""Run only SAM 3 so the mask can be verified before spending SAM 3D time."""
	image = original if original is not None else image
	if image is None:
	return None, "Upload an image first."
	prompt = (prompt or "").strip()
	if not points and not prompt:
	return None, "Enter an object prompt, or click the object in the image."
	if SETUP_ERROR:
	return None, f"Space setup is incomplete: {SETUP_ERROR}"
	request_id = uuid.uuid4().hex[:8]
	_request_log(request_id, "SAM 3 test handler entered")
	started = time.monotonic()
	try:
	rgb = image.convert("RGB")
	mask, count, score = _pick_mask(rgb, prompt, points, granularity, request_id)
	if mask is None:
	return None, "SAM 3 ran but found no matching object."
	preview = _overlay(np.asarray(rgb, dtype=np.uint8), mask)
	if points:
	preview = np.asarray(_draw_points(Image.fromarray(preview), points))
	return preview, (
	f"SAM 3 OK ({'clicks' if points else 'text'}): {count} candidate(s), "
	f"score {score:.3f}, "
	f"{time.monotonic() - started:.1f}s (includes model load)."
	)
	except Exception as exc:
	traceback.print_exc()
	return None, f"SAM 3 failed: {type(exc).__name__}: {exc}"


	@spaces.GPU(duration=100)
	def reconstruct(
	image: Image.Image \| None,
	prompt: str,
	points: list[list[float]],
	original: Image.Image \| None,
	granularity: str,
	) -> tuple[Any, Any, Any, str]:
	"""Gradio's direct ZeroGPU event handler; keep it directly decorated."""
	image = original if original is not None else image
	if image is None:
	return None, None, None, "Upload an image first."
	prompt = (prompt or "").strip()
	if not points and not prompt:
	return (
	None,
	None,
	None,
	"Enter a short object description such as 'chair', or click the object.",
	)
	if SETUP_ERROR:
	return None, None, None, f"Space setup is incomplete: {SETUP_ERROR}"

	request_id = uuid.uuid4().hex[:8]
	_request_log(
	request_id,
	"GPU handler entered",
	"maximum_gpu_seconds=150",
	)
	started = time.monotonic()
	rgb = image.convert("RGB")
	image_array = np.asarray(rgb, dtype=np.uint8)
	stage = "SAM 3 segmentation"
	_request_log(
	request_id,
	"GPU work started",
	f"image={rgb.width}x{rgb.height}; max_gpu_seconds=150",
	)
	try:
	_request_log(request_id, "stage started", stage)
	stage_started = time.monotonic()
	mask, count, score = _pick_mask(rgb, prompt, points, granularity, request_id)
	_request_log(
	request_id,
	"SAM 3 stage completed",
	f"matching_instances={count}; selected_score={score}; stage_seconds={time.monotonic() - stage_started:.1f}; total_seconds={time.monotonic() - started:.1f}",
	)
	if mask is None:
	message = (
	"SAM 3 found no object at those clicks."
	if points
	else f"SAM 3 did not find '{prompt}' in this image."
	)
	_request_log(request_id, "request completed without a match", message)
	return None, None, image_array, message

	preview = _overlay(image_array, mask)
	if points:
	preview = np.asarray(_draw_points(Image.fromarray(preview), points))
	stage = "SAM 3D reconstruction"
	_request_log(request_id, "stage started", stage)
	stage_started = time.monotonic()
	result = _load_sam3d()(
	image=image_array,
	mask=mask,
	seed=42,
	)
	_request_log(
	request_id,
	"stage completed",
	f"{stage}; stage_seconds={time.monotonic() - stage_started:.1f}; total_seconds={time.monotonic() - started:.1f}",
	)
	stage = "GLB export"
	_request_log(request_id, "stage started", stage)
	stage_started = time.monotonic()
	model_file = _export_glb(result)
	_request_log(
	request_id,
	"stage completed",
	f"{stage}; stage_seconds={time.monotonic() - stage_started:.1f}; total_seconds={time.monotonic() - started:.1f}",
	)
	elapsed = int(time.monotonic() - started)
	status_text = f"Complete: {count} candidate mask(s); GLB ready in {elapsed}s."
	_request_log(request_id, "request completed", f"elapsed={elapsed}s")
	return model_file, model_file, preview, status_text
	except Exception as exc:
	elapsed = int(time.monotonic() - started)
	_request_log(
	request_id,
	"request failed",
	f"stage={stage}; elapsed={elapsed}s; error={type(exc).__name__}: {exc}",
	)
	traceback.print_exc()
	error_lower = str(exc).lower()
	timeout_hint = " ZeroGPU may have reached its runtime limit; check the run logs." if (
	isinstance(exc, TimeoutError)
	or any(word in error_lower for word in ("timeout", "timed out", "deadline", "time limit"))
	) else ""
	return (
	None,
	None,
	image_array,
	f"Failed during {stage} after {elapsed}s: {type(exc).__name__}: {exc}.{timeout_hint}",
	)


	def _readiness() -> str:
	if SETUP_ERROR:
	return f"Setup: {SETUP_ERROR}"
	if CONFIG_PATH is None:
	return "Setup: SAM 3D checkpoint is unavailable."
	return "Checkpoints downloaded. Models load inside the ZeroGPU call."


	def _on_upload(image: Image.Image \| None):
	"""A new upload resets the click state and re-enables the text prompt."""
	return image, [], gr.update(interactive=True), _point_note([])


	def _on_clear():
	return None, [], gr.update(interactive=True), _point_note([])


	def _point_note(points: list[list[float]]) -> str:
	if not points:
	return "No clicks: the object prompt is used."
	keep = sum(1 for _, _, label in points if label)
	return (
	f"{keep} keep (green) / {len(points) - keep} exclude (red) click(s). "
	"The text prompt is ignored while clicks exist."
	)


	def _redraw(image, original, points):
	base = original if original is not None else image
	shown = _draw_points(base, points) if (base is not None and points) else base
	return (
	shown,
	base,
	points,
	gr.update(interactive=not points),
	_point_note(points),
	)


	def _add_point(image, original, points, kind, evt: gr.SelectData):
	if image is None and original is None:
	return _redraw(image, original, points)
	x, y = evt.index[0], evt.index[1]
	points = [*points, [float(x), float(y), 1 if kind.startswith("Keep") else 0]]
	return _redraw(image, original, points)


	def _undo_point(image, original, points):
	return _redraw(image, original, points[:-1])


	def _clear_points(image, original, points):
	return _redraw(image, original, [])


	with gr.Blocks(title="SAM 3D Lab") as demo:
	gr.Markdown(
	"""
	# SAM 3D Lab

	Upload a photo and pick one object, then reconstruct it with SAM 3 and
	SAM 3D Objects. Pick it either by typing a prompt or by clicking it
	(green = keep this, red = exclude this, so you can isolate a single part).
	Checkpoints download when the Space starts; models load inside each request,
	which asks for up to ~150 seconds of ZeroGPU time (typically ~75-130 s used).
	"""
	)
	readiness = gr.Markdown(_readiness())
	original_state = gr.State(None)
	points_state = gr.State([])
	with gr.Row():
	with gr.Column():
	input_image = gr.Image(label="Input image (click to add points)", type="pil")
	point_note = gr.Markdown(_point_note([]))
	click_kind = gr.Radio(
	["Keep (green +)", "Exclude (red -)"],
	value="Keep (green +)",
	label="Next click adds",
	)
	with gr.Row():
	undo_button = gr.Button("Undo last click")
	clear_button = gr.Button("Clear clicks")
	granularity = gr.Radio(
	list(GRANULARITY),
	value=GRANULARITY[0],
	label="Single-click mask size (part vs whole)",
	)
	prompt = gr.Textbox(
	label="Object prompt (used only when there are no clicks)",
	placeholder="chair, coffee mug, shoe...",
	value="chair",
	)
	with gr.Row():
	check_button = gr.Button("1. GPU check (cheap)")
	sam_button = gr.Button("2. SAM 3 only")
	run_button = gr.Button("3. Segment and reconstruct", variant="primary")
	with gr.Column():
	preview = gr.Image(label="SAM 3 mask preview", interactive=False)
	status = gr.Textbox(label="Status", lines=3, interactive=False)
	with gr.Row():
	model = gr.Model3D(label="3D preview", height=460)
	download = gr.File(label="Download GLB")

	point_outputs = [input_image, original_state, points_state, prompt, point_note]
	input_image.upload(
	_on_upload,
	inputs=[input_image],
	outputs=[original_state, points_state, prompt, point_note],
	)
	input_image.clear(
	_on_clear,
	inputs=None,
	outputs=[original_state, points_state, prompt, point_note],
	)
	input_image.select(
	_add_point,
	inputs=[input_image, original_state, points_state, click_kind],
	outputs=point_outputs,
	)
	undo_button.click(
	_undo_point,
	inputs=[input_image, original_state, points_state],
	outputs=point_outputs,
	)
	clear_button.click(
	_clear_points,
	inputs=[input_image, original_state, points_state],
	outputs=point_outputs,
	)

	job_inputs = [input_image, prompt, points_state, original_state, granularity]
	check_button.click(gpu_check, inputs=None, outputs=[status])
	sam_button.click(sam3_only, inputs=job_inputs, outputs=[preview, status])
	run_button.click(
	reconstruct,
	inputs=job_inputs,
	outputs=[model, download, preview, status],
	)


	if __name__ == "__main__":
	demo.queue(default_concurrency_limit=1).launch()