zero_gpu.py

"""Allocate a shared GPU only for reconstruction; export on the CPU afterward."""
import spaces
import torch
from infinisplat import Pipeline
from quota import GPU_SECONDS, require_free_quota
from owner_quota import sponsored_gpu_context

_pipeline = None


def initialize(model):
    global _pipeline
    if not model:
        raise RuntimeError('Configure INFINISPLAT_MODEL before starting ZeroGPU.')
    _pipeline = Pipeline.from_pretrained(model, device='cuda:0')


def to_cpu(value):
    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):
        return [to_cpu(item) for item in value]
    if isinstance(value, tuple):
        return tuple(to_cpu(item) for item in value)
    return value


@spaces.GPU(duration=GPU_SECONDS, size='large')
def _reconstruct_on_gpu(files, cameras, mode, size, target_wh=None):
    if _pipeline is None:
        raise RuntimeError('Model initialization is incomplete.')
    scene = _pipeline.reconstruct(files, cameras=cameras, pose_mode=mode,
                                  long_side=size, max_views=50, target_wh=target_wh)
    # The demo only exports splats and cameras. Keep dense depth, RGB and raw
    # points on the GPU until this call ends instead of transferring them.
    return to_cpu({key: scene[key] for key in
                   ('splats', 'c2w', 'intrinsics', 'metadata', 'camera_prior')
                   if key in scene})


def reconstruct_on_gpu(files, cameras, mode, size, target_wh=None):
    # Run on the CPU before the decorator can reserve any shared GPU time.
    # The owner explicitly sponsors visitor runs. Keep the included-quota guard
    # before scheduling and never send the owner's long-lived token to browsers.
    require_free_quota()
    with sponsored_gpu_context():
        return _reconstruct_on_gpu(files, cameras, mode, size, target_wh=target_wh)