| import modal |
|
|
| app = modal.App("surflo-ctrl") |
| VOL = modal.Volume.from_name("surflo-vol", create_if_missing=True) |
|
|
| img = ( |
| modal.Image.from_registry("nvidia/cuda:12.4.1-devel-ubuntu22.04", add_python="3.10") |
| .apt_install("software-properties-common", "curl", "git", "build-essential", |
| "cmake", "ninja-build") |
| .run_commands( |
| "add-apt-repository -y ppa:deadsnakes/ppa", |
| "apt-get update -y", |
| "apt-get install -y python3.10 python3.10-venv python3.10-dev python3.10-distutils", |
| ) |
| ) |
|
|
| PAYLOAD = r''' |
| import os, subprocess, sys, time |
| |
| VENV = "/vol/venv" |
| PY = VENV + "/bin/python" |
| REPO = "/vol/Surflo" |
| CKPT = "/vol/checkpoints/surflo_v0.pt" |
| SAMPLE = REPO + "/media/sample" |
| log = [] |
| |
| |
| def sh(cmd, label, timeout=7200): |
| t0 = time.time() |
| env = dict(os.environ) |
| env.pop("PYTHONPATH", None) |
| try: |
| p = subprocess.run(["bash", "-lc", cmd], capture_output=True, text=True, |
| timeout=timeout, env=env, cwd="/") |
| except subprocess.TimeoutExpired: |
| log.append("=== %s TIMEOUT" % label) |
| return 1 |
| dt = time.time() - t0 |
| log.append("=== %s rc=%d %.0fs\n%s" % (label, p.returncode, dt, |
| ((p.stdout or "") + (p.stderr or ""))[-3000:])) |
| print("=== %s rc=%d %.0fs" % (label, p.returncode, dt), flush=True) |
| return p.returncode |
| |
| |
| def save(): |
| open("/vol/ctrl_run.txt", "w").write("\n".join(log)) |
| |
| |
| # --- wait for the build to have produced a working stack ------------------- |
| READY = """ |
| import importlib, sys |
| need = ['torch','diff_gaussian_rasterization_surflo','fused_ssim','geodel','nvdiffrast.torch'] |
| bad = [] |
| for m in need: |
| try: |
| importlib.import_module(m) |
| except Exception as e: |
| bad.append(m) |
| print('READY' if not bad else 'NOT_READY ' + ','.join(bad)) |
| """ |
| for attempt in range(60): |
| if os.path.exists("/vol/build310.txt"): |
| log.append("build310.txt present") |
| break |
| log.append("wait_build %d" % attempt) |
| print("wait_build %d" % attempt, flush=True) |
| time.sleep(60) |
| |
| for attempt in range(20): |
| r = subprocess.run([PY, "-c", READY], capture_output=True, text=True) |
| out = (r.stdout or "") + (r.stderr or "") |
| log.append("wait %d: %s" % (attempt, out.strip()[:200])) |
| print("wait %d: %s" % (attempt, out.strip()[:120]), flush=True) |
| if "READY" in out and "NOT_READY" not in out: |
| break |
| time.sleep(60) |
| else: |
| log.append("!! stack never became ready") |
| save() |
| sys.exit("not ready") |
| |
| # --- VALID CONTROL: guided on Surflo's own real photos --------------------- |
| n = len([f for f in os.listdir(SAMPLE) if f.lower().endswith((".jpg", ".jpeg", ".png"))]) |
| log.append("sample images: %d" % n) |
| |
| OUT = "/vol/outputs/ctrl_guided" |
| sh("cd %s && %s scripts/infer.py mode=guided guided=default ckpt=%s " |
| "source.image_folder=%s source.n_images=16 num_query_points=100000 output_dir=%s" |
| % (REPO, PY, CKPT, SAMPLE, OUT), "ctrl_guided_default", timeout=7200) |
| |
| listing = "" |
| for root, _d, files in os.walk(OUT): |
| for f in sorted(files): |
| listing += " %s %s\n" % (os.path.join(root, f), |
| "{:,}".format(os.path.getsize(os.path.join(root, f)))) |
| log.append("--- outputs ---\n" + listing) |
| save() |
| ''' |
|
|
| import subprocess |
|
|
|
|
| @app.function(timeout=28800, volumes={"/vol": VOL}, image=img, gpu="L4") |
| def run_ctrl() -> str: |
| import subprocess as sp |
| with open("/vol/_ctrl.py", "w") as fh: |
| fh.write(PAYLOAD) |
| p = sp.run(["bash", "-lc", "/usr/bin/python3.10 /vol/_ctrl.py"], |
| capture_output=True, text=True, timeout=28800) |
| VOL.commit() |
| return ((p.stdout or "") + (p.stderr or ""))[-4000:] |
|
|
|
|
| @app.local_entrypoint() |
| def main(): |
| print(run_ctrl.remote()) |
|
|
