File size: 3,731 Bytes
b209cd2 | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | 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())
|