from __future__ import annotations import base64 import gc import io import json import os import shutil import subprocess import sys import tempfile import time import zlib from pathlib import Path from threading import Lock from uuid import uuid4 import numpy as np from fastapi.testclient import TestClient from PIL import Image, ImageOps PROJECT_REPOSITORY = "https://github.com/DrStrangel0ve/3dprintpic.git" PROJECT_REVISION = "481aedec2788843b52be374b6a2c85df9571864d" TRIPOSG_REPOSITORY = "https://github.com/VAST-AI-Research/TripoSG.git" TRIPOSG_SOURCE_REVISION = "fc5c40990181e2a756c4e0b1c2f4d6b5202faf8c" TRIPOSG_MODEL_REVISION = "2c1c516d22d58db486a058d98d31bb6177344e06" DEPTH_MODEL = "depth-anything/Depth-Anything-V2-Large-hf" DEPTH_MODEL_REVISION = "7581137eff8d4e94f6e796d3baea0e9fa79b22d2" SAM3_MODEL = "facebook/sam3" SAM3_MODEL_REVISION = "3c879f39826c281e95690f02c7821c4de09afae7" SPACE_DIR = Path(__file__).resolve().parent RUNTIME_ROOT = Path(os.getenv("THREEDPRINTPIC_RUNTIME_DIR", Path(tempfile.gettempdir()) / "3dprintpic-space")) OUTPUT_DIR = RUNTIME_ROOT / "output" SOURCE_DIR = RUNTIME_ROOT / "source" TRIPOSG_DIR = RUNTIME_ROOT / "TripoSG" ASSET_DIR = RUNTIME_ROOT / "assets" _SOURCE_LOCK = Lock() _TRIPOSG_LOCK = Lock() _FACE_ASSET_LOCK = Lock() _FACE_ASSET_STATUS = None SAM3_HOVER_MAP_MAX_DIMENSION = 1024 SAM3_HOVER_MAP_VERSION = 1 SAM3_PRECOMPUTE_STATE_VERSION = 2 SAM3_PRECOMPUTE_TTL_SECONDS = 20 * 60 SAM3_PRECOMPUTE_MAX_SOURCE_PIXELS = 40_000_000 SAM3_PRECOMPUTE_MAX_MASKS = 256 SAM3_PRECOMPUTE_MAX_COMPRESSED_BYTES = 32 * 1024 * 1024 SAM3_HOVER_SPECIFICITY_WEIGHT = 0.20 SAM3_HOVER_CONTEXT_MIN_COVERAGE = 0.25 SAM3_HOVER_CONTEXT_LARGE_COVERAGE = 0.55 SAM3_HOVER_CONTEXT_MIN_BORDER_SIDES = 3 LOCAL_RELIEF_DETAIL_MULTIPLIERS = { 384: 1.5, 512: 2.0, 768: 3.0, 900: 4.0, } LOCAL_RELIEF_PRINTER_EDGE_MM = 256 def _run_git(*args: str, cwd: Path | None = None) -> str: completed = subprocess.run( ["git", *args], cwd=cwd, check=True, capture_output=True, text=True, timeout=300, ) return completed.stdout.strip() def _repository_root_if_local() -> Path | None: explicit = os.getenv("THREEDPRINTPIC_SOURCE_DIR") candidates = [Path(explicit).expanduser()] if explicit else [] candidates.extend((SPACE_DIR.parent, Path.cwd())) for candidate in candidates: resolved = candidate.resolve() if (resolved / "backend" / "main.py").is_file(): return resolved return None def _clone_exact_repository(repository: str, revision: str, destination: Path) -> Path: with _SOURCE_LOCK: if (destination / ".git").is_dir(): current = _run_git("rev-parse", "HEAD", cwd=destination) if current == revision: return destination shutil.rmtree(destination) destination.parent.mkdir(parents=True, exist_ok=True) _run_git("clone", "--filter=blob:none", "--no-checkout", repository, str(destination)) _run_git("checkout", "--detach", revision, cwd=destination) current = _run_git("rev-parse", "HEAD", cwd=destination) if current != revision: raise RuntimeError(f"Source checkout mismatch: expected {revision}, got {current}") return destination def ensure_project_source() -> Path: local = _repository_root_if_local() if local is not None: return local return _clone_exact_repository(PROJECT_REPOSITORY, PROJECT_REVISION, SOURCE_DIR) PROJECT_DIR = ensure_project_source() if str(PROJECT_DIR) not in sys.path: sys.path.insert(0, str(PROJECT_DIR)) OUTPUT_DIR.mkdir(parents=True, exist_ok=True) ASSET_DIR.mkdir(parents=True, exist_ok=True) os.environ.setdefault("OUTPUT_DIR", str(OUTPUT_DIR)) os.environ.setdefault("DEPTH_PROVIDER", "transformers") os.environ.setdefault("DEPTH_MODEL", DEPTH_MODEL) os.environ.setdefault("SELECTION_ALLOW_MODEL_DOWNLOAD", "1") os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") os.environ.setdefault("THREEDPRINTPIC_ASSET_CACHE_DIR", str(ASSET_DIR)) from backend import main as backend_main # noqa: E402 from backend import face_depth_refinement as backend_face_refinement # noqa: E402 from backend import gnm_face_foundation as backend_gnm_face # noqa: E402 from backend.benchmark.direct_mesh import ( # noqa: E402 convert_mesh_to_stl, postprocess_mesh_for_stl, repair_mesh_for_printable_stl, ) from backend.stl_diagnostics import json_safe_stl_diagnostics, stl_diagnostics # noqa: E402 BACKEND_CLIENT = TestClient(backend_main.app) def image_dimensions_mm(image_path: str | Path | None, long_edge_mm: float) -> tuple[float, float]: if not image_path: return float(long_edge_mm), float(long_edge_mm) with Image.open(image_path) as image: width, height = ImageOps.exif_transpose(image).size scale = float(long_edge_mm) / float(max(width, height, 1)) return round(width * scale, 2), round(height * scale, 2) def local_relief_dimensions_mm( image_path: str | Path | None, print_scale_percent: float, ) -> tuple[float, float]: scale_percent = float(print_scale_percent) if scale_percent < 10.0 or scale_percent > 100.0 or scale_percent % 5.0 != 0.0: raise ValueError("Print size must match the local 10-100% production scale") long_edge_mm = int(LOCAL_RELIEF_PRINTER_EDGE_MM * scale_percent / 100.0) return image_dimensions_mm(image_path, long_edge_mm) def _safe_file(path: Path) -> str: resolved = path.resolve() if not resolved.is_file(): raise FileNotFoundError(f"Expected output was not created: {resolved.name}") return str(resolved) def _response_error(response) -> RuntimeError: try: payload = response.json() except Exception: payload = response.text if isinstance(payload, dict): detail = payload.get("detail") or payload else: detail = payload return RuntimeError(f"Generation failed ({response.status_code}): {detail}") def _diagnostic_summary(diagnostics: dict, *, model: str) -> dict: aliases = { "is_watertight": ("stl_is_watertight", "is_watertight"), "is_volume": ("stl_is_volume", "is_volume"), "is_manifold": ("stl_is_manifold", "is_manifold"), "winding_consistent": ("stl_winding_consistent", "winding_consistent"), "component_count": ("stl_component_count", "component_count"), "nonmanifold_edge_count": ("stl_nonmanifold_edge_count", "nonmanifold_edge_count"), "degenerate_face_count": ("stl_degenerate_face_count", "degenerate_face_count"), "positive_volume": ("stl_positive_volume", "positive_volume"), "face_count": ("stl_faces", "face_count"), "vertex_count": ("stl_vertices", "vertex_count"), "normalized_bbox_complexity_log1p": ( "stl_faces_per_normalized_bbox_volume_log1p", "normalized_bbox_complexity_log1p", ), } summary = {"model": model} for output_key, candidates in aliases.items(): for candidate in candidates: if candidate in diagnostics: summary[output_key] = diagnostics[candidate] break bbox = [diagnostics.get(f"stl_bbox_{axis}") for axis in "xyz"] if all(value is not None for value in bbox): summary["bbox_extents"] = bbox hard_checks = { "watertight": summary.get("is_watertight") is True, "volume": summary.get("is_volume") is True, "manifold": summary.get("is_manifold") is True, "winding": summary.get("winding_consistent") is True, "single_component": summary.get("component_count") == 1, "nonmanifold_edges": summary.get("nonmanifold_edge_count") == 0, "degenerate_faces": summary.get("degenerate_face_count") == 0, "positive_volume": summary.get("positive_volume") is True, } summary["stl_passes_hard_checks"] = all(hard_checks.values()) summary["stl_failed_checks"] = [name for name, passed in hard_checks.items() if not passed] return summary def _require_face_native_runtime(platform_name: str | None = None) -> dict[str, str]: platform_name = os.name if platform_name is None else platform_name if platform_name == "nt": return {} import ctypes import ctypes.util required = ( ("glesv2", "GLESv2", "libGLESv2.so.2", "libgles2"), ("egl", "EGL", "libEGL.so.1", "libegl1"), ) libraries = {} for key, lookup_name, soname, package in required: library = ctypes.util.find_library(lookup_name) or soname try: ctypes.CDLL(library) except OSError as exc: raise RuntimeError( f"Face parity native runtime preflight failed: {soname} is " f"unavailable; install the Debian {package} package" ) from exc libraries[key] = library return libraries def ensure_face_assets() -> dict: global _FACE_ASSET_STATUS if _FACE_ASSET_STATUS is not None: return json.loads(json.dumps(_FACE_ASSET_STATUS)) with _FACE_ASSET_LOCK: if _FACE_ASSET_STATUS is not None: return json.loads(json.dumps(_FACE_ASSET_STATUS)) try: native_libraries = _require_face_native_runtime() resolved = { "mediapipe_face_landmarker": ( backend_face_refinement._resolve_face_landmarker_model(), backend_face_refinement.FACE_LANDMARKER_MODEL_SHA256, ), "yunet_face_detector": ( backend_face_refinement._resolve_yunet_model(), backend_face_refinement.YUNET_MODEL_SHA256, ), "gnm_head_model": ( backend_gnm_face.resolve_gnm_model(), backend_gnm_face.GNM_MODEL_SHA256, ), "gnm_head_landmarks": ( backend_gnm_face.resolve_gnm_landmarks(), backend_gnm_face.GNM_LANDMARKS_SHA256, ), } except Exception as exc: raise RuntimeError( f"Face parity asset preflight failed: {type(exc).__name__}: {exc}" ) from exc _FACE_ASSET_STATUS = { "verified": True, "cache": "writable-runtime-assets", "native_libraries": native_libraries, "assets": { name: { "filename": path.name, "bytes": int(path.stat().st_size), "sha256": sha256, } for name, (path, sha256) in resolved.items() }, } return json.loads(json.dumps(_FACE_ASSET_STATUS)) def _selection_requests_person(selection: dict | None) -> bool: if not isinstance(selection, dict): return False labels = selection.get("labels") if not isinstance(labels, (list, tuple)): return False return any(str(label).strip().casefold() == "person" for label in labels) def _require_deterministic_depth_parity(result: dict) -> dict: depth_metadata = result.get("depth_metadata") if not isinstance(depth_metadata, dict): depth_metadata = {} summary = { "requested_precision": result.get("requested_depth_inference_precision"), "effective_precision": result.get("depth_inference_precision"), "deterministic_cuda": depth_metadata.get("deterministic_cuda"), } failures = [] if summary["requested_precision"] != "float32": failures.append("requested precision was not float32") if summary["effective_precision"] != "float32": failures.append("effective precision was not float32") if summary["deterministic_cuda"] is not True: failures.append("deterministic CUDA was not active") if failures: raise RuntimeError( "Hosted depth parity check failed: " + "; ".join(failures) ) summary["status"] = "verified" return summary def _require_face_parity(result: dict, selection: dict | None) -> dict: selected_person_required = _selection_requests_person(selection) face_refinement = result.get("face_refinement") if not isinstance(face_refinement, dict): face_refinement = {} faces = face_refinement.get("faces") faces = list(faces) if isinstance(faces, list) else [] refined_faces = [ face for face in faces if isinstance(face, dict) and face.get("status") == "refined" ] mediapipe_faces = [ face for face in refined_faces if "mediapipe-face-landmarker" in str(face.get("detector") or "") and int(face.get("landmark_count") or 0) >= 468 ] deterministic_depth_faces = [ face for face in refined_faces if isinstance(face.get("depth_inference"), dict) and face["depth_inference"].get("requested_precision") == "float32" and face["depth_inference"].get("effective_precision") == "float32" and face["depth_inference"].get("deterministic_cuda") is True ] detector_errors = face_refinement.get("detector_errors") detector_errors = ( list(detector_errors) if isinstance(detector_errors, list) else [] ) summary = { "required_for_selected_person": selected_person_required, "applied": face_refinement.get("applied") is True, "detected_faces": int(face_refinement.get("detected_faces") or 0), "refined_faces": int(face_refinement.get("refined_faces") or 0), "mediapipe_landmark_faces": len(mediapipe_faces), "deterministic_fp32_depth_faces": len(deterministic_depth_faces), "landmark_shape_prior_faces": sum( 1 for face in refined_faces if isinstance(face.get("landmark_shape_prior"), dict) and face["landmark_shape_prior"].get("enabled") is True ), "detectors": sorted( { str(face.get("detector")) for face in refined_faces if face.get("detector") } ), "detector_errors": detector_errors, } face_parity_required = ( selected_person_required or summary["detected_faces"] > 0 ) summary["required"] = face_parity_required if not face_parity_required: summary["status"] = "observed" return summary failures = [] if not summary["applied"]: failures.append("face refinement was not applied") if summary["detected_faces"] < 1: failures.append("no selected face was detected") if summary["refined_faces"] != summary["detected_faces"]: failures.append("not every detected face was refined") if len(refined_faces) != summary["refined_faces"]: failures.append("refined-face telemetry was incomplete") if summary["mediapipe_landmark_faces"] != summary["refined_faces"]: failures.append("MediaPipe 468+ landmark guidance was not used for every face") if summary["deterministic_fp32_depth_faces"] != summary["refined_faces"]: failures.append("deterministic FP32 depth was not used for every face crop") if detector_errors: failures.append("a face detector reported an error") if failures: diagnostics = ( "; detector diagnostics: " + " | ".join(detector_errors) if detector_errors else "" ) raise RuntimeError( "Hosted face parity check failed: " + "; ".join(failures) + diagnostics ) summary["status"] = "verified" return summary def _require_sam3_token() -> None: if not os.getenv("HF_TOKEN"): raise RuntimeError( "SAM 3 needs the Space owner's read-only HF_TOKEN secret before object selection can run." ) def _load_selection_image(image_path: str | Path) -> Image.Image: with Image.open(image_path) as source: return ImageOps.exif_transpose(source).convert("RGB") def _save_selection_job( image_path: str | Path, image: Image.Image, mask: Image.Image, *, resolved_model: str, model_status: str, labels: list[str], mask_count: int = 1, ) -> dict: width, height = image.size if resolved_model != SAM3_MODEL or not model_status.startswith("sam3"): raise RuntimeError("SAM 3 did not produce a verified selection mask") job_id = uuid4().hex job_dir = OUTPUT_DIR / "selection" / job_id job_dir.mkdir(parents=True, exist_ok=False) selected = backend_main.selected_image_from_mask(image, mask, background_mode="neutral") overlay = backend_main.selection_overlay(image, mask) mask_pixels = int(np.count_nonzero(np.asarray(mask.convert("L")) > 0)) metadata = { "job_id": job_id, "source_filename": Path(image_path).name, "source_fingerprint": backend_main.selection_source_fingerprint(image), "mask_count": int(mask_count), "model_id": resolved_model, "selection_model_status": model_status, "model_status": "composed-clicked-masks", "selection_labels": labels, "background_mode": "neutral", "selection_infill_mode": "none", "selection_infill": {"mode": "none", "enabled": False}, "face_detection_source": "neutral-selection-cutout", "image_size": {"width": width, "height": height}, "mask_pixels": mask_pixels, "mask_coverage": mask_pixels / float(max(1, width * height)), } image.save(job_dir / "source.png") selected.save(job_dir / "selected_image.png") selected.save(job_dir / "selection_face_detection.png") mask.save(job_dir / "selection_mask.png") overlay.save(job_dir / "selection_overlay.png") backend_main.selection_tint(mask).save(job_dir / "selection_tint.png") (job_dir / "selection.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") return { "job_id": job_id, "selected": _safe_file(job_dir / "selected_image.png"), "overlay": _safe_file(job_dir / "selection_overlay.png"), "mask": backend_main.output_relative_path(job_dir / "selection_mask.png"), "labels": labels, "model": resolved_model, "model_status": model_status, "mask_coverage": metadata["mask_coverage"], } def _selection_job( image_path: str | Path, point_x: float, point_y: float, ) -> dict: _require_sam3_token() image = _load_selection_image(image_path) width, height = image.size normalized_point = { "x": float(np.clip(point_x / max(width, 1), 0.0, 1.0)), "y": float(np.clip(point_y / max(height, 1), 0.0, 1.0)), } mask, resolved_model, model_status, labels = backend_main.sam3_person_aware_selection_mask( image, [normalized_point], device="cuda", ) return _save_selection_job( image_path, image, mask, resolved_model=resolved_model, model_status=model_status, labels=labels, ) def select_object(image_path: str | Path | None, point_x: float, point_y: float) -> dict: if not image_path: raise ValueError("Upload a photo before selecting an object") return _selection_job(image_path, point_x, point_y) def _png_data_url(image: Image.Image) -> str: buffer = io.BytesIO() image.save(buffer, format="PNG", optimize=True) encoded = base64.b64encode(buffer.getvalue()).decode("ascii") return f"data:image/png;base64,{encoded}" def cleanup_expired_selection_precomputes( max_age_seconds: int = SAM3_PRECOMPUTE_TTL_SECONDS, *, now: float | None = None, ) -> int: cutoff = (time.time() if now is None else float(now)) - float(max_age_seconds) removed = 0 with backend_main.SELECTION_PRECOMPUTE_LOCK: expired = [ cache_id for cache_id, entry in backend_main.SELECTION_PRECOMPUTE_CACHE.items() if float(entry.get("created_at_epoch", 0.0)) < cutoff ] for cache_id in expired: backend_main.SELECTION_PRECOMPUTE_CACHE.pop(cache_id, None) removed += 1 return removed def _discard_precompute_source_pixels(precompute_id: str) -> None: with backend_main.SELECTION_PRECOMPUTE_LOCK: cached = backend_main.SELECTION_PRECOMPUTE_CACHE.get(precompute_id) if cached is not None: cached.pop("image", None) def _sam3_hover_region_map( masks: np.ndarray, scores: np.ndarray, labels: list[str], image_size: tuple[int, int], *, max_dimension: int = SAM3_HOVER_MAP_MAX_DIMENSION, ) -> tuple[Image.Image, dict[str, dict], tuple[int, int]]: import cv2 width, height = image_size scale = min(1.0, float(max_dimension) / float(max(width, height, 1))) map_width = max(1, round(width * scale)) map_height = max(1, round(height * scale)) mask_array = np.asarray(masks, dtype=bool) score_array = np.asarray(scores, dtype=np.float32).reshape(-1) if len(score_array) != len(mask_array) or len(labels) != len(mask_array): raise RuntimeError("SAM 3 returned inconsistent instance metadata") if not len(mask_array): raise RuntimeError("SAM 3 found no selectable objects in this image") winner_scores = np.full((map_height, map_width), -np.inf, dtype=np.float32) winners = np.full((map_height, map_width), -1, dtype=np.int32) for instance_index, mask in enumerate(mask_array): if (map_width, map_height) == (width, height): resized = mask else: resized = cv2.resize( mask.astype(np.uint8), (map_width, map_height), interpolation=cv2.INTER_NEAREST, ).astype(bool) coverage = float(np.mean(resized)) border_sides = sum( bool(np.any(edge)) for edge in (resized[0, :], resized[-1, :], resized[:, 0], resized[:, -1]) ) scene_context = bool( str(labels[instance_index]).strip().lower() != "person" and ( coverage >= SAM3_HOVER_CONTEXT_LARGE_COVERAGE or ( coverage >= SAM3_HOVER_CONTEXT_MIN_COVERAGE and border_sides >= SAM3_HOVER_CONTEXT_MIN_BORDER_SIDES ) ) ) if scene_context: continue hit_priority = float(score_array[instance_index]) - ( SAM3_HOVER_SPECIFICITY_WEIGHT * coverage ) update = resized & (hit_priority > winner_scores) winner_scores[update] = hit_priority winners[update] = instance_index covered = winners >= 0 region_map = np.zeros((map_height, map_width), dtype=np.uint32) regions: dict[str, dict] = {} minimum_region_pixels = max(4, round(map_width * map_height * 0.00001)) next_region_id = 1 for instance_index in range(len(mask_array)): winning_pixels = (winners == instance_index) & covered component_count, component_labels, stats, _centroids = cv2.connectedComponentsWithStats( winning_pixels.astype(np.uint8), connectivity=8, ) for component_index in range(1, component_count): area = int(stats[component_index, cv2.CC_STAT_AREA]) if area < minimum_region_pixels: continue if next_region_id >= 2**24: raise RuntimeError("SAM 3 produced too many hover regions") region_map[component_labels == component_index] = next_region_id regions[str(next_region_id)] = { "label": str(labels[instance_index]), "score": round(float(score_array[instance_index]), 6), "pixels": area, "instance_index": instance_index, } next_region_id += 1 if not regions: raise RuntimeError("SAM 3 found no selectable objects in this image") encoded = np.stack( ( region_map & 255, (region_map >> 8) & 255, (region_map >> 16) & 255, ), axis=-1, ).astype(np.uint8) return Image.fromarray(encoded, mode="RGB"), regions, (map_width, map_height) def prepare_object_selection(image_path: str | Path | None) -> tuple[str, dict]: if not image_path: raise ValueError("Upload a photo before selecting an object") _require_sam3_token() cleanup_expired_selection_precomputes() image = _load_selection_image(image_path) masks = np.empty((0, image.height, image.width), dtype=bool) scores = np.empty((0,), dtype=np.float32) labels: list[str] = [] resolved_model = "" try: masks, scores, labels, resolved_model = backend_main.compute_sam3_selection_instances( image, device="cuda", ) if resolved_model != SAM3_MODEL: raise RuntimeError("SAM 3 did not produce verified selection instances") finally: release_gpu_models() hover_map, regions, hover_size = _sam3_hover_region_map( masks, scores, labels, image.size, ) source_fingerprint = backend_main.selection_source_fingerprint(image) public_regions = { region_id: { "label": metadata["label"], "score": metadata["score"], "pixels": metadata["pixels"], } for region_id, metadata in regions.items() } used_instance_indices = sorted( {int(metadata["instance_index"]) for metadata in regions.values()} ) if image.width * image.height > SAM3_PRECOMPUTE_MAX_SOURCE_PIXELS: raise RuntimeError("The uploaded image is too large for cached object selection") if len(used_instance_indices) > SAM3_PRECOMPUTE_MAX_MASKS: raise RuntimeError("SAM 3 found too many objects to cache safely") compact_instance_indices = { instance_index: compact_index for compact_index, instance_index in enumerate(used_instance_indices) } compressed_masks: list[bytes] = [] compressed_bytes = 0 for instance_index in used_instance_indices: packed_mask = np.packbits( np.asarray(masks[instance_index], dtype=bool), axis=1, ) payload = zlib.compress(packed_mask.tobytes(order="C"), level=6) compressed_bytes += len(payload) if compressed_bytes > SAM3_PRECOMPUTE_MAX_COMPRESSED_BYTES: raise RuntimeError("The SAM 3 object map is too large to cache safely") compressed_masks.append(payload) precompute_id = uuid4().hex manifest = { "version": SAM3_HOVER_MAP_VERSION, "width": hover_size[0], "height": hover_size[1], "source_width": image.width, "source_height": image.height, "hit_map": _png_data_url(hover_map), "regions": public_regions, } state = { "state_version": SAM3_PRECOMPUTE_STATE_VERSION, "precompute_id": precompute_id, "source_fingerprint": source_fingerprint, "image_size": [image.width, image.height], "hover_size": [hover_size[0], hover_size[1]], "region_count": len(regions), "region_instances": { region_id: compact_instance_indices[int(metadata["instance_index"])] for region_id, metadata in regions.items() }, "model": resolved_model, # ZeroGPU executes GPU-decorated functions in forked workers. Returning # packed masks in gr.State is what carries them back to the app process; # process-local globals disappear when the worker exits. "selection_precompute": { "kind": "sam3-hover-session-v2", "compressed_masks": compressed_masks, "compressed_bytes": compressed_bytes, "mask_width": image.width, "labels": [str(labels[index]) for index in used_instance_indices], "model_id": resolved_model, "image_size": [image.width, image.height], "created_at_epoch": time.time(), }, } return json.dumps(manifest, separators=(",", ":")), state def _inline_selection_precompute(precompute_state: dict) -> dict | None: cached = precompute_state.get("selection_precompute") if not isinstance(cached, dict): return None if int(precompute_state.get("state_version", 0)) != SAM3_PRECOMPUTE_STATE_VERSION: raise ValueError("The object map uses an unsupported session format; refresh it") created_at = float(cached.get("created_at_epoch", 0.0)) if not np.isfinite(created_at) or time.time() - created_at > SAM3_PRECOMPUTE_TTL_SECONDS: raise ValueError("The object map expired; wait for SAM 3 to refresh it") image_size = cached.get("image_size") or [] if len(image_size) != 2: raise ValueError("The cached SAM 3 object map has invalid dimensions") width, height = (int(image_size[0]), int(image_size[1])) mask_width = int(cached.get("mask_width", 0)) compressed_masks = cached.get("compressed_masks") labels = cached.get("labels") or [] expected_bytes = (width + 7) // 8 valid_compressed_masks = isinstance(compressed_masks, (list, tuple)) and all( isinstance(payload, bytes) for payload in compressed_masks ) compressed_bytes = ( sum(len(payload) for payload in compressed_masks) if valid_compressed_masks else SAM3_PRECOMPUTE_MAX_COMPRESSED_BYTES + 1 ) if ( width <= 0 or height <= 0 or width * height > SAM3_PRECOMPUTE_MAX_SOURCE_PIXELS or mask_width != width or not valid_compressed_masks or len(compressed_masks) > SAM3_PRECOMPUTE_MAX_MASKS or compressed_bytes > SAM3_PRECOMPUTE_MAX_COMPRESSED_BYTES or len(labels) != len(compressed_masks) or str(cached.get("model_id")) != SAM3_MODEL ): raise ValueError("The cached SAM 3 object map is inconsistent") return { "compressed_masks": list(compressed_masks), "packed_row_bytes": expected_bytes, "mask_width": mask_width, "labels": [str(label) for label in labels], "model_id": str(cached["model_id"]), "image_size": (width, height), } def _selection_candidate(cached: dict, instance_index: int) -> np.ndarray: mask_count = ( len(cached["compressed_masks"]) if "compressed_masks" in cached else len(cached.get("packed_masks", [])) ) if not 0 <= instance_index < mask_count: raise ValueError("The highlighted SAM 3 region has an invalid instance index") width, height = cached["image_size"] if "compressed_masks" in cached: try: raw_mask = zlib.decompress(cached["compressed_masks"][instance_index]) except zlib.error as exc: raise ValueError("The cached SAM 3 object mask is corrupt") from exc expected_size = height * int(cached["packed_row_bytes"]) if len(raw_mask) != expected_size: raise ValueError("The cached SAM 3 object mask has an invalid size") packed_mask = np.frombuffer(raw_mask, dtype=np.uint8).reshape( height, int(cached["packed_row_bytes"]), ) else: packed_mask = cached["packed_masks"][instance_index] return np.unpackbits( packed_mask, axis=1, count=int(cached["mask_width"]), ).astype(bool) def _selection_component_at_point( candidate: np.ndarray, point: dict[str, float], precompute_state: dict, width: int, height: int, ) -> np.ndarray: pixel_point = backend_main.selection_points_to_pixels([point], width, height)[0] px = int(np.clip(round(pixel_point[0]), 0, max(0, width - 1))) py = int(np.clip(round(pixel_point[1]), 0, max(0, height - 1))) if not candidate[py, px]: hover_width, hover_height = precompute_state.get("hover_size") or [width, height] source_pixels_per_hover_pixel = max( width / max(1, int(hover_width)), height / max(1, int(hover_height)), ) radius = max(2, int(np.ceil(source_pixels_per_hover_pixel)) + 1) left = max(0, px - radius) right = min(width, px + radius + 1) top = max(0, py - radius) bottom = min(height, py + radius + 1) local_y, local_x = np.nonzero(candidate[top:bottom, left:right]) if not len(local_x): raise ValueError("No cached SAM 3 object covers that point; hover over a highlighted object") distances = (local_x + left - px) ** 2 + (local_y + top - py) ** 2 nearest = int(np.argmin(distances)) pixel_point = [float(local_x[nearest] + left), float(local_y[nearest] + top)] return backend_main.seeded_sam2_component(candidate, [pixel_point]) def select_precomputed_objects( image_path: str | Path | None, precompute_state: dict | None, selections: list[dict] | None, ) -> dict: if not image_path: raise ValueError("Upload a photo before selecting an object") if not precompute_state or not ( precompute_state.get("selection_precompute") or precompute_state.get("precompute_id") ): raise ValueError("Wait for SAM 3 to finish finding objects") if not isinstance(selections, list) or not selections: raise ValueError("Click at least one highlighted object before finishing") if len(selections) > SAM3_PRECOMPUTE_MAX_MASKS: raise ValueError("Too many highlighted objects were selected") image = _load_selection_image(image_path) fingerprint = backend_main.selection_source_fingerprint(image) if fingerprint != precompute_state.get("source_fingerprint"): raise ValueError("The uploaded image changed; wait for SAM 3 to refresh the object map") cached = _inline_selection_precompute(precompute_state) if cached is None: cleanup_expired_selection_precomputes() cached = backend_main.get_selection_precompute(str(precompute_state["precompute_id"])) region_instances = precompute_state.get("region_instances") or {} width, height = cached["image_size"] combined = np.zeros((height, width), dtype=bool) labels: list[str] = [] candidates: dict[int, np.ndarray] = {} seen_regions: set[int] = set() for selection in selections: try: region_id = int(selection["region_id"]) coordinates = np.asarray([selection["x"], selection["y"]], dtype=np.float64) instance_index = int(region_instances[str(region_id)]) except (KeyError, TypeError, ValueError) as exc: raise ValueError("A highlighted SAM 3 region is no longer available") from exc if region_id in seen_regions: continue if not np.all(np.isfinite(coordinates)): raise ValueError("Selection coordinates must be finite") seen_regions.add(region_id) point = { "x": float(np.clip(coordinates[0], 0.0, 1.0)), "y": float(np.clip(coordinates[1], 0.0, 1.0)), } if instance_index not in candidates: candidates[instance_index] = _selection_candidate(cached, instance_index) component = _selection_component_at_point( candidates[instance_index], point, precompute_state, width, height, ) combined |= component label = str(cached["labels"][instance_index]) if label not in labels: labels.append(label) if not np.any(combined): raise ValueError("The selected SAM 3 objects produced an empty mask") mask = Image.fromarray((combined.astype(np.uint8) * 255), mode="L") model_status = ( "sam3-concept-precomputed-point" if len(seen_regions) == 1 else "sam3-concept-precomputed-multi-point" ) return _save_selection_job( image_path, image, mask, resolved_model=str(cached["model_id"]), model_status=model_status, labels=labels, mask_count=len(seen_regions), ) def select_precomputed_object( image_path: str | Path | None, precompute_state: dict | None, normalized_x: float, normalized_y: float, region_id: int | None = None, ) -> dict: return select_precomputed_objects( image_path, precompute_state, [ { "x": normalized_x, "y": normalized_y, "region_id": region_id, } ], ) def release_gpu_models() -> None: try: backend_main.release_selection_models() except Exception: pass gc.collect() try: import torch if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception: pass def _depth_model_source() -> str: from huggingface_hub import snapshot_download return snapshot_download( repo_id=DEPTH_MODEL, revision=DEPTH_MODEL_REVISION, token=os.getenv("HF_TOKEN") or None, ) def generate_relief( image_path: str | Path | None, scope: str, selection: dict | None, print_scale_percent: float, relief_height_mm: float, base_thickness_mm: float, detail_samples: int, ) -> tuple[str, str, str, dict]: if not image_path: raise ValueError("Upload a photo before generating a relief") selected = str(scope).lower().startswith("select") if selected and not selection: raise ValueError("Click the object in the image before generating the selected-object relief") detail_samples = int(detail_samples) if detail_samples not in LOCAL_RELIEF_DETAIL_MULTIPLIERS: raise ValueError("Mesh detail must match a local frontend production preset") mesh_resolution_multiplier = LOCAL_RELIEF_DETAIL_MULTIPLIERS[detail_samples] x_mm, y_mm = local_relief_dimensions_mm(image_path, print_scale_percent) face_assets = ensure_face_assets() depth_model_source = _depth_model_source() data = { # Keep this request contract aligned with the local Next.js frontend. # Infer and refine from the complete photograph, then apply the exact # selection only when the final printable surface is emitted. "selection_mode": "context", "selection_subject_lock": "true" if selected else "false", "selection_emission_only": "true" if selected else "false", "depth_provider": "transformers", "depth_model": depth_model_source, "device": "auto", "depth_downsample_sharpening": "0.35", "depth_inference_precision": "float32", "target_dimension": str(int(detail_samples)), "z_scale": str(float(relief_height_mm)), "base_thickness_mm": str(float(base_thickness_mm)), "max_xy_size": str(max(x_mm, y_mm)), "detail_basis_mm": str(LOCAL_RELIEF_PRINTER_EDGE_MM), "invert": "false", "relief_polarity": "raised-print", "mesh_resolution_multiplier": str(mesh_resolution_multiplier), "printer_profile": "Bambu Lab P1S", "printer_max_x_mm": str(LOCAL_RELIEF_PRINTER_EDGE_MM), "printer_max_y_mm": str(LOCAL_RELIEF_PRINTER_EDGE_MM), "printer_max_z_mm": str(LOCAL_RELIEF_PRINTER_EDGE_MM), "printer_clearance_mm": "0", "print_scale_percent": str(float(print_scale_percent)), "sigma": "0.35", "detail_boost": "0.8", "printable_feature_depth_mm": "0.4", "feature_bridge_depth_mm": "0.8", "background_detail_boost": "2.4", "background_photo_detail_mm": "0.60", "trim_top_background": "true", "relief_gamma": "0.75", "base_border_px": "2", "detail_radius": "2.0", "low_percentile": "1.0", "high_percentile": "99.0", "max_relief_slope": "2.0", "nozzle_diameter_mm": "0.4", "minimum_feature_mm": "0.8", "face_refinement_mode": "auto", "face_detail_strength": "1.0", "face_feather_ratio": "0.20", "face_max_correction_ratio": "0.08", } if selected: data["selection_job_id"] = selection["job_id"] release_gpu_models() request_image_path = (selection.get("selected") or image_path) if selected else image_path with open(request_image_path, "rb") as image_file: response = BACKEND_CLIENT.post( "/process_image", files={"file": (Path(request_image_path).name, image_file, "image/png")}, data=data, ) if response.status_code != 200: raise _response_error(response) result = response.json() depth_parity = _require_deterministic_depth_parity(result) face_parity = _require_face_parity(result, selection) if result.get("selection_mode") != data["selection_mode"]: raise RuntimeError("Relief backend did not honor the local frontend selection mode") if selected: selection_context = result.get("selection_depth_context") if not isinstance(selection_context, dict) or not ( selection_context.get("method") == "full_scene_subject_locked_background_v1" and selection_context.get("subject_surface_locked") is True and selection_context.get("emission_scope") == "selected-mask-only" and result.get("selection_subject_lock") is True and result.get("selection_emission_only") is True and result.get("selection_crop") is None ): raise RuntimeError("Selected relief did not preserve the local full-source selected-emission contract") selection_emission = result.get("relief_postprocess", {}).get( "selection_emission" ) backing_connector = ( selection_emission.get("backing_connector") if isinstance(selection_emission, dict) else None ) backing_foundation = ( selection_emission.get("backing_foundation") if isinstance(selection_emission, dict) else None ) enclosed_hole_fill = ( selection_emission.get("enclosed_hole_fill") if isinstance(selection_emission, dict) else None ) if not isinstance(selection_emission, dict) or not ( selection_emission.get("enabled") is True and selection_emission.get("method") == "full_scene_depth_grounded_closed_hole_free_selection_emission_v3" and selection_emission.get("retained_unselected_pixels") == 0 and selection_emission.get("removed_selected_pixels") == 0 and selection_emission.get("unsupported_selected_mesh_pixels") == 0 and selection_emission.get("retained_selection_ratio") == 1.0 and isinstance(backing_foundation, dict) and backing_foundation.get("accepted") is True and backing_foundation.get("method") == "column_grounded_backing_foundation_v1" and backing_foundation.get("bounded_to_selection_bbox") is True and isinstance(backing_connector, dict) and backing_connector.get("accepted") is True and backing_connector.get("within_bridge_budget") is True and backing_connector.get("unsupported_selected_mesh_pixels") == 0 and isinstance(enclosed_hole_fill, dict) and enclosed_hole_fill.get("accepted") is True and enclosed_hole_fill.get("method") == "exterior_flood_enclosed_hole_fill_v1" and selection_emission.get("closed_hole_pixels_after") == 0 ): raise RuntimeError( "Selected relief failed exact bounded-emission validation" ) job_dir = OUTPUT_DIR / result["job_id"] stl_path = OUTPUT_DIR / result["stl_model"] preview_path = job_dir / "output_relief_preview.png" if not preview_path.is_file(): preview_path = job_dir / "output_depth_preview.png" diagnostics_path = OUTPUT_DIR / result["diagnostics"] diagnostics = result.get("stl_diagnostics", {}) summary = _diagnostic_summary(diagnostics, model=DEPTH_MODEL) summary["dimensions_mm"] = { "x": x_mm, "y": y_mm, "z": float(relief_height_mm) + float(base_thickness_mm), } summary["relief_height_mm"] = float(relief_height_mm) summary["base_thickness_mm"] = float(base_thickness_mm) summary["scope"] = "selected-objects-full-source-depth" if selected else "full-scene" summary["selection_mode"] = result.get("selection_mode", data["selection_mode"]) summary["local_relief_parity"] = ( "full-source-deterministic-fp32-face-parity-v3" ) summary["depth_parity"] = depth_parity summary["face_refinement"] = face_parity summary["face_assets"] = face_assets summary["inpainting"] = False return ( _safe_file(stl_path), _safe_file(stl_path), _safe_file(preview_path), {"summary": summary, "full_report": _safe_file(diagnostics_path)}, ) def generate_diorama( image_path: str | Path | None, scope: str, selection: dict | None, max_size_mm: float, scene_depth_mm: float, base_thickness_mm: float, ) -> tuple[str, str, str, str, dict]: if not image_path: raise ValueError("Upload a photo before generating a scene diorama") selected = str(scope).lower().startswith("select") if selected and not selection: raise ValueError("Click the object in the image before building selected diorama layers") release_gpu_models() depth_model_source = _depth_model_source() data = { "mask_paths_json": json.dumps([selection["mask"]] if selection else []), "selection_labels_json": json.dumps([selection.get("labels", [])] if selection else []), "depth_provider": "transformers", "depth_model": depth_model_source, "device": "cuda", "max_size_mm": str(float(max_size_mm)), "scene_depth_mm": str(float(scene_depth_mm)), "base_thickness_mm": str(float(base_thickness_mm)), } with open(image_path, "rb") as image_file: response = BACKEND_CLIENT.post( "/process_scene_diorama", files={"file": (Path(image_path).name, image_file, "image/png")}, data=data, ) if response.status_code != 200: raise _response_error(response) result = response.json() stl_path = OUTPUT_DIR / result["stl_model"] glb_path = OUTPUT_DIR / result["scene_model"] preview_path = OUTPUT_DIR / result["preview"] diagnostics = result.get("stl_diagnostics", {}) summary = _diagnostic_summary(diagnostics, model=DEPTH_MODEL) summary["scene_mode"] = "layered-camera-free-diorama" summary["selected_layers"] = 1 if selected else 0 summary["inpainting"] = False return ( _safe_file(glb_path), _safe_file(stl_path), _safe_file(glb_path), _safe_file(preview_path), summary, ) def _patch_triposg_source(source_dir: Path) -> None: inference_path = source_dir / "triposg" / "inference_utils.py" vae_path = source_dir / "triposg" / "models" / "autoencoders" / "autoencoder_kl_triposg.py" inference_text = inference_path.read_text(encoding="utf-8") inference_text = inference_text.replace( "from diso import DiffDMC", "try:\n from diso import DiffDMC\nexcept ImportError:\n DiffDMC = None", 1, ) inference_path.write_text(inference_text, encoding="utf-8") vae_text = vae_path.read_text(encoding="utf-8") vae_text = vae_text.replace( " q = self.proj_query(q)", " q = self.proj_query(q.to(dtype=self.proj_query.weight.dtype))", 1, ) vae_path.write_text(vae_text, encoding="utf-8") def ensure_triposg_source() -> Path: with _TRIPOSG_LOCK: source = _clone_exact_repository(TRIPOSG_REPOSITORY, TRIPOSG_SOURCE_REVISION, TRIPOSG_DIR) _patch_triposg_source(source) return source def _load_triposg_models(): import torch from huggingface_hub import snapshot_download source = ensure_triposg_source() scripts_dir = source / "scripts" for path in (source, scripts_dir): if str(path) not in sys.path: sys.path.insert(0, str(path)) from triposg.pipelines.pipeline_triposg import TripoSGPipeline triposg_weights = Path( snapshot_download( repo_id="VAST-AI/TripoSG", revision=TRIPOSG_MODEL_REVISION, token=os.getenv("HF_TOKEN") or None, ) ) pipeline = TripoSGPipeline.from_pretrained( triposg_weights, torch_dtype=torch.float16, low_cpu_mem_usage=True, ).to("cuda") return pipeline def _prepare_selected_mesh_image(selection: dict, output_path: Path, canvas_size: int = 512) -> Image.Image: selection_dir = OUTPUT_DIR / "selection" / selection["job_id"] source_path = selection_dir / "source.png" mask_path = selection_dir / "selection_mask.png" _safe_file(source_path) _safe_file(mask_path) with Image.open(source_path) as source_file, Image.open(mask_path) as mask_file: source = ImageOps.exif_transpose(source_file).convert("RGB") mask = mask_file.convert("L") bbox = mask.getbbox() if bbox is None: raise ValueError("The selected object mask is empty") left, top, right, bottom = bbox margin = max(2, int(round(max(right - left, bottom - top) * 0.08))) crop_box = ( max(0, left - margin), max(0, top - margin), min(source.width, right + margin), min(source.height, bottom + margin), ) source_crop = source.crop(crop_box) mask_crop = mask.crop(crop_box) selected_crop = Image.composite(source_crop, Image.new("RGB", source_crop.size, "white"), mask_crop) selected_crop.thumbnail((int(canvas_size * 0.9), int(canvas_size * 0.9)), Image.Resampling.LANCZOS) canvas = Image.new("RGB", (canvas_size, canvas_size), "white") canvas.paste( selected_crop, ((canvas_size - selected_crop.width) // 2, (canvas_size - selected_crop.height) // 2), ) canvas.save(output_path) return canvas def generate_full_mesh( image_path: str | Path | None, scope: str, selection: dict | None, max_dimension_mm: float, seed: int, ) -> tuple[str, str, str, dict]: if not image_path: raise ValueError("Upload a photo before generating a full mesh") if not str(scope).lower().startswith("select") or not selection: raise ValueError("Full Mesh requires Select object; click the object before generating") import torch import trimesh release_gpu_models() job_id = uuid4().hex job_dir = OUTPUT_DIR / "full-mesh" / job_id job_dir.mkdir(parents=True, exist_ok=False) raw_glb = job_dir / "output_mesh_raw.glb" repaired_mesh = job_dir / "output_mesh_repaired.ply" final_glb = job_dir / "output_mesh.glb" final_stl = job_dir / "output_model.stl" metrics: dict = {} started = time.perf_counter() torch.cuda.reset_peak_memory_stats() pipeline = None try: pipeline = _load_triposg_models() prepared = _prepare_selected_mesh_image(selection, job_dir / "mesh_input.png") with torch.inference_mode(): samples = pipeline( image=prepared, generator=torch.Generator(device="cuda").manual_seed(int(seed)), num_inference_steps=50, guidance_scale=7.0, use_flash_decoder=False, ).samples[0] mesh = trimesh.Trimesh( samples[0].astype(np.float32), np.ascontiguousarray(samples[1]), process=False, ) if not len(mesh.vertices) or not len(mesh.faces): raise RuntimeError("TripoSG returned an empty mesh") mesh.export(raw_glb) repair_mesh_for_printable_stl( raw_glb, repaired_mesh, "printable", max_normalized_face_density_log1p=9.95, preconditioner="adaptive-voxel-close", voxel_resolution=128, voxel_fill_method="orthographic", smoothing_iterations=2, allow_convex_hull_fallback=False, metrics=metrics, ) postprocess_mesh_for_stl( repaired_mesh, final_glb, target_max_dimension=float(max_dimension_mm), target_faces=10000, max_normalized_face_density_log1p=9.95, preserve_printability=True, ) convert_mesh_to_stl(final_glb, final_stl) diagnostics = json_safe_stl_diagnostics(stl_diagnostics(final_stl)) summary = _diagnostic_summary(diagnostics, model="VAST-AI/TripoSG") summary.update( { "source_revision": TRIPOSG_SOURCE_REVISION, "model_revision": TRIPOSG_MODEL_REVISION, "foreground_model": SAM3_MODEL, "foreground_model_revision": SAM3_MODEL_REVISION, "runtime_seconds": round(time.perf_counter() - started, 3), "peak_vram_gib": round(torch.cuda.max_memory_allocated() / (1024**3), 3), "repair": metrics, "scope": "selected-object", "inpainting": False, } ) (job_dir / "diagnostics.json").write_text( json.dumps({"stl": diagnostics, "summary": summary}, indent=2, allow_nan=False), encoding="utf-8", ) return _safe_file(final_stl), _safe_file(final_stl), _safe_file(final_glb), summary finally: del pipeline gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() def cleanup_expired_outputs(max_age_seconds: int = 6 * 60 * 60) -> int: cleanup_expired_selection_precomputes() if not OUTPUT_DIR.is_dir(): return 0 now = time.time() removed = 0 for child in OUTPUT_DIR.iterdir(): if not child.is_dir(): continue candidates = list(child.iterdir()) if child.name in {"selection", "full-mesh"} else [child] for candidate in candidates: if not candidate.is_dir(): continue try: if now - candidate.stat().st_mtime > max_age_seconds: shutil.rmtree(candidate) removed += 1 except OSError: continue return removed