worker/handler.py

"""RunPod serverless handler: video (base64) -> Instant4D pipeline -> .splatv.

Runs inside the worker image (worker/Dockerfile), which layers the RunPod SDK
and SAM 2 on top of the proven 4dgs Space image. The pipeline module is the
same one the Docker Space runs — one codebase, two deployment targets.

Input  : {"video_b64": ..., "max_frames": 96, "target_fps": 24}
Output : {"filename": ..., "splatv_b64": ...} inline, OR
         {"filename": ..., "splatv_url": ...} when the endpoint has
         BUCKET_ENDPOINT_URL/BUCKET_ACCESS_KEY_ID/BUCKET_SECRET_ACCESS_KEY
         configured - RunPod does NOT upload handler return values to S3
         automatically, so this handler calls
         ``runpod.serverless.utils.rp_upload.upload_file_to_bucket`` itself
         whenever
         those env vars are present, and only falls back to inline base64
         (rejected past 180 MB) when they are not.
Progress: runpod.serverless.progress_update(job, data) carries the last log
          lines, which the Space's /status poller surfaces in the UI.
Optional: SAM2_SPLIT=1 swaps the dynamic/static split inside Instant4D's
          prune.py for the SAM2-video-mask split (see sam2_pipeline_hook.py);
          default is the existing Otsu/ceiling heuristic pipeline.py already
          patches in.
"""
from __future__ import annotations

import base64
import os
import tempfile

import runpod

import pipeline  # the 4dgs pipeline module baked into the image

if os.environ.get('SAM2_SPLIT', '').strip() in ('1', 'true', 'True'):
    import sam2_pipeline_hook
    sam2_pipeline_hook.install(pipeline, log=lambda line: print(line, flush=True))


def handler(job):
    inp = job.get('input') or {}
    if not inp.get('video_b64'):
        return {'error': 'no video_b64 in input'}

    tail: list[str] = []

    def log(line: str) -> None:
        print(line, flush=True)
        tail.append(line)
        if len(tail) > 200:
            tail.pop(0)
        runpod.serverless.progress_update(job, {'log_tail': tail[-8:]})

    with tempfile.TemporaryDirectory() as td:
        video = os.path.join(td, 'clip.mp4')
        with open(video, 'wb') as f:
            f.write(base64.b64decode(inp['video_b64']))
        out = pipeline.run_pipeline(
            video, 'clip',
            int(inp.get('max_frames', 96)), int(inp.get('target_fps', 24)),
            td, log)
        filename = os.path.basename(out)

        if os.environ.get('BUCKET_ENDPOINT_URL', '').strip():
            # RunPod's own uploader is opt-in code, not an automatic hook on
            # the return value: rp_upload.upload_file() reads the BUCKET_*
            # env vars and does the S3 PUT + presign itself.
            try:
                from runpod.serverless.utils import rp_upload
                url = rp_upload.upload_file_to_bucket(filename, out)
                log(f'uploaded to bucket -> {url}')
                return {'filename': filename, 'splatv_url': url}
            except Exception as exc:  # noqa: BLE001 - fall back to inline
                log(f'bucket upload failed ({exc}); falling back to inline base64')

        with open(out, 'rb') as f:
            data = f.read()
    if len(data) > 180 * 1024 * 1024:
        return {'error': f'.splatv is {len(data) / 1e6:.0f} MB, too large to inline; '
                         'configure BUCKET_ENDPOINT_URL/BUCKET_ACCESS_KEY_ID/'
                         'BUCKET_SECRET_ACCESS_KEY on the endpoint'}
    return {'filename': filename,
            'splatv_b64': base64.b64encode(data).decode()}


runpod.serverless.start({'handler': handler})