Spaces:
Running on Zero
Running on Zero
| """Input preparation, checkpoint selection, inference, and portable result bundles.""" | |
| from dataclasses import dataclass | |
| import gc | |
| import json | |
| import math | |
| import os | |
| from pathlib import Path | |
| import secrets | |
| import shutil | |
| import tempfile | |
| import threading | |
| import time | |
| import warnings | |
| import zipfile | |
| import numpy as np | |
| from PIL import Image | |
| from src.utils.camera_io import IMAGE_EXTENSIONS, list_image_files, load_cameras_json | |
| ROOT = Path(__file__).resolve().parents[1] | |
| RGB_MODE = 'RGB-only' | |
| POSE_MODE = 'Pose-conditioned' | |
| class Options: | |
| mode: str = RGB_MODE | |
| width: int = 896 | |
| height: int = 896 | |
| opacity: float = 0.05 | |
| sky: bool = False | |
| align_cameras: bool = True | |
| max_images: int | None = None | |
| voxel_size: float = 0.0018 | |
| max_scale_pct: float = 99.5 | |
| filter_flying_points_infer: bool = False | |
| def validate(self): | |
| if not math.isfinite(self.max_scale_pct) or not 0 < self.max_scale_pct <= 100: | |
| raise ValueError('max_scale_pct must be greater than 0 and at most 100.') | |
| if not math.isfinite(self.voxel_size) or not 0.0015 <= self.voxel_size <= 0.0020: | |
| raise ValueError('Voxel size must be between 0.0015 and 0.0020.') | |
| if self.max_images is not None and (not isinstance(self.max_images, int) or self.max_images < 0): | |
| raise ValueError('Max images must be 0 (all) or a positive integer.') | |
| if self.mode not in (RGB_MODE, POSE_MODE): | |
| raise ValueError('Choose RGB-only or Pose-conditioned inference.') | |
| for value in (self.width, self.height): | |
| if not isinstance(value, int) or value <= 0 or value % 14: | |
| raise ValueError('Width and height must be positive multiples of 14.') | |
| if not math.isfinite(self.opacity) or not 0 <= self.opacity <= 1: | |
| raise ValueError('Opacity must be between 0 and 1.') | |
| def ordered_images(files, minimum=1): | |
| paths = sorted((Path(p) for p in files or []), key=lambda p: p.name) | |
| if len(paths) < minimum: | |
| raise ValueError('Upload at least one image of your scene.') | |
| if any(not p.is_file() or p.suffix.lower() not in IMAGE_EXTENSIONS for p in paths): | |
| raise ValueError('Use supported image files: PNG, JPEG, WebP, BMP, or TIFF.') | |
| stems = [p.stem for p in paths] | |
| if len(stems) != len(set(stems)): | |
| raise ValueError('Image filename stems must be unique, including across file extensions.') | |
| return paths | |
| def uploaded_resolution(files): | |
| """Default to 896 square without exceeding any uploaded image dimension.""" | |
| width = height = 896 | |
| for path in ordered_images(files, minimum=0): | |
| with Image.open(path) as image: | |
| width, height = min(width, image.width), min(height, image.height) | |
| width, height = width // 14 * 14, height // 14 * 14 | |
| if min(width, height) < 14: | |
| raise ValueError('Each uploaded image must be at least 14 pixels wide and high.') | |
| return width, height | |
| def stage_inputs(files, camera, options, destination): | |
| """Preserve filenames so camera IDs and exported image IDs remain consistent.""" | |
| options.validate() | |
| paths = ordered_images(files) | |
| if options.max_images: | |
| paths = paths[:options.max_images] | |
| destination = Path(destination) | |
| images = destination / 'images' | |
| images.mkdir(parents=True, exist_ok=False) | |
| if options.mode == POSE_MODE: | |
| if not camera: | |
| raise ValueError('Pose-conditioned inference requires prior_camera.json.') | |
| extrinsics, intrinsics = load_cameras_json(camera) | |
| for path in paths: | |
| if path.stem not in extrinsics or path.stem not in intrinsics: | |
| raise ValueError(f'Camera JSON is missing poses or intrinsics for {path.stem}.') | |
| with Image.open(path) as image: | |
| cx, cy = intrinsics[path.stem][:2, 2] | |
| if not (0 < cx < image.width and 0 < cy < image.height): | |
| raise ValueError(f'Camera principal point lies outside image {path.name}.') | |
| camera_path = destination / 'prior_camera.json' | |
| shutil.copyfile(camera, camera_path) | |
| else: | |
| # Switching back to RGB-only never silently uses a previously uploaded prior. | |
| camera_path = None | |
| for path in paths: | |
| with Image.open(path) as image: | |
| image.verify() | |
| shutil.copyfile(path, images / path.name) | |
| return images, camera_path | |
| def extract_video_frames(video, count, destination): | |
| import cv2 | |
| if not video or not Path(video).is_file(): | |
| raise ValueError('Upload a video first.') | |
| if not 2 <= int(count) <= 32: | |
| raise ValueError('Extract between 2 and 32 frames.') | |
| capture = cv2.VideoCapture(str(video)) | |
| destination = Path(destination) | |
| destination.mkdir(parents=True, exist_ok=False) | |
| try: | |
| total = int(capture.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| if total < 2: | |
| raise ValueError('The video could not be decoded or has fewer than two frames.') | |
| paths = [] | |
| for index in np.linspace(0, total - 1, min(int(count), total), dtype=int): | |
| capture.set(cv2.CAP_PROP_POS_FRAMES, int(index)) | |
| ok, frame = capture.read() | |
| if not ok: | |
| continue | |
| path = destination / f'frame_{index:08d}.jpg' | |
| if cv2.imwrite(str(path), frame): | |
| paths.append(str(path)) | |
| if len(paths) < 2: | |
| raise ValueError('Fewer than two frames could be decoded; try a different video.') | |
| return paths | |
| finally: | |
| capture.release() | |
| def inspect_panorama(path): | |
| """Validate the same ERP convention consumed by scripts/convert_panorama.py.""" | |
| if not path or not Path(path).is_file(): | |
| raise ValueError('Upload one full 360° panorama first.') | |
| try: | |
| with warnings.catch_warnings(): | |
| warnings.simplefilter('ignore', Image.DecompressionBombWarning) | |
| with Image.open(path) as image: | |
| if image.format not in ('JPEG', 'PNG', 'WEBP'): | |
| raise ValueError('Use a JPG, PNG, or WebP panorama.') | |
| if image.width != 2 * image.height: | |
| raise ValueError('Use a full 360° × 180° equirectangular panorama with a 2:1 aspect ratio.') | |
| if image.getexif().get(274, 1) != 1: | |
| raise ValueError('Apply EXIF orientation to the panorama before uploading it.') | |
| size = image.size | |
| # Reading PNG EXIF may load the image; verify a freshly opened stream. | |
| with Image.open(path) as image: | |
| image.verify() | |
| return size | |
| except (OSError, Image.DecompressionBombError) as error: | |
| raise ValueError('The panorama could not be decoded. Upload a valid JPG, PNG, or WebP image.') from error | |
| def resolve_weights(config_path, student=None, geometry=None): | |
| from omegaconf import OmegaConf | |
| cfg = OmegaConf.load(config_path) | |
| repo = os.environ.get('SPLATANYTHING_MODEL_REPO') | |
| revision = os.environ.get('SPLATANYTHING_MODEL_REVISION', 'main') | |
| choices = [ | |
| ('SplatAnything_pretrained', student or os.environ.get('SPLATANYTHING_CKPT'), | |
| 'SPLATANYTHING_MODEL_FILE'), | |
| ('worldmirror2_pretrained', geometry or os.environ.get('WORLDMIRROR2_CKPT'), | |
| 'SPLATANYTHING_GEOMETRY_FILE'), | |
| ] | |
| for key, override, file_var in choices: | |
| path = Path(override or cfg.wrapper[key]).expanduser() | |
| if not path.is_absolute(): | |
| path = (Path.cwd() if override else ROOT) / path | |
| if not path.is_file() and repo and not override: | |
| from huggingface_hub import hf_hub_download | |
| filename = os.environ.get(file_var, Path(cfg.wrapper[key]).name) | |
| path = Path(hf_hub_download(repo, filename, revision=revision)) | |
| if not path.is_file(): | |
| raise FileNotFoundError( | |
| f'{key} not found: {path}. Configure local checkpoints, ' | |
| 'SPLATANYTHING_CKPT / WORLDMIRROR2_CKPT, or SPLATANYTHING_MODEL_REPO.') | |
| # Hub snapshot links point to extensionless blobs. Keep the logical | |
| # filename so safetensors is decoded with the correct loader. | |
| cfg.wrapper[key] = str(path.absolute()) | |
| return cfg | |
| def to_cpu(value): | |
| import torch | |
| if isinstance(value, torch.Tensor): | |
| return value.detach().cpu() | |
| if isinstance(value, dict): | |
| return {key: to_cpu(item) for key, item in value.items()} | |
| if isinstance(value, (list, tuple)): | |
| return type(value)(to_cpu(item) for item in value) | |
| return value | |
| class Pipeline: | |
| def __init__(self, config_path=None, device='cuda', student=None, geometry=None, sky_path=None): | |
| self.config_path = Path(config_path or ROOT / 'configs/inference.yaml') | |
| self.device = device | |
| self.student_path, self.geometry_path = student, geometry | |
| self._explicit_sky_path = sky_path or os.environ.get('SPLATANYTHING_SKY_CKPT') | |
| self.sky_path = self._explicit_sky_path or str(ROOT / 'checkpoints/skyseg.onnx') | |
| self.model = self.teacher = self.segmenter = self.config = None | |
| self.lock = threading.RLock() | |
| def initialize(self): | |
| with self.lock: | |
| self._initialize() | |
| def _initialize(self): | |
| """Called before GPU scheduling in ZeroGPU and lazily for ordinary local runs.""" | |
| if self.model is not None: | |
| return | |
| import torch | |
| from omegaconf import OmegaConf | |
| from src.models.models.gaussian_worldmirror import GaussianWorldMirror | |
| from src.models.models.geometry_worldmirror import GeometryWorldMirror | |
| from src.utils.checkpoints import load_model_weights | |
| cfg = resolve_weights(self.config_path, self.student_path, self.geometry_path) | |
| if str(self.device).startswith('cuda') and not torch.cuda.is_available(): | |
| raise RuntimeError('CUDA is unavailable. Select a GPU runtime or use --device cpu for a small smoke test.') | |
| # Keep the conditioning layers resident, while selecting the exact CLI token | |
| # layout per request. A shared lock protects mode and renderer state. | |
| model = GaussianWorldMirror(**OmegaConf.to_container(cfg.wrapper.model), enable_cond=True) | |
| load_model_weights(model, cfg.wrapper.SplatAnything_pretrained) | |
| model.to(self.device).eval().requires_grad_(False) | |
| teacher = GeometryWorldMirror(enable_cond=True) | |
| load_model_weights(teacher, cfg.wrapper.worldmirror2_pretrained) | |
| teacher.to(self.device).eval().requires_grad_(False) | |
| self.model, self.teacher, self.config = model, teacher, cfg | |
| gc.collect() | |
| def _set_mode(self, enabled): | |
| for model in (self.model, self.teacher): | |
| model.enable_cond = enabled | |
| encoder = model.visual_geometry_transformer | |
| if not hasattr(encoder, '_demo_conditioned_patch_start'): | |
| encoder._demo_conditioned_patch_start = encoder.patch_start_idx | |
| encoder.enable_cond = enabled | |
| encoder.patch_start_idx = encoder._demo_conditioned_patch_start - (0 if enabled else 2) | |
| self.model.gs_renderer.use_gt_camera = enabled | |
| def prepare_inputs(self, images, camera, options): | |
| """Run CPU image preparation and sky segmentation before GPU scheduling.""" | |
| from src.utils.scene_io import build_views | |
| from src.utils.sky_segmentation import SkySegmenter | |
| options.validate() | |
| with self.lock: | |
| if options.sky and self.segmenter is None: | |
| path = Path(self.sky_path).expanduser() | |
| if not path.is_file(): | |
| if self._explicit_sky_path: | |
| raise ValueError('Sky segmentation model not found at the configured SPLATANYTHING_SKY_CKPT path.') | |
| from huggingface_hub import hf_hub_download | |
| model_repo = os.environ.get('SPLATANYTHING_MODEL_REPO') | |
| repo = os.environ.get('SPLATANYTHING_SKY_REPO', model_repo or 'ritianyu/SplatAnything') | |
| revision = os.environ.get('SPLATANYTHING_SKY_REVISION') or ( | |
| os.environ.get('SPLATANYTHING_MODEL_REVISION', 'main') if repo == model_repo else 'main') | |
| filename = os.environ.get('SPLATANYTHING_SKY_FILE', 'skyseg.onnx') | |
| try: | |
| path = Path(hf_hub_download(repo, filename, revision=revision)) | |
| except Exception as error: | |
| raise ValueError(f'Could not download the sky segmentation model from {repo}. Check the sky model settings.') from error | |
| self.sky_path = str(path.absolute()) | |
| self.segmenter = SkySegmenter(self.sky_path, device='cpu') | |
| if options.sky: | |
| self.segmenter._ensure_session() | |
| views, _, paths = build_views( | |
| str(images), str(camera) if options.mode == POSE_MODE else None, | |
| (options.width, options.height), 'cpu', cam_align=options.align_cameras, | |
| sky_segmenter=self.segmenter if options.sky else None) | |
| return views, paths | |
| def predict(self, images, camera, options, prepared=None): | |
| import torch | |
| from src.utils.scene_io import build_views | |
| with self.lock: | |
| self.initialize() | |
| enabled = options.mode == POSE_MODE | |
| self._set_mode(enabled) | |
| options.validate() | |
| self.model.gs_renderer.voxel_size = options.voxel_size | |
| self.model.gs_renderer.max_scale_pct = options.max_scale_pct | |
| self.model.gs_renderer.filter_flying_points_infer = options.filter_flying_points_infer | |
| if options.sky and prepared is None: | |
| prepared = self.prepare_inputs(images, camera, options) | |
| if prepared is not None: | |
| cpu_views, paths = prepared | |
| views = {key: value.to(self.device) for key, value in cpu_views.items()} | |
| else: | |
| views, _, paths = build_views( | |
| str(images), str(camera) if enabled else None, | |
| (options.width, options.height), self.device, | |
| cam_align=options.align_cameras) | |
| flags = [1, 0, 1] if enabled else [0, 0, 0] | |
| self.model.gs_renderer.filter_sky_infer = options.sky | |
| try: | |
| with torch.inference_mode(): | |
| geometry = self.teacher(views, cond_flags=flags) | |
| predictions = self.model(views, geometry, cond_flags=flags) | |
| result = to_cpu(predictions), to_cpu(views), paths | |
| del geometry, predictions, views | |
| return result | |
| finally: | |
| # Token/rotary caches are device-local and must not retain CUDA | |
| # allocations across a ZeroGPU reservation or a mode switch. | |
| for model in (self.model, self.teacher): | |
| encoder = model.visual_geometry_transformer | |
| if hasattr(encoder, 'pos_getter'): | |
| encoder.pos_getter.position_cache.clear() | |
| for module in model.modules(): | |
| if hasattr(module, 'frequency_cache'): | |
| module.frequency_cache.clear() | |
| if str(self.device).startswith('cuda'): | |
| torch.cuda.empty_cache() | |
| def export_preview(predictions, paths, output, options, config): | |
| """Write only the Gaussians and small metadata needed before showing the viewer.""" | |
| from omegaconf import OmegaConf | |
| from src.utils.scene_io import export_gaussians | |
| output = Path(output) | |
| output.mkdir(parents=True, exist_ok=True) | |
| splats = predictions['splats'] | |
| if not len(splats['means'][0]) or not (splats['opacities'][0] >= options.opacity).any(): | |
| raise ValueError('No Gaussians remain after filtering. Lower opacity or use more overlapping views.') | |
| export_gaussians(predictions, output, options.opacity, None) | |
| cameras = {'extrinsics': [], 'intrinsics': []} | |
| for i, image in enumerate(paths): | |
| for name, key in [('extrinsics', 'export_camera_poses'), ('intrinsics', 'export_camera_intrs')]: | |
| cameras[name].append({'camera_id': Path(image).stem, | |
| 'matrix': predictions[key][0, i].tolist()}) | |
| (output / 'cameras.json').write_text(json.dumps(cameras, indent=2) + '\n') | |
| # Downloaded metadata contains portable filenames rather than server paths. | |
| public_config = OmegaConf.create(OmegaConf.to_container(config)) | |
| for key in ('worldmirror2_pretrained', 'SplatAnything_pretrained'): | |
| public_config.wrapper[key] = Path(public_config.wrapper[key]).name | |
| public_config.wrapper.enable_cond = options.mode == POSE_MODE | |
| public_config.wrapper.model.use_gt_camera = options.mode == POSE_MODE | |
| public_config.wrapper.infer_opacity_thresh = options.opacity | |
| public_config.wrapper.model.gs_voxel_size = options.voxel_size | |
| public_config.wrapper.model.max_scale_pct = options.max_scale_pct | |
| public_config.wrapper.model.filter_flying_points_infer = options.filter_flying_points_infer | |
| OmegaConf.save(public_config, output / 'config.yaml') | |
| metadata = dict(mode=options.mode, resolution=[options.width, options.height], | |
| images=[Path(p).name for p in paths], opacity=options.opacity, | |
| sky_seg=options.sky, camera_alignment=options.align_cameras, | |
| max_images=options.max_images, voxel_size=options.voxel_size, | |
| max_scale_pct=options.max_scale_pct, | |
| filter_flying_points_infer=options.filter_flying_points_infer) | |
| (output / 'run.json').write_text(json.dumps(metadata, indent=2) + '\n') | |
| return output / 'gaussians.ply' | |
| def export_auxiliary(predictions, views, paths, output): | |
| """Export full-resolution raw points, depth and RGB after the viewer is ready. | |
| Paths are used only for image names: staged input folders may already be gone. | |
| All tensors passed here must be CPU tensors, independent of the GPU request. | |
| """ | |
| from src.utils.scene_io import export_points, export_depth_rgb | |
| export_points(predictions, Path(output)) | |
| export_depth_rgb(predictions, views, paths, Path(output)) | |
| def export_result(predictions, views, paths, output, options, config): | |
| """Synchronous export for callers that explicitly need the complete result.""" | |
| ply = export_preview(predictions, paths, output, options, config) | |
| export_auxiliary(predictions, views, paths, output) | |
| return ply | |
| def package_result(output): | |
| output = Path(output) | |
| archive = output / 'result.zip' | |
| with tempfile.NamedTemporaryFile(dir=output.parent, suffix='.zip', delete=False) as temporary: | |
| staged = Path(temporary.name) | |
| try: | |
| with zipfile.ZipFile(staged, 'w', zipfile.ZIP_DEFLATED, compresslevel=1) as handle: | |
| for path in sorted(output.rglob('*')): | |
| if path.is_file() and path != archive and path.name not in {'preview.compressed.ply', 'viewer-settings.json'}: | |
| compressed = path.suffix.lower() in {'.png', '.jpg', '.jpeg', '.webp'} | |
| handle.write(path, path.relative_to(output), | |
| compress_type=zipfile.ZIP_STORED if compressed else zipfile.ZIP_DEFLATED) | |
| staged.replace(archive) | |
| finally: | |
| staged.unlink(missing_ok=True) | |
| return archive | |
| class ResultStore: | |
| """Unpredictable per-result links, with a bounded lifetime and no arbitrary paths.""" | |
| def __init__(self, root=None, lifetime=24 * 3600): | |
| self.root = Path(root or tempfile.mkdtemp(prefix='splatanything-results-')).resolve() | |
| self.root.mkdir(parents=True, exist_ok=True) | |
| self.lifetime = lifetime | |
| self.records = {} | |
| self.exports = {} | |
| self.lock = threading.Lock() | |
| def reserve(self): | |
| with self.lock: | |
| now = time.time() | |
| for token, (path, expiry) in list(self.records.items()): | |
| task = self.exports.get(token) | |
| if expiry < now and (task is None or task['done'].is_set()): | |
| shutil.rmtree(path, ignore_errors=True) | |
| del self.records[token] | |
| self.exports.pop(token, None) | |
| token = secrets.token_urlsafe(32) | |
| path = self.root / secrets.token_hex(16) | |
| path.mkdir() | |
| self.records[token] = path, now + self.lifetime | |
| return token, path | |
| def begin_export(self, token): | |
| with self.lock: | |
| self.exports[token] = {'done': threading.Event(), 'html_done': threading.Event(), 'error': False} | |
| def finish_html(self, token): | |
| with self.lock: | |
| self.exports[token]['html_done'].set() | |
| def wait_for_html(self, token): | |
| with self.lock: | |
| task = self.exports.get(token) | |
| if task is not None: | |
| task['html_done'].wait() | |
| def finish_export(self, token, failed=False): | |
| with self.lock: | |
| task = self.exports[token] | |
| task['error'] = failed | |
| task['html_done'].set() | |
| task['done'].set() | |
| def export_status(self, token): | |
| with self.lock: | |
| record = self.records.get(token) | |
| if record is None or record[1] < time.time(): | |
| return None | |
| task = self.exports.get(token) | |
| if task is None: | |
| return 'ready' | |
| if not task['done'].is_set(): | |
| return 'preparing' | |
| return 'failed' if task['error'] else 'ready' | |
| def wait_for_export(self, token): | |
| with self.lock: | |
| task = self.exports.get(token) | |
| if task is not None: | |
| task['done'].wait() | |
| return self.export_status(token) | |
| def resolve(self, token, filename): | |
| if filename not in {'scene.html', 'gaussians.ply', 'result.zip', 'cameras.json', | |
| 'preview.compressed.ply', 'viewer-settings.json'}: | |
| return None | |
| with self.lock: | |
| record = self.records.get(token) | |
| if record is None or record[1] < time.time(): | |
| return None | |
| path = record[0] / filename | |
| return path if path.is_file() else None | |