revert to clean eager app.py (remove all AOTI/dynamo machinery); 480p24 has ample headroom, eager is sufficient; align with local run_server

app.py

@@ -255,116 +255,6 @@ except Exception as e:  # noqa: BLE001
 # run_session_blocking (in xvideo/serving/zerogpu_engine.py) reuses the vendored
 # gate/PE/streaming logic; it is NOT a modification of the existing server.
 # ---------------------------------------------------------------------------
-_AOTI_STATE = {}
-
-
-def _fork_load_aoti() -> None:
-    """Inside a GPU fork: fetch (idempotent, /data-cached) and load the dev-box
-    precompiled AOTInductor VAE encode packages (same sm_120 arch + torch 2.9.1),
-    then monkeypatch the encode entrypoints. Executes via dlopen — zero dynamo,
-    which is inert in ZeroGPU forks. Fetch happens IN THE FORK so a download
-    problem can never wedge the app boot."""
-    if _AOTI_STATE.get("installed") or _RUNTIME is None:
-        return
-    from torch._inductor import aoti_load_package
-    aoti_dir = _DATA / "aoti"
-    aoti_dir.mkdir(parents=True, exist_ok=True)
-    RES = ((480, 840), (840, 480), (720, 1248), (1248, 720))
-    wanted = []
-    for H, W in RES:
-        for t in (9, 1):
-            wanted.append(("encode", t, H, W, f"vae_encode_t{t}_{H}x{W}.pt2"))
-        for t in (2, 1):
-            wanted.append(("decode", t, H, W, f"vae_decode_t{t}_{H}x{W}.pt2"))
-    fetch = bool(_AOTI_STATE.pop("fetch", False))
-    if fetch:
-        for kind, t, H, W, fname in wanted:
-            p = aoti_dir / fname
-            if p.is_file():
-                continue
-            try:
-                from huggingface_hub import hf_hub_download
-                hf_hub_download("wxDai/joyomni-aoti-sm120", fname, repo_type="dataset",
-                                local_dir=str(aoti_dir))
-                print(f"[fork] fetched aoti {fname}", flush=True)
-            except Exception:  # noqa: BLE001
-                _AOTI_STATE.setdefault("log", []).append(f"fetch_fail {fname}"[:120])
-
-    enc_files, dec_files = {}, {}
-    for kind, t, H, W, fname in wanted:
-        p = aoti_dir / fname
-        if not p.is_file():
-            continue
-        if kind == "encode":
-            enc_files[(t, H, W)] = str(p)
-        else:
-            dec_files[(t, H * 2 // 3 // 16, W * 2 // 3 // 16)] = str(p)
-    if not enc_files and not dec_files:
-        return
-    loaded = {}
-
-    def _get(files, key):
-        if key in loaded:
-            return loaded[key]
-        path = files.get(key)
-        if path is None:
-            loaded[key] = None
-            return None
-        try:
-            loaded[key] = aoti_load_package(path)
-            _AOTI_STATE.setdefault("log", []).append(f"lazy-loaded {path.rsplit('/', 1)[-1]}")
-            print(f"[fork] aoti lazy-loaded {path.rsplit('/', 1)[-1]}", flush=True)
-        except Exception as e:  # noqa: BLE001
-            loaded[key] = None
-            _AOTI_STATE.setdefault("log", []).append(f"load_fail {key}: {e!r}"[:180])
-        return loaded[key]
-
-    def _unwrap_fn(fn):
-        while hasattr(fn, "_torchdynamo_orig_callable"):
-            fn = fn._torchdynamo_orig_callable
-        return fn
-
-    seen = set()
-    for vae in (_RUNTIME.pipeline.vae, _RUNTIME.pseudo_encode_vae, _RUNTIME.decode_vae):
-        if vae is None or id(vae) in seen:
-            continue
-        seen.add(id(vae))
-        if enc_files and hasattr(vae, "_encode"):
-            fb_e = _unwrap_fn(vae._encode)
-
-            def _encode_dispatch(x, _fb=fb_e):
-                fn = None
-                if x.shape[0] == 1 and x.shape[1] == 3:
-                    fn = _get(enc_files, (int(x.shape[2]), int(x.shape[3]), int(x.shape[4])))
-                if fn is None:
-                    return _fb(x)
-                try:
-                    return fn(x)
-                except Exception:  # noqa: BLE001
-                    return _fb(x)
-
-            vae._encode = _encode_dispatch
-        if dec_files and hasattr(vae, "_decode"):
-            fb_d = _unwrap_fn(vae._decode)
-
-            def _decode_dispatch(z, _fb=fb_d):
-                fn = None
-                if z.shape[0] == 1:
-                    fn = _get(dec_files, (int(z.shape[2]), int(z.shape[3]), int(z.shape[4])))
-                if fn is None:
-                    return _fb(z)
-                try:
-                    return fn(z)
-                except Exception:  # noqa: BLE001
-                    return _fb(z)
-
-            vae._decode = _decode_dispatch
-    _AOTI_STATE["installed"] = True
-    _AOTI_STATE.setdefault("log", []).append(
-        f"registered enc={len(enc_files)} dec={len(dec_files)} (lazy)")
-    print(f"[fork] aoti lazy dispatch registered enc={len(enc_files)} dec={len(dec_files)}", flush=True)
-
-
 def _fork_enable_compile() -> None:
     """ZeroGPU's CUDA-emulation layer disables torch dynamo in the parent process;
     forks inherit that, so every torch.compile wrapper silently runs eager (measured:
@@ -413,13 +303,17 @@ def _fork_enable_compile() -> None:
                     setattr(_vae, _attr, _plain)
                     _unwrapped += _n
         globals()["_FORK_UNWRAPPED"] = _unwrapped
-        print(f"[fork] unwrapped {_unwrapped} poisoned compile wrappers; VAE runs eager+AOTI (dynamo-independent)", flush=True)
+        _vc._configured.clear()
+        _vc._configured_encode.clear()
+        _vc._configured_encode_dynamic.clear()
+        if _RUNTIME is not None:
+            _vc.maybe_setup_decode(_RUNTIME.decode_vae)
+            _vc.maybe_setup_encode(_RUNTIME.pipeline.vae)
+            if _RUNTIME.pseudo_encode_vae is not _RUNTIME.pipeline.vae:
+                _vc.maybe_setup_encode(_RUNTIME.pseudo_encode_vae)
+        print(f"[fork] dynamo re-enabled; unwrapped {_unwrapped} poisoned compile wrappers; VAE re-wrapped fresh", flush=True)
     except Exception as e:  # noqa: BLE001
         print(f"[fork] compile re-enable failed: {e!r}", flush=True)
-    try:
-        _fork_load_aoti()
-    except Exception as e:  # noqa: BLE001
-        print(f"[fork] aoti install failed: {e!r}", flush=True)
 
 
 def make_gpu_session(in_q, out_q, paths):
@@ -462,7 +356,6 @@ def _run_session_in_fork(in_q, out_q, paths):
 
 @spaces.GPU(duration=290, size="xlarge")
 def warmup():
-    _AOTI_STATE["fetch"] = True
     """Warm up the pipeline inside the fork (real GPU): both orientations at both
     resolution tiers, same coverage as the standalone server's load-time warmup.
     720p is included so its VAE compile artifacts land in /data — without this,
@@ -488,13 +381,6 @@ def gpu_probe():
 
     _fork_enable_compile()
     out = {"device": torch.cuda.get_device_name(0)}
-    try:
-        _adir = _DATA / "aoti"
-        out["aoti_state"] = {"installed": _AOTI_STATE.get("installed"),
-                             "log": _AOTI_STATE.get("log", [])[-6:],
-                             "files": sorted(p.name for p in _adir.glob("*.pt2")) if _adir.is_dir() else []}
-    except Exception as e:  # noqa: BLE001
-        out["aoti_state"] = {"err": repr(e)[:120]}
     try:
         import os as _os
         out["env_torch"] = {k: v for k, v in _os.environ.items()
@@ -647,8 +533,25 @@ def gpu_probe():
         out["census_err"] = repr(e)[:150]
 
     try:
+        import torch._dynamo.utils as _du
         from xvideo.models.vae import vae_compile as _vc
+        vae_d = _RUNTIME.decode_vae
         vae_e = _RUNTIME.pipeline.vae
+        tgt = getattr(vae_d, "_decode", None) or vae_d.decode
+        out["decode_wrap"] = repr(type(tgt))[:90]
+        out["decode_is_compiled"] = hasattr(tgt, "_torchdynamo_orig_callable") or "Optimized" in repr(type(tgt))
+        lat_c = int(getattr(vae_d, "latent_channels", 16) or 16)
+        z = torch.zeros(1, lat_c, 2, 60, 105, device=dev, dtype=torch.bfloat16)
+        z = _vc.prep_input(z)
+        with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
+            for _ in range(3):
+                vae_d.decode(z, return_dict=False)[0]
+            torch.cuda.synchronize()
+            t0 = _t.perf_counter()
+            for _ in range(15):
+                vae_d.decode(z, return_dict=False)[0]
+            torch.cuda.synchronize()
+        out["vae_decode_direct_ms"] = round((_t.perf_counter() - t0) / 15 * 1000, 2)
         x = torch.zeros(1, 3, 9, 480, 840, device=dev, dtype=torch.bfloat16)
         x = _vc.prep_input(x)
         with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
@@ -660,27 +563,15 @@ def gpu_probe():
                 vae_e.encode(x)
             torch.cuda.synchronize()
         out["vae_encode_direct_ms"] = round((_t.perf_counter() - t0) / 10 * 1000, 2)
+        fr = _du.counters.get("frames", {})
+        out["dynamo_frames"] = {"total": fr.get("total"), "ok": fr.get("ok")}
+        st = _du.counters.get("stats", {})
+        out["dynamo_unique_graphs"] = st.get("unique_graphs")
+        gb = _du.counters.get("graph_break", {})
+        out["dynamo_graph_breaks"] = sum(gb.values()) if gb else 0
     except Exception as e:  # noqa: BLE001
         out["vae_probe_err"] = repr(e)[:200]
 
-    try:
-        from xvideo.models.vae import vae_compile as _vc
-        vae_d = _RUNTIME.decode_vae
-        lat_c = int(getattr(vae_d, "latent_channels", 16) or 16)
-        z = torch.zeros(1, lat_c, 2, 60, 105, device=dev, dtype=torch.bfloat16)
-        z = _vc.prep_input(z)
-        with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
-            for _ in range(3):
-                vae_d.decode(z, return_dict=False)[0]
-            torch.cuda.synchronize()
-            t0 = _t.perf_counter()
-            for _ in range(10):
-                vae_d.decode(z, return_dict=False)[0]
-            torch.cuda.synchronize()
-        out["vae_decode_direct_ms"] = round((_t.perf_counter() - t0) / 10 * 1000, 2)
-    except Exception as e:  # noqa: BLE001
-        out["vae_decode_err"] = repr(e)[:200]
-
     try:
         import torch.nn.functional as F
         torch.backends.cudnn.benchmark = True
@@ -805,8 +696,7 @@ async def do_warmup():
 
 
 @spaces.GPU(duration=120, size="xlarge")
-def engine_bench(fresh_cache: bool = False, pace: float = 12.0, chunks: int = 16,
-                 height: int = 480, width: int = 840):
+def engine_bench(fresh_cache: bool = False, pace: float = 12.0, chunks: int = 16):
     """Engine-only throughput probe inside the fork: no session wrapper, no
     queues — the exact mirror of the dev-box bench. Separates 'fork environment'
     from 'session machinery' when chasing Space-vs-local speed gaps. Samples SM
@@ -843,13 +733,13 @@ def engine_bench(fresh_cache: bool = False, pace: float = 12.0, chunks: int = 16
 
     _th.Thread(target=_clk, daemon=True).start()
 
-    settings = StreamingSettings(height=height, width=width, num_inference_steps=2,
+    settings = StreamingSettings(height=480, width=840, num_inference_steps=2,
                                  seed=42, max_temporal_ids=8, profile_timings=True,
                                  output_codec="mjpeg")
     sess = _RUNTIME.create_v2v_session(
         "Turn the scene into an oil painting with warm colors.", settings=settings)
     rng = np.random.default_rng(7)
-    base = rng.integers(0, 256, size=(height, width, 3), dtype=np.uint8)
+    base = rng.integers(0, 256, size=(480, 840, 3), dtype=np.uint8)
     recs = []
     t0 = _t.time()
     i = 0
@@ -860,7 +750,7 @@ def engine_bench(fresh_cache: bool = False, pace: float = 12.0, chunks: int = 16
                 now = _t.time()
                 if target > now:
                     _t.sleep(target - now)
-            frame = _Image.fromarray(np.roll(base, (i * 7) % width, axis=1), mode="RGB")
+            frame = _Image.fromarray(np.roll(base, (i * 7) % 840, axis=1), mode="RGB")
             for r in sess.push_frame(frame, {"seq": i + 1, "t_capture_ms": 0.0}):
                 p = r.profile or {}
                 recs.append({k: round(float(p[k]), 4) for k in
@@ -888,7 +778,6 @@ def engine_bench(fresh_cache: bool = False, pace: float = 12.0, chunks: int = 16
         v = [c[k] for c in st if k in c]
         return round(sum(v) / len(v), 4) if v else None
     return {"chunks": len(recs), "fresh_cache": bool(fresh_cache), "pace": pace,
-            "aoti_log": _AOTI_STATE.get("log", [])[-8:],
             "dit": _m("dit_denoise_s"), "enc": _m("vae_encode_s"),
             "ref": _m("reference_prepare_s"), "dec": _m("vae_decode_s"),
             "jpg": _m("jpeg_encode_s"), "f2t": _m("frames_to_tensor_s"),
@@ -906,13 +795,11 @@ async def do_gpu_probe():
 
 
 @app.get("/engine-bench")
-async def do_engine_bench(fresh_cache: int = 0, pace: float = 12.0, chunks: int = 16,
-                          height: int = 480, width: int = 840):
+async def do_engine_bench(fresh_cache: int = 0, pace: float = 12.0, chunks: int = 16):
     loop = asyncio.get_event_loop()
     try:
         return {"result": await loop.run_in_executor(
-            None, lambda: engine_bench(bool(fresh_cache), float(pace), int(chunks),
-                                       int(height), int(width)))}
+            None, lambda: engine_bench(bool(fresh_cache), float(pace), int(chunks)))}
     except Exception as e:  # noqa: BLE001
         return JSONResponse({"error": str(e)[:300]}, status_code=500)