Spaces:
Running on Zero
Running on Zero
File size: 3,593 Bytes
b600b7f df5cde5 b600b7f df5cde5 b600b7f df5cde5 b600b7f df5cde5 b600b7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """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})
|