"""WASD-controlled push_cube world model (AR shortcut-forcing DiT-S, 5-step inference). HF Space edition: loads a stripped fp16 checkpoint + Wan-VAE from the repo, initial frames from frames/*.png. Same stdlib HTTP server + HTML page as the cluster demo. CPU or GPU (auto). """ import os, sys, io, json, time, threading, socket, base64, glob from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import numpy as np import torch import yaml ROOT = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, ROOT) # weights live in a separate public model repo (Space repo has a 1 GB cap) from huggingface_hub import hf_hub_download WREPO = "lamamkh/push-cube-wasd-worldmodel-weights" CKPT_PATH = hf_hub_download(WREPO, "model_fp16.pt") VAE_PATH = hf_hub_download(WREPO, "Wan2.1_VAE.pth") os.environ["WAN_VAE_PATH"] = VAE_PATH from acwm.model.interface import get_dynamics_class DEVICE = "cuda" if torch.cuda.is_available() else "cpu" if DEVICE == "cpu": torch.set_num_threads(max(1, os.cpu_count() or 2)) PORT = int(os.environ.get("PORT", "7860")) K_STEPS = 5 MAX_LATENT = 10 print(f"[wasd] device={DEVICE}", flush=True) cfg = yaml.safe_load(open(os.path.join(ROOT, "push_cube_ar.yaml"))) mc = cfg["model_config"]; mc["action_dim"] = 2 mc["use_flash_attn"] = torch.cuda.is_available() WM = get_dynamics_class(cfg["dynamics_class"])(cfg["model_name"], mc).to(DEVICE).eval() _ck = torch.load(CKPT_PATH, map_location="cpu", weights_only=False) _sd = {k: (v.float() if v.is_floating_point() else v) for k, v in _ck["model_state_dict"].items()} WM.load_state_dict(_sd, strict=False) WM = WM.float() print(f"[wasd] loaded step={_ck.get('step','?')}", flush=True) ACR = WM.model.action_compress_rate KMAX = WM.k_max FRAMES = sorted(glob.glob(os.path.join(ROOT, "frames", "*.png"))) def load_frame(idx): import imageio.v2 as iio img = iio.imread(FRAMES[idx % len(FRAMES)]).astype(np.float32) / 255.0 return torch.from_numpy(img) # [H,W,3] @torch.no_grad() def ar_one_frame(z_all, a_curr): h, w, D = z_all.shape[2], z_all.shape[3], z_all.shape[4] t_len = z_all.shape[1] + 1 K = K_STEPS; d = 1.0 / K; d_min = 1.0 / KMAX z = torch.randn(1, 1, h, w, D, device=DEVICE) for i in range(K): seq = torch.cat([z_all, z], dim=1) sig = torch.ones(1, t_len, device=DEVICE); sig[:, -1] = i / K ds = torch.full((1, t_len), d_min, device=DEVICE); ds[:, -1] = d b = WM.model(seq, sig, ds, a_curr) z = z + b[:, -1:] * d return z @torch.no_grad() def decode_last(z_all): vid = WM.vae.decode_to_pixel(z_all.permute(0, 1, 4, 2, 3).contiguous()) vid = ((vid + 1) / 2).clamp(0, 1)[0].permute(0, 2, 3, 1).cpu().numpy() return vid[-1] def png_b64(img01): import imageio.v2 as iio buf = io.BytesIO() iio.imwrite(buf, (np.clip(img01, 0, 1) * 255).astype(np.uint8), format="png") return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode() LOCK = threading.Lock() S = {"z": None, "acts": []} @torch.no_grad() def reset(idx): f0 = load_frame(int(idx)) o0 = f0.unsqueeze(0).to(DEVICE) # [1,H,W,3] S["z"] = WM.encode_obs(o0.unsqueeze(1)); S["acts"] = [] return png_b64(f0.numpy()), 1 @torch.no_grad() def step(dx, dy): v = np.array([dx, dy], np.float32); n = float(np.linalg.norm(v)) a = v / n if n > 1e-6 else np.array([1.0, 0.0], np.float32) S["acts"].extend([a] * ((ACR + 1) if len(S["acts"]) == 0 else ACR)) a_curr = torch.tensor(np.stack(S["acts"]), device=DEVICE).unsqueeze(0).float() S["z"] = torch.cat([S["z"], ar_one_frame(S["z"], a_curr)], dim=1) if S["z"].shape[1] > MAX_LATENT: S["z"] = S["z"][:, 1:]; S["acts"] = S["acts"][ACR:] return png_b64(decode_last(S["z"])), S["z"].shape[1], [float(a[0]), float(a[1])] PAGE = """push_cube · WASD world model
world model frame
press R or Reset to startthinking…

Drive the world model

WASD (or tap the keys) pushes the disk. Each press imagines ~4 frames with 5-step denoising in an AR shortcut DiT — everything is generated, no simulator. On the free CPU tier a step takes a while; on GPU it's ~0.3 s.

W
A
S
D

""" class H(BaseHTTPRequestHandler): def log_message(self, *a): pass def _json(self, obj, code=200): b = json.dumps(obj).encode() self.send_response(code); self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(b))); self.end_headers(); self.wfile.write(b) def do_GET(self): b = PAGE.encode() self.send_response(200); self.send_header("Content-Type", "text/html; charset=utf-8") self.send_header("Content-Length", str(len(b))); self.end_headers(); self.wfile.write(b) def do_POST(self): n = int(self.headers.get("Content-Length", 0)) req = json.loads(self.rfile.read(n) or b"{}") try: with LOCK: if self.path.endswith("/reset"): frame, nlat = reset(req.get("idx", 0)); self._json({"frame": frame, "n": nlat}) elif self.path.endswith("/step"): if S["z"] is None: self._json({"error": "reset first"}, 400); return frame, nlat, a = step(float(req.get("dx", 0)), float(req.get("dy", 0))) self._json({"frame": frame, "n": nlat, "a": a}) else: self._json({"error": "unknown"}, 404) except Exception as e: import traceback; traceback.print_exc() self._json({"error": str(e)}, 500) if __name__ == "__main__": print(f"[wasd] serving :{PORT}", flush=True) ThreadingHTTPServer(("0.0.0.0", PORT), H).serve_forever()