Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/back_01.png filter=lfs diff=lfs merge=lfs -text
+examples/back_264.png filter=lfs diff=lfs merge=lfs -text
+examples/bamboo_L.png filter=lfs diff=lfs merge=lfs -text
+examples/watercolor_L.png filter=lfs diff=lfs merge=lfs -text

LICENSE_ACTIONROPE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2026 The ActionRoPE Authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,13 +1,62 @@
 ---
-title: Actionrope
-emoji: 📚
-colorFrom: gray
+title: ActionRoPE
+emoji: 🎮
+colorFrom: indigo
 colorTo: red
 sdk: gradio
 sdk_version: 6.30.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Action-controlled game world model video generation
+python_version: "3.12"
+startup_duration_timeout: 1h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# ActionRoPE — action as coordinates for planar-camera game world models
+
+Demo of [**teawhite/ActionRoPE**](https://huggingface.co/teawhite/ActionRoPE) —
+*ActionRoPE: action as coordinates for spatially consistent control in planar-camera
+game world models* ([paper 2610.25861](https://huggingface.co/papers/2610.25861),
+[GitHub](https://github.com/TeaWhiteBro/ActionRoPE)).
+
+ActionRoPE turns a video-diffusion world model into a **controllable** one without
+touching its weights: the camera displacement of each frame is written directly into
+the rotary positions of the DiT's tokens (world coordinates), so tokens that look at
+the same world location share the same positional address across time. Two
+companions make it work: **CNS** (coordinate-conditioned noise scheduling), which
+noises each latent cell according to how "known" its world region is, and a tiny
+zero-initialized **known/new input channel**.
+
+## How to drive it
+
+- **First frame** — any image; it is centre-cropped to 26:15 and resized to 832×480
+  (the paper's resolution).
+- **Actions** — a comma-separated plan over the 80 frames after the first one, each
+  segment `<direction>[*speed]:<frames>`:
+  - directions: `idle`, `up`, `down`, `left`, `right`, `up-left`, `up-right`,
+    `down-left`, `down-right` (aliases `u/d/l/r/ul/ur/dl/dr`)
+  - speed `1.0` is the game's nominal walking speed: **69.2 px/s horizontally,
+    46.0 px/s vertically** at 832×480; `*1.5` walks 50% faster.
+  - e.g. `right:40,left:40` (walk right 2.5 s, then back), or
+    `up-right*1.5:32,left:40,down-left:8`.
+- The output is an **81-frame / 5 s / 16 fps / 832×480** clip.
+
+Advanced: seed, sampling steps (paper: 50), and the CNS on/off switch.
+
+## Notes
+
+- Runs on ZeroGPU; the Wan2.2-TI2V-5B backbone (5 B DiT, ~10 GB fp16) and the
+  Wan2.2 VAE are downloaded once at startup (~13 GB total). The ActionRoPE
+  checkpoint is a *full* DiT state dict, so the base model's shards are never
+  fetched.
+- Sampling uses the paper's evaluation path verbatim: 50 Euler steps of the shift-5
+  flow-matching schedule, no guidance, the shipped NULL-prompt UMT5 embedding, VAE
+  run in fp32.
+
+## Credits
+
+- Code: [TeaWhiteBro/ActionRoPE](https://github.com/TeaWhiteBro/ActionRoPE)
+  (Apache-2.0), vendored under `actionrope/` with only the training modules removed.
+- Weights: [teawhite/ActionRoPE](https://huggingface.co/teawhite/ActionRoPE)
+  (Wan2.2-TI2V-5B backbone: Apache-2.0) and [Wan-AI/Wan2.2-TI2V-5B](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B) (VAE).
+- The example first frames in `examples/` are extracted from the authors' own
+  released videos on the paper's project page.

actionrope/__init__.py

@@ -0,0 +1,2 @@
+"""ActionRoPE: Action as Coordinates for Spatially Consistent Control in Planar-Camera Game World Models."""
+__version__ = "1.0.0"

actionrope/checkpoint.py

@@ -0,0 +1,44 @@
+"""safetensors checkpoints with a small, uniform metadata header.
+
+Every ActionRoPE checkpoint stores (as strings):
+
+    model            "wan2.2-ti2v-5b" | "skyreels-v2-df-1.3b" | "ltx-2.5-22b" | "minimax-h3-33b"
+    format           "full" (complete DiT state dict) or "lora" (LoRA factors + new modules only)
+    mask_channel     "true" / "false"   known/new input channel
+    cns              "true" / "false"   Coordinate-Conditioned Noise Scheduling
+    cns_shift        base shift S of the sampling schedule
+    cns_shift_known  shift S_known of known tokens
+    cns_alpha        exponent alpha (1 in the paper)
+    step             training step
+    lora_rank, lora_alpha, lora_targets   (format = "lora" only)
+"""
+from __future__ import annotations
+
+import json
+import os
+
+import torch
+
+
+def save_checkpoint(path: str, state_dict: dict[str, torch.Tensor], metadata: dict) -> None:
+    from safetensors.torch import save_file
+    meta = {k: (json.dumps(v) if isinstance(v, (list, dict)) else str(v).lower() if isinstance(v, bool) else str(v))
+            for k, v in metadata.items()}
+    save_file({k: v.detach().contiguous().cpu() for k, v in state_dict.items()}, path, metadata=meta)
+    os.chmod(path, 0o644)      # safetensors follows the umask, which may make the file private
+
+
+def read_metadata(path: str) -> dict[str, str]:
+    from safetensors import safe_open
+    with safe_open(path, "pt") as f:
+        return dict(f.metadata() or {})
+
+
+def meta_bool(meta: dict, key: str, default: bool = False) -> bool:
+    v = meta.get(key)
+    return default if v is None else str(v).strip().lower() in ("1", "true", "yes")
+
+
+def meta_float(meta: dict, key: str, default: float | None = None) -> float | None:
+    v = meta.get(key)
+    return default if v is None or v == "" else float(v)

actionrope/cns.py

@@ -0,0 +1,24 @@
+"""Coordinate-Conditioned Noise Scheduling (CNS).
+
+Every backbone samples with a shifted flow-matching schedule sigma = S u / (1 + (S - 1) u), u in [0, 1]. CNS gives each
+token its own shift according to its knownness kappa (1 = inside the first-frame footprint, 0 = newly revealed):
+
+    u       = sigma / (S - (S - 1) sigma)                 (invert the base schedule)
+    S_tok   = S * (S_known / S) ** (alpha * kappa)
+    sigma_t = S_tok u / (1 + (S_tok - 1) u)
+
+Known tokens follow a schedule with the smaller shift S_known, so they reach low noise earlier during sampling;
+newly revealed tokens keep the base schedule (sigma_t = sigma for kappa = 0). sigma_t(1) = 1 and sigma_t(0) = 0 for
+every kappa. The same function is used to noise the training samples and to integrate the sampler.
+"""
+from __future__ import annotations
+
+import torch
+
+
+def cns_sigma(sigma, kappa: torch.Tensor, shift: float, shift_known: float, alpha: float = 1.0) -> torch.Tensor:
+    """Per-token noise level for the global noise level ``sigma`` (float or tensor) and knownness ``kappa``."""
+    S = float(shift)
+    u = sigma / (S - (S - 1.0) * sigma)
+    S_tok = S * (float(shift_known) / S) ** (alpha * kappa)
+    return S_tok * u / (1.0 + (S_tok - 1.0) * u)

actionrope/coords.py

@@ -0,0 +1,41 @@
+"""Known / newly revealed regions in world coordinates (backbone-agnostic).
+
+A latent (or token) grid cell (i, j) of latent frame k looks at the world position (i + dy_k / p, j + dx_k / p), where
+(dx_k, dy_k) is the camera offset of that latent frame in pixels and p the number of pixels per grid cell. The cell is
+*known* when this position lies inside the footprint of the first frame, [-0.5, h - 0.5] x [-0.5, w - 0.5] (closed),
+and *newly revealed* otherwise. All functions here are pure geometry and work on any grid.
+"""
+from __future__ import annotations
+
+import torch
+
+
+def footprint_distance(offset_px: torch.Tensor, h: int, w: int, px_per_cell: float) -> torch.Tensor:
+    """Distance (in grid cells) of every cell's world position to the first-frame footprint; 0 inside.
+
+    offset_px: [b, f, 2] camera offset (dx, dy) of each latent frame in pixels. Returns float64 [b, f, h, w].
+    """
+    off = offset_px.to(torch.float64) / px_per_cell
+    b, f = off.shape[:2]
+    dx, dy = off[..., 0].view(b, f, 1, 1), off[..., 1].view(b, f, 1, 1)
+    wi = torch.arange(h, device=off.device, dtype=torch.float64).view(1, 1, h, 1) + dy
+    wj = torch.arange(w, device=off.device, dtype=torch.float64).view(1, 1, 1, w) + dx
+    dh = torch.clamp(torch.maximum(-0.5 - wi, wi - (h - 0.5)), min=0.0)
+    dw = torch.clamp(torch.maximum(-0.5 - wj, wj - (w - 0.5)), min=0.0)
+    return torch.sqrt(dh * dh + dw * dw)
+
+
+def known_mask(offset_px: torch.Tensor, h: int, w: int, px_per_cell: float) -> torch.Tensor:
+    """Boolean [b, f, h, w]: True where the cell's world position lies inside the first-frame footprint."""
+    return footprint_distance(offset_px, h, w, px_per_cell) <= 0
+
+
+def knownness(offset_px: torch.Tensor, h: int, w: int, px_per_cell: float) -> torch.Tensor:
+    """kappa as float32 [b, 1, f, h, w]: 1 for known cells, 0 for newly revealed cells."""
+    return known_mask(offset_px, h, w, px_per_cell).to(torch.float32).unsqueeze(1)
+
+
+def loss_weight_map(offset_px: torch.Tensor, h: int, w: int, px_per_cell: float, new_weight: float = 2.0) -> torch.Tensor:
+    """Per-cell loss weight, float32 [b, 1, f, h, w]: 1 for known cells, ``new_weight`` for newly revealed cells."""
+    known = known_mask(offset_px, h, w, px_per_cell)
+    return torch.where(known, 1.0, float(new_weight)).to(torch.float32).unsqueeze(1)

actionrope/data.py

@@ -0,0 +1,105 @@
+"""Dataset index shared by all backbones.
+
+The dataset is described by one JSON-lines annotation file, one row per 81-frame clip:
+
+    {"clip": "000123",                      # clip id = file stem of the video and of every latent file
+     "video": "clips/000123.mp4",           # path relative to the annotation file
+     "split": "train",                      # "train" or "val"
+     "kind": "plain",                       # "plain", or "burn" / "hardcut" for scene-transition clips
+     "valid": true,                         # false if the camera track of the clip is unreliable
+     "rt_ratio": 0.98,                      # real-time ratio of the recording (plain clips)
+     "sources": ["000045", "000981"],       # transition clips: the two clips they were composed from
+     "frame_offset_px": [[0.0, 0.0], ...],  # 81 x (dx, dy) camera offsets in px (see actionrope.geometry)
+     "prompts": ["In a 2.5D top down view, the player is in ...", ...]}   # 21 per-cell scene prompts (no actions)
+
+The 21 prompts follow the Wan-family latent frames (``geometry.cell_frames``); they only differ inside transition
+clips. Latents are pre-computed per backbone with ``scripts/encode_latents.py`` into ``<latent_dir>/<clip>.pt``.
+"""
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+TRANSITION_KINDS = ("burn", "hardcut")
+
+
+@dataclass
+class Clip:
+    clip: str
+    split: str
+    kind: str
+    valid: bool
+    rt_ratio: float | None
+    sources: tuple[str, ...]
+    frame_offset_px: np.ndarray       # float64 [81, 2]
+    prompts: list[str]
+    video: str | None = None
+
+
+def load_annotations(path: str) -> list[Clip]:
+    root = os.path.dirname(os.path.abspath(path))
+    clips = []
+    with open(path, encoding="utf-8") as fh:
+        for line in fh:
+            if not line.strip():
+                continue
+            r = json.loads(line)
+            clips.append(Clip(
+                clip=r["clip"], split=r["split"], kind=r.get("kind", "plain"), valid=bool(r.get("valid", True)),
+                rt_ratio=r.get("rt_ratio"), sources=tuple(r.get("sources") or ()),
+                frame_offset_px=np.asarray(r["frame_offset_px"], np.float64), prompts=list(r["prompts"]),
+                video=os.path.join(root, r["video"]) if r.get("video") else None))
+    return clips
+
+
+def training_clips(clips: list[Clip], min_rt_ratio: float = 0.8, exclude_val_sources: bool = True,
+                   verbose: bool = True) -> list[Clip]:
+    """Training split after the filters of the paper:
+
+    * drop clips whose camera track is invalid;
+    * drop transition clips composed from a validation clip (their frames would leak into training);
+    * drop plain clips recorded below ``min_rt_ratio`` of real time (the game stuttered, so motion per frame drops).
+    """
+    val = {c.clip for c in clips if c.split == "val"}
+    out, n_invalid, n_leak, n_slow = [], 0, 0, 0
+    for c in clips:
+        if c.split != "train":
+            continue
+        if not c.valid:
+            n_invalid += 1
+        elif c.kind in TRANSITION_KINDS:
+            if exclude_val_sources and any(s in val for s in c.sources):
+                n_leak += 1
+            else:
+                out.append(c)
+        elif min_rt_ratio > 0 and float(1.0 if c.rt_ratio is None else c.rt_ratio) < min_rt_ratio:
+            n_slow += 1
+        else:
+            out.append(c)
+    if verbose:
+        n_train = sum(c.split == "train" for c in clips)
+        print(f"[data] train clips: {n_train} - invalid {n_invalid} - leaking transitions {n_leak} - "
+              f"rt_ratio < {min_rt_ratio} {n_slow} = {len(out)}", flush=True)
+    return out
+
+
+def validation_clips(clips: list[Clip], min_rt_ratio: float = 0.9, n: int | None = 16, seed: int = 0) -> list[Clip]:
+    """Validation clips used to monitor the training loss (``n`` random clips with rt_ratio >= ``min_rt_ratio``)."""
+    import random
+    pool = [c for c in clips if c.split == "val" and c.valid and c.kind not in TRANSITION_KINDS
+            and float(1.0 if c.rt_ratio is None else c.rt_ratio) >= min_rt_ratio]
+    pool.sort(key=lambda c: c.clip)
+    if n is not None and n < len(pool):
+        random.Random(seed).shuffle(pool)
+        pool = sorted(pool[:n], key=lambda c: c.clip)
+    return pool
+
+
+def load_text_table(path: str) -> dict[str, torch.Tensor]:
+    """{prompt: embedding} written by ``scripts/build_text_table.py``."""
+    d = torch.load(path, map_location="cpu", weights_only=False)
+    return d["table"] if isinstance(d, dict) and "table" in d and isinstance(d["table"], dict) else d

actionrope/geometry.py

@@ -0,0 +1,98 @@
+"""Clip geometry shared by all backbones.
+
+Conventions
+-----------
+* A clip has 81 frames at 16 fps and 832x480 pixels (5 s).
+* Screen axes: x to the right, y downwards, in pixels of the 832x480 frame.
+* ``frame_offset_px[i] = (dx, dy)`` is the on-screen displacement of the player (= of the camera, which follows the
+  player) at frame i relative to frame 0, so ``frame_offset_px[0] == (0, 0)``. The static background moves by
+  ``-frame_offset_px``.
+* Every backbone maps the 81 per-frame offsets to its own latent frames (see ``frames_to_cells`` for the Wan-family
+  causal VAE; LTX-2.5 and MiniMax-H3 have their own grouping in their model files).
+
+Actions
+-------
+In a planar-camera game an action moves the camera by a known amount. We describe an action sequence as segments
+``<direction>[*<speed>]:<n_frames>``, e.g. ``"right:40,left:40"`` (walk right for 2.5 s, then back). ``speed`` is a
+multiplier of the nominal walking speed of the training data. Segments start at frame 1; frames not covered by any
+segment are idle.
+"""
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+FRAMES = 81
+FPS = 16
+WIDTH, HEIGHT = 832, 480
+
+# Nominal walking speed of the player in the training data, in px/s at 832x480 and 16 fps (x, y). It is the median
+# speed over straight-walk training clips; the evaluation commands are multiples of the horizontal value.
+WALK_SPEED_PX_S = (69.24106392992304, 46.012630371093756)
+
+_S = math.sqrt(0.5)
+DIRECTIONS = {
+    "idle": (0.0, 0.0),
+    "up": (0.0, -1.0),
+    "up-left": (-_S, -_S),
+    "left": (-1.0, 0.0),
+    "down-left": (-_S, _S),
+    "down": (0.0, 1.0),
+    "down-right": (_S, _S),
+    "right": (1.0, 0.0),
+    "up-right": (_S, -_S),
+}
+_ALIASES = {"u": "up", "d": "down", "l": "left", "r": "right", "ul": "up-left", "ur": "up-right",
+            "dl": "down-left", "dr": "down-right", "stop": "idle", "none": "idle"}
+
+
+def parse_actions(spec: str, n_frames: int = FRAMES) -> np.ndarray:
+    """Action string -> per-frame velocity [n_frames, 2] in px/frame (frame 0 has no velocity).
+
+    ``spec`` is a comma-separated list of ``<direction>[*<speed>]:<n_frames>`` segments, e.g. ``"right:40,left:40"``
+    or ``"up-right*1.5:24,idle:8,left:48"``. Directions: idle, up, down, left, right, up-left, up-right, down-left,
+    down-right (or u, d, l, r, ul, ur, dl, dr). The segments may cover at most ``n_frames - 1`` frames.
+    """
+    vel = np.zeros((n_frames, 2), np.float64)
+    frame = 1
+    for seg in filter(None, (s.strip() for s in spec.split(","))):
+        try:
+            head, count = seg.rsplit(":", 1)
+            name, _, speed = head.partition("*")
+            name = _ALIASES.get(name.strip().lower(), name.strip().lower())
+            d = np.array(DIRECTIONS[name], np.float64)
+            mult, count = float(speed) if speed else 1.0, int(count)
+        except (KeyError, ValueError) as e:
+            raise ValueError(f"cannot parse action segment {seg!r} (expected <direction>[*<speed>]:<n_frames>)") from e
+        if count < 0 or frame + count > n_frames:
+            raise ValueError(f"action segments cover more than {n_frames - 1} frames: {spec!r}")
+        vel[frame:frame + count] = d * np.array(WALK_SPEED_PX_S) * mult / FPS
+        frame += count
+    return vel
+
+
+def offsets_from_velocity(vel: np.ndarray) -> np.ndarray:
+    """Per-frame velocity [F, 2] (px/frame) -> per-frame camera offset [F, 2] with offset[0] = 0."""
+    vel = np.asarray(vel, np.float64).copy()
+    vel[0] = 0.0
+    return np.cumsum(vel, axis=0)
+
+
+def offsets_from_actions(spec: str, n_frames: int = FRAMES) -> np.ndarray:
+    """Action string -> per-frame camera offset ``frame_offset_px`` [n_frames, 2]."""
+    return offsets_from_velocity(parse_actions(spec, n_frames))
+
+
+def cell_frames(k: int) -> list[int]:
+    """Frames covered by latent frame (cell) k of the Wan-family causal VAE: cell 0 = frame 0, cell k = frames
+    4k-3 .. 4k."""
+    return [0] if k == 0 else list(range(4 * k - 3, 4 * k + 1))
+
+
+def frames_to_cells(frame_offset_px: np.ndarray) -> np.ndarray:
+    """Per-frame offsets [81, 2] -> per-cell offsets [21, 2]: the mean over the frames of each cell."""
+    off = np.asarray(frame_offset_px, np.float64)
+    if off.shape != (FRAMES, 2):
+        raise ValueError(f"frame_offset_px must have shape ({FRAMES}, 2), got {off.shape}")
+    return np.stack([off[cell_frames(k)].mean(0) for k in range((FRAMES - 1) // 4 + 1)], 0)

actionrope/media.py

@@ -0,0 +1,51 @@
+"""Image / video input and output."""
+from __future__ import annotations
+
+import numpy as np
+
+from .geometry import FPS, HEIGHT, WIDTH
+
+
+def read_frames(path: str, n: int | None = None) -> np.ndarray:
+    """Decode a video into uint8 RGB frames [F, H, W, 3] (PyAV)."""
+    import av
+    out = []
+    with av.open(path) as container:
+        for frame in container.decode(video=0):
+            out.append(frame.to_ndarray(format="rgb24"))
+            if n is not None and len(out) >= n:
+                break
+    return np.stack(out, 0)
+
+
+def load_first_frame(path: str) -> np.ndarray:
+    """Image (or the first frame of a video) -> uint8 [480, 832, 3].
+
+    Inputs of another size are centre-cropped to the 26:15 aspect ratio and resized (bicubic), so the pixel scale of
+    the scene - and therefore the walking speed in pixels - is not distorted."""
+    from PIL import Image
+    if path.lower().endswith((".mp4", ".mov", ".mkv", ".webm", ".avi")):
+        img = read_frames(path, 1)[0]
+    else:
+        img = np.array(Image.open(path).convert("RGB"))
+    if img.shape[:2] != (HEIGHT, WIDTH):
+        h, w = img.shape[:2]
+        tw, th = (w, round(w * HEIGHT / WIDTH)) if w * HEIGHT <= h * WIDTH else (round(h * WIDTH / HEIGHT), h)
+        x0, y0 = (w - tw) // 2, (h - th) // 2
+        img = np.array(Image.fromarray(img[y0:y0 + th, x0:x0 + tw]).resize((WIDTH, HEIGHT), Image.BICUBIC))
+    return img
+
+
+def write_video(path: str, frames, fps: float = FPS, crf: int | None = 16, quality: float | None = None) -> None:
+    """Write uint8 RGB frames with libx264. ``quality`` selects imageio's quality scale instead of a CRF (the Wan2.2
+    evaluation videos of the paper were written with quality=9, the other backbones with crf=16)."""
+    import imageio
+    frames = [np.asarray(f) for f in frames]
+    if quality is not None:
+        writer = imageio.get_writer(path, fps=fps, quality=quality)
+        for f in frames:
+            writer.append_data(f)
+        writer.close()
+    else:
+        import imageio.v3 as iio
+        iio.imwrite(path, np.stack(frames, 0), fps=fps, codec="libx264", output_params=["-crf", str(int(crf))])

actionrope/models/__init__.py

@@ -0,0 +1,2 @@
+"""Backbone integrations (Space build: wan only)."""
+MODELS = ("wan",)

actionrope/models/umt5.py

@@ -0,0 +1,47 @@
+"""UMT5-XXL prompt encoder of the Wan family (Wan2.2-TI2V-5B and SkyReels-V2 use the same encoder).
+
+Prompts are tokenised to 512 tokens and the embeddings after the true sequence length are zeroed, exactly as
+DiffSynth's ``WanVideoUnit_PromptEmbedder`` does. Training reads pre-computed embeddings from a text table
+(``scripts/build_text_table.py``) so the 11 GB encoder does not have to be loaded.
+"""
+from __future__ import annotations
+
+import os
+
+import torch
+
+T5_WEIGHTS = "models_t5_umt5-xxl-enc-bf16.pth"
+T5_TOKENIZER = os.path.join("google", "umt5-xxl")
+
+
+class UMT5Encoder:
+    def __init__(self, model_dir: str, device: str = "cuda"):
+        from diffsynth.pipelines.wan_video import ModelConfig, WanVideoPipeline
+        self.device = device
+        self.pipe = WanVideoPipeline.from_pretrained(
+            torch_dtype=torch.bfloat16, device=device, redirect_common_files=False,
+            model_configs=[ModelConfig(path=os.path.join(model_dir, T5_WEIGHTS))],
+            tokenizer_config=ModelConfig(path=os.path.join(model_dir, T5_TOKENIZER)))
+        self.pipe.load_models_to_device(["text_encoder"])
+
+    @torch.no_grad()
+    def __call__(self, prompts: list[str]) -> list[torch.Tensor]:
+        """-> one bf16 [512, 4096] CPU tensor per prompt."""
+        ids, mask = self.pipe.tokenizer(prompts, return_mask=True, add_special_tokens=True)
+        ids, mask = ids.to(self.device), mask.to(self.device)
+        seq_lens = mask.gt(0).sum(dim=1).long()
+        emb = self.pipe.text_encoder(ids, mask)
+        for j, n in enumerate(seq_lens):
+            emb[j, n:] = 0
+        return [e.to(torch.bfloat16).cpu().clone() for e in emb]
+
+
+def build_text_table(prompts: list[str], model_dir: str, device: str = "cuda", batch_size: int = 32) -> dict:
+    """{prompt: bf16 [512, 4096]} for the sorted, de-duplicated prompts (batches of ``batch_size``)."""
+    enc = UMT5Encoder(model_dir, device)
+    keys = sorted(set(prompts))
+    table = {}
+    for i in range(0, len(keys), batch_size):
+        chunk = keys[i:i + batch_size]
+        table.update(zip(chunk, enc(chunk)))
+    return table

actionrope/models/wan.py

@@ -0,0 +1,393 @@
+"""Wan2.2-TI2V-5B with ActionRoPE + CNS (the main backbone of the paper).
+
+Geometry: the Wan2.2 VAE compresses 16x in space and 4x in time (latent frame 0 = video frame 0, latent frame k =
+frames 4k-3 .. 4k), so an 81-frame 832x480 clip becomes 21 latent frames of 48 x 30 x 52. The DiT patch is (1, 2, 2):
+one token covers 32 x 32 pixels, one latent cell 16 x 16 pixels.
+
+* Camera offsets. The offset of a latent frame is the mean of the per-frame offsets of the frames it covers
+  (``geometry.frames_to_cells``); the token positions are shifted by offset / 32 (ActionRoPE, Sec. 3.1), the known/new
+  regions are evaluated on the 30 x 52 latent grid (offset / 16).
+* Known/new input channel. A zero-initialised Conv3d(1, 3072) adds the known map to the patch embedding (15,360
+  parameters).
+* Training (Sec. 4.1). Full fine-tuning of the DiT with the flow-matching loss of Wan2.2: one timestep index per
+  sample from the 1000-step training schedule (shift 5), loss weighted by DiffSynth's Wan timestep weighting. Latent
+  frame 0 is the clean first frame (timestep 0, no loss); newly revealed latent cells have loss weight ``new_weight``.
+  With CNS every latent cell is noised to its own level sigma_tok(sigma, kappa) and the DiT receives 1000 sigma_tok
+  as per-token timestep (kappa average-pooled to the 2 x 2 token grid). Each latent frame cross-attends to its own
+  scene prompt (prompts change inside scene-transition clips); with probability ``scene_dropout`` the whole clip
+  uses NULL_PROMPT.
+* Sampling. 50 Euler steps of the shift-5 schedule, no classifier-free guidance, NULL_PROMPT. Latent frame 0 is
+  written back to the encoded first frame after every step. With CNS each latent cell is integrated along its own
+  sigma_tok axis.
+"""
+from __future__ import annotations
+
+import os
+import random
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from ..checkpoint import meta_bool, meta_float, read_metadata
+from ..cns import cns_sigma
+from ..coords import knownness, loss_weight_map
+from ..data import load_annotations, load_text_table, training_clips, validation_clips
+from ..geometry import FPS, frames_to_cells
+from ..prompts import NULL_PROMPT
+from .wan_dit import MASK_EMBEDDING, arope_forward, install_mask_channel
+
+MODEL_NAME = "wan2.2-ti2v-5b"
+DIT_FILES = [f"diffusion_pytorch_model-0000{i}-of-00003.safetensors" for i in (1, 2, 3)]
+VAE_FILE = "Wan2.2_VAE.pth"
+
+LATENT_SHAPE = (48, 21, 30, 52)          # (channels, latent frames, height, width) of an 81-frame 832x480 clip
+LAT_H, LAT_W = LATENT_SHAPE[2], LATENT_SHAPE[3]
+PX_PER_LATENT = 16                        # VAE stride
+PX_PER_TOKEN = 32                         # VAE stride x DiT patch size
+SHIFT = 5.0                               # Wan2.2 flow-matching shift (training and sampling)
+STEPS = 50
+VAE_TILE = dict(tile_size=(30, 52), tile_stride=(15, 26))   # only used if tiled=True; frames are encoded whole
+
+# Evaluation videos of the paper were written with imageio's quality scale 9 (libx264, crf 5).
+WRITE_KW = {"quality": 9}
+NULL_PROMPT_FILE = "umt5_null_prompt.pt"      # {NULL_PROMPT: UMT5 embedding}, shipped with the checkpoint
+
+
+# ----------------------------------------------------------------------------- helpers
+def cell_offsets(frame_offset_px: np.ndarray) -> torch.Tensor:
+    """Per-frame camera offsets [81, 2] -> per-latent-frame offsets float32 [21, 2] (mean over the frames of a cell)."""
+    return torch.from_numpy(frames_to_cells(frame_offset_px)).to(torch.float32)
+
+
+def known_input(offset_px: torch.Tensor) -> torch.Tensor:
+    """Input of the known/new channel: float32 [b, 1, 21, 30, 52], 1 = known, 0 = newly revealed."""
+    return knownness(offset_px, LAT_H, LAT_W, PX_PER_LATENT)
+
+
+def cns_kappa(offset_px: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Knownness on the latent grid [b, 1, f, 30, 52] (noise levels) and on the token grid [b, f, 15, 26]
+    (timestep embedding; the mean over the 2 x 2 latent cells of a token)."""
+    kappa = knownness(offset_px, LAT_H, LAT_W, PX_PER_LATENT)
+    return kappa, torch.nn.functional.avg_pool2d(kappa[:, 0], kernel_size=2)
+
+
+def load_pipeline(model_dir: str, device: str = "cuda", vae: bool = True, torch_dtype=torch.bfloat16):
+    """DiffSynth WanVideoPipeline with the DiT (and the VAE); no text encoder."""
+    from diffsynth.pipelines.wan_video import ModelConfig, WanVideoPipeline
+    configs = [ModelConfig(path=[os.path.join(model_dir, f) for f in DIT_FILES])]
+    if vae:
+        configs.append(ModelConfig(path=os.path.join(model_dir, VAE_FILE)))
+    return WanVideoPipeline.from_pretrained(torch_dtype=torch_dtype, device=device, model_configs=configs,
+                                            tokenizer_config=None, redirect_common_files=False)
+
+
+def load_checkpoint(dit: nn.Module, path: str) -> dict:
+    """Load a full ActionRoPE checkpoint into ``dit`` (strict); installs the known/new channel if the checkpoint has
+    one. Returns the checkpoint metadata."""
+    from safetensors.torch import load_file
+    meta = read_metadata(path)
+    sd = load_file(path)
+    has_mask = any(k.startswith(MASK_EMBEDDING + ".") for k in sd)
+    if meta.get("mask_channel") is not None and meta_bool(meta, "mask_channel") != has_mask:
+        raise ValueError(f"{path}: metadata mask_channel={meta['mask_channel']} but the tensors say {has_mask}")
+    if has_mask:
+        install_mask_channel(dit)
+    dit.load_state_dict(sd, strict=True)
+    return meta
+
+
+def build_text_table(prompts: list[str], model_dir: str, device: str = "cuda") -> dict[str, torch.Tensor]:
+    """{prompt: bf16 [512, 4096]} UMT5 embeddings (the encoder of the Wan family)."""
+    from .umt5 import build_text_table as _build
+    return _build(prompts, model_dir, device)
+
+
+class PromptEmbeddings:
+    """Prompt -> UMT5 embedding. Looks the prompt up in an optional embedding file ({prompt: tensor}, e.g. the text
+    table of ``scripts/build_text_table.py``) and otherwise encodes it with UMT5 from ``model_dir`` (loaded lazily)."""
+
+    def __init__(self, model_dir: str, device: str, table: str | dict | None = None):
+        self.model_dir, self.device = model_dir, device
+        self.table = dict(load_text_table(table) if isinstance(table, str) else (table or {}))
+        self.encoder = None
+
+    def __call__(self, prompt: str) -> torch.Tensor:
+        if prompt not in self.table:
+            if self.encoder is None:
+                from .umt5 import UMT5Encoder
+                self.encoder = UMT5Encoder(self.model_dir, self.device)
+            self.table[prompt] = self.encoder([prompt])[0]
+        return self.table[prompt].unsqueeze(0).to(device=self.device, dtype=torch.bfloat16)   # [1, 512, 4096]
+
+
+# ----------------------------------------------------------------------------- VAE
+class _vae_fp32:
+    """Run the VAE in fp32 inside the block."""
+
+    def __init__(self, pipe):
+        self.pipe = pipe
+
+    def __enter__(self):
+        self.pipe.vae.model.float()
+
+    def __exit__(self, *exc):
+        self.pipe.vae.model.to(self.pipe.torch_dtype)
+
+
+@torch.no_grad()
+def encode_first_frame(pipe, img: np.ndarray) -> torch.Tensor:
+    """uint8 [480, 832, 3] -> latent [1, 48, 1, 30, 52]: the pipeline's (bf16-loaded) VAE run in fp32, not tiled, as in
+    the paper's evaluation. ``LatentEncoder`` loads the VAE weights in full fp32 instead (closer to the training data)."""
+    from PIL import Image
+    x = pipe.preprocess_image(Image.fromarray(img), torch_dtype=torch.float32).transpose(0, 1)   # [3, 1, H, W]
+    with _vae_fp32(pipe):
+        z = pipe.vae.encode([x], device=pipe.device, tiled=False, **VAE_TILE)
+    return z.to(dtype=pipe.torch_dtype, device=pipe.device)
+
+
+@torch.no_grad()
+def decode(pipe, latents: torch.Tensor) -> np.ndarray:
+    """latent [1, 48, 21, 30, 52] -> uint8 frames [81, 480, 832, 3] (not tiled)."""
+    video = pipe.vae.decode(latents, device=pipe.device, tiled=False, **VAE_TILE)
+    return np.stack([np.array(f) for f in pipe.vae_output_to_video(video)], 0)
+
+
+# ----------------------------------------------------------------------------- sampling
+@torch.no_grad()
+def sample(dit, z0: torch.Tensor, context: torch.Tensor, offset_px: torch.Tensor, mask_input: torch.Tensor | None, *,
+           steps: int = STEPS, shift: float = SHIFT, seed: int = 0, cns: dict | None = None) -> torch.Tensor:
+    """Flow-matching Euler sampler (no guidance). z0 [1, 48, 1, 30, 52] first-frame latent, context [1, L, 4096],
+    offset_px [1, 21, 2]. ``cns`` = {"shift_known", "alpha"} enables CNS with base shift ``shift``."""
+    from diffsynth.diffusion.flow_match import FlowMatchScheduler
+    sched = FlowMatchScheduler("Wan")
+    sched.set_timesteps(steps, shift=shift)
+    gen = torch.Generator("cpu").manual_seed(seed)
+    latents = torch.randn((1, z0.shape[1], LATENT_SHAPE[1], LAT_H, LAT_W), generator=gen, dtype=torch.float32)
+    latents = latents.to(dtype=z0.dtype, device=z0.device)
+    latents[:, :, 0:1] = z0
+    if cns is not None:
+        kappa, kappa_tok = cns_kappa(offset_px)
+        kappa, kappa_tok = kappa.to(latents.device), kappa_tok.to(latents.device)
+
+        def sigma_tok(s, k):
+            return cns_sigma(s, k, shift, cns["shift_known"], cns["alpha"])
+    for i, t in enumerate(sched.timesteps):
+        timestep = t.unsqueeze(0).to(dtype=latents.dtype, device=latents.device)
+        tmap = None
+        if cns is not None:
+            tmap = (1000.0 * sigma_tok(float(sched.sigmas[i]), kappa_tok)).to(latents.dtype)
+        pred = arope_forward(dit, latents, timestep, context, offset_px=offset_px, px_per_token=PX_PER_TOKEN,
+                             mask_input=mask_input, timestep_map=tmap, first_frame_cond=True)
+        if cns is not None:
+            s_i = float(sched.sigmas[i])
+            s_n = float(sched.sigmas[i + 1]) if i + 1 < steps else 0.0
+            latents = latents + pred * (sigma_tok(s_n, kappa) - sigma_tok(s_i, kappa)).to(latents.dtype)
+        else:
+            latents = sched.step(pred, sched.timesteps[i], latents)
+        latents[:, :, 0:1] = z0
+    return latents
+
+
+class Generator:
+    """First frame + per-frame camera offsets -> 81-frame clip with a full ActionRoPE checkpoint.
+
+    Prompt embeddings are looked up in ``text_table`` (a {prompt: embedding} file such as the training text table;
+    default: ``umt5_null_prompt.pt`` next to the checkpoint, which holds NULL_PROMPT) and otherwise computed with the
+    UMT5 encoder of ``model_dir`` (loaded on first use). The lookup is the path of the paper's evaluation: it avoids
+    the 11 GB encoder, and the evaluation clips are reproduced bit-exactly only if no other model has run on the GPU
+    in the same process before the first sampling call (e.g. running UMT5 first changes cuDNN's attention kernel
+    choice, a bf16-rounding-level difference).
+    """
+
+    def __init__(self, ckpt: str, model_dir: str, device: str = "cuda", text_table: str | dict | None = None, **kw):
+        self.device = device
+        self.pipe = load_pipeline(model_dir, device, vae=True)
+        self.dit = self.pipe.dit
+        self.meta = load_checkpoint(self.dit, ckpt)
+        self.dit.eval().requires_grad_(False)
+        self.mask_channel = hasattr(self.dit, MASK_EMBEDDING)
+        self.shift = meta_float(self.meta, "cns_shift", SHIFT)
+        self.cns = ({"shift_known": meta_float(self.meta, "cns_shift_known"), "alpha": meta_float(self.meta, "cns_alpha", 1.0)}
+                    if meta_bool(self.meta, "cns") else None)
+        if text_table is None:
+            shipped = os.path.join(os.path.dirname(os.path.abspath(ckpt)), NULL_PROMPT_FILE)
+            text_table = shipped if os.path.isfile(shipped) else None
+        self.prompts = PromptEmbeddings(model_dir, device, text_table)
+
+    @torch.no_grad()
+    def __call__(self, first_frame: np.ndarray, frame_offset_px: np.ndarray, prompt: str = NULL_PROMPT, seed: int = 0,
+                 steps: int = STEPS, shift: float | None = None, cns: bool | None = None) -> tuple[np.ndarray, float]:
+        z0 = encode_first_frame(self.pipe, first_frame)
+        context = self.prompts(prompt)
+        offset = cell_offsets(frame_offset_px).unsqueeze(0).to(self.device)
+        mask = known_input(offset) if self.mask_channel else None
+        use_cns = self.cns is not None if cns is None else bool(cns)
+        if use_cns and self.cns is None:
+            raise ValueError("cns=True needs a checkpoint trained with CNS (no cns_shift_known in the metadata)")
+        latents = sample(self.dit, z0, context, offset, mask, steps=steps, shift=self.shift if shift is None else shift,
+                         seed=seed, cns=self.cns if use_cns else None)
+        return decode(self.pipe, latents), float(FPS)
+
+    def eval_frames(self, frames: np.ndarray) -> np.ndarray:
+        return frames
+
+
+# ----------------------------------------------------------------------------- latents
+class LatentEncoder:
+    """81 uint8 frames -> {"latent": bf16 [48, 21, 30, 52]} with the Wan2.2 VAE in fp32 (not tiled)."""
+
+    def __init__(self, model_dir: str, device: str = "cuda"):
+        from diffsynth.pipelines.wan_video import ModelConfig, WanVideoPipeline
+        self.device = device
+        self.pipe = WanVideoPipeline.from_pretrained(torch_dtype=torch.float32, device=device, tokenizer_config=None,
+                                                     model_configs=[ModelConfig(path=os.path.join(model_dir, VAE_FILE))],
+                                                     redirect_common_files=False)
+
+    @torch.no_grad()
+    def __call__(self, frames: np.ndarray, clip=None) -> dict:
+        x = torch.from_numpy(np.array(frames, dtype=np.uint8)).permute(3, 0, 1, 2).float().div_(127.5).sub_(1.0)
+        z = self.pipe.vae.encode([x], device=self.device, tiled=False, **VAE_TILE)[0]
+        if tuple(z.shape) != LATENT_SHAPE:
+            raise ValueError(f"latent shape {tuple(z.shape)}, expected {LATENT_SHAPE} (frames must be 81 x 480 x 832)")
+        return {"latent": z.to(torch.bfloat16).cpu()}
+
+
+# ----------------------------------------------------------------------------- training
+class LatentDataset(torch.utils.data.Dataset):
+    """Pre-computed latents + per-latent-frame scene prompts (looked up in the text table) + camera offsets."""
+
+    def __init__(self, clips, latent_dir: str, table: dict, scene_dropout: float = 0.0):
+        self.clips, self.latent_dir, self.table = clips, latent_dir, table
+        self.scene_dropout = float(scene_dropout)
+        # prompt identity -> id: latent frames with equal ids share one text embedding in the forward pass
+        self.prompt_id = {k: i for i, k in enumerate(sorted(table))}
+        missing = sorted({p for c in clips for p in c.prompts} - set(table))
+        if missing:
+            raise KeyError(f"{len(missing)} prompts are missing from the text table, e.g. {missing[0]!r}")
+        if self.scene_dropout > 0 and NULL_PROMPT not in table:
+            raise KeyError(f"scene dropout needs {NULL_PROMPT!r} in the text table")
+
+    def __len__(self):
+        return len(self.clips)
+
+    def __getitem__(self, i):
+        c = self.clips[i]
+        latent = torch.load(os.path.join(self.latent_dir, c.clip + ".pt"), map_location="cpu", weights_only=False)["latent"]
+        if tuple(latent.shape) != LATENT_SHAPE:
+            raise ValueError(f"{c.clip}: latent shape {tuple(latent.shape)}, expected {LATENT_SHAPE}")
+        drop = self.scene_dropout > 0 and random.random() < self.scene_dropout   # scene dropout: whole clip
+        prompts = [NULL_PROMPT] * len(c.prompts) if drop else c.prompts
+        return {"clip": c.clip,
+                "latent": latent.to(torch.bfloat16).contiguous(),
+                "context": torch.stack([self.table[p] for p in prompts], 0).to(torch.bfloat16),
+                "context_ids": torch.tensor([self.prompt_id[p] for p in prompts], dtype=torch.int64),
+                "offset_px": cell_offsets(c.frame_offset_px)}
+
+
+def collate(samples: list[dict]) -> dict:
+    return {k: torch.stack([s[k] for s in samples], 0) if torch.is_tensor(samples[0][k]) else [s[k] for s in samples]
+            for k in samples[0]}
+
+
+class TrainModule(nn.Module):
+    """DiT + ActionRoPE/CNS flow-matching loss; ``forward`` returns (loss, stats) as required by DeepSpeed."""
+
+    def __init__(self, dit, mask_channel: bool = True, new_weight: float = 2.0, cns: bool = True, cns_shift: float = SHIFT,
+                 cns_shift_known: float = 2.0, cns_alpha: float = 1.0, gradient_checkpointing: bool = True):
+        super().__init__()
+        from diffsynth.diffusion.flow_match import FlowMatchScheduler
+        self.dit = dit
+        self.mask_channel, self.new_weight = bool(mask_channel), float(new_weight)
+        self.cns, self.cns_shift = bool(cns), float(cns_shift)
+        self.cns_shift_known, self.cns_alpha = float(cns_shift_known), float(cns_alpha)
+        self.gradient_checkpointing = bool(gradient_checkpointing)
+        self.scheduler = FlowMatchScheduler("Wan")             # sigma = 5 u / (1 + 4 u) over 1000 training steps
+        self.scheduler.set_timesteps(1000, training=True)
+
+    def forward(self, batch: dict, timestep_id: int, generator: torch.Generator | None = None, train: bool = True):
+        x0 = batch["latent"]                                                    # bf16 [b, 48, 21, 30, 52]
+        b, c = x0.shape[:2]
+        offset_px = batch["offset_px"].to(x0.device)                           # float32 [b, 21, 2]
+        timestep = self.scheduler.timesteps[timestep_id].view(1)
+        noise = (torch.randn_like(x0) if generator is None
+                 else torch.randn(x0.shape, generator=generator, device=x0.device, dtype=x0.dtype))
+        timestep_map = None
+        if self.cns:
+            # CNS: every latent cell is noised to sigma_tok(sigma, kappa); the target stays eps - x0
+            sigma = float(self.scheduler.sigmas[timestep_id])
+            kappa, kappa_tok = cns_kappa(offset_px)
+            sig = cns_sigma(sigma, kappa, self.cns_shift, self.cns_shift_known, self.cns_alpha).to(x0.dtype)
+            x_t = (1.0 - sig) * x0 + sig * noise
+            timestep_map = (1000.0 * cns_sigma(sigma, kappa_tok, self.cns_shift, self.cns_shift_known,
+                                               self.cns_alpha)).to(x0.dtype)
+        else:
+            x_t = self.scheduler.add_noise(x0, noise, timestep)
+        target = self.scheduler.training_target(x0, noise, timestep)
+        x_t[:, :, 0:1] = x0[:, :, 0:1]                                          # latent frame 0 = clean first frame
+        weight = loss_weight_map(offset_px, LAT_H, LAT_W, PX_PER_LATENT, self.new_weight)
+        mask = known_input(offset_px) if self.mask_channel else None
+        pred = arope_forward(self.dit, x_t, timestep.to(device=x0.device, dtype=x0.dtype), batch["context"],
+                             offset_px=offset_px, px_per_token=PX_PER_TOKEN, mask_input=mask, timestep_map=timestep_map,
+                             context_ids=batch["context_ids"], first_frame_cond=True,
+                             use_gradient_checkpointing=self.gradient_checkpointing and train)
+        pred, target, weight = pred[:, :, 1:], target[:, :, 1:], weight[:, :, 1:]  # no loss on the first frame
+        se = (pred.float() - target.float()) ** 2
+        loss_w = (se * weight).sum() / (weight.sum() * c)                      # weight broadcasts over channels
+        loss = loss_w * self.scheduler.training_weight(timestep).to(loss_w.device)
+        return loss, {"loss_w": loss_w.detach(), "mse": se.mean().detach(), "timestep": float(timestep)}
+
+
+class TrainTask:
+    """Training task for ``actionrope.trainer.train`` (see ``configs/train/wan.yaml``)."""
+
+    def __init__(self, cfg: dict):
+        self.cfg = cfg
+        self._table = None
+        from diffsynth.diffusion.flow_match import FlowMatchScheduler
+        self._sched = FlowMatchScheduler("Wan")
+        self._sched.set_timesteps(1000, training=True)
+
+    def _method(self) -> dict:
+        c = self.cfg
+        return dict(mask_channel=bool(c.get("mask_channel", True)), new_weight=float(c.get("new_weight", 2.0)),
+                    cns=bool(c.get("cns", True)), cns_shift=float(c.get("cns_shift", SHIFT)),
+                    cns_shift_known=float(c.get("cns_shift_known", 2.0)), cns_alpha=float(c.get("cns_alpha", 1.0)))
+
+    def build_module(self) -> nn.Module:
+        m = self._method()
+        dit = load_pipeline(self.cfg["model_dir"], device="cpu", vae=False).dit
+        if m["mask_channel"]:
+            install_mask_channel(dit)
+        dit.train().requires_grad_(True)
+        return TrainModule(dit, gradient_checkpointing=bool(self.cfg.get("gradient_checkpointing", True)), **m)
+
+    def datasets(self):
+        c = self.cfg
+        if self._table is None:
+            self._table = load_text_table(c["text_table"])
+        clips = load_annotations(c["annotations"])
+        train = training_clips(clips, min_rt_ratio=float(c.get("min_rt_ratio", 0.8)))
+        train_ds = LatentDataset(train, c["latent_dir"], self._table, scene_dropout=float(c.get("scene_dropout", 0.1)))
+        val_ds = None
+        if int(c.get("val_every", 0)) > 0 and int(c.get("val_clips", 16)) > 0:
+            val = validation_clips(clips, min_rt_ratio=float(c.get("val_min_rt_ratio", 0.9)), n=int(c.get("val_clips", 16)),
+                                   seed=int(c["seed"]))
+            val_ds = LatentDataset(val, c["latent_dir"], self._table, scene_dropout=0.0)
+        return train_ds, val_ds
+
+    collate = staticmethod(collate)
+
+    def on_prepared(self, module, device):
+        # the RoPE tables are plain tensors (not buffers): move them once instead of on every step
+        module.dit.freqs = tuple(t.to(device) for t in module.dit.freqs)
+
+    def export(self, state_dict: dict, step: int):
+        tensors = {k[len("dit."):]: v for k, v in state_dict.items() if k.startswith("dit.")}
+        m = self._method()
+        meta = {"model": MODEL_NAME, "format": "full", "mask_channel": m["mask_channel"], "cns": m["cns"],
+                "cns_shift": m["cns_shift"], "cns_shift_known": m["cns_shift_known"], "cns_alpha": m["cns_alpha"]}
+        return tensors, meta
+
+    def timestep_value(self, tid: int) -> float:
+        return float(self._sched.timesteps[int(tid)])

actionrope/models/wan_dit.py

@@ -0,0 +1,181 @@
+"""ActionRoPE forward pass for the Wan-family video DiT (``diffsynth.models.wan_video_dit.WanModel``).
+
+Used by Wan2.2-TI2V-5B (``wan.py``) and SkyReels-V2-DF-1.3B (``skyreels.py``), which share the DiT architecture.
+
+ActionRoPE (Sec. 3.1): token (i, j) of latent frame k gets the rotary position (k, i + dy_k / s, j + dx_k / s) instead
+of its screen position (k, i, j), where (dx_k, dy_k) is the camera offset of latent frame k in pixels and s the number
+of pixels per token (32 for Wan2.2, 16 for SkyReels-V2). Tokens that see the same world location therefore share the
+same positional address across time. The rotary frequencies are the DiT's own; only the positions change, so no
+parameters are added on the control path.
+
+``arope_forward`` re-implements the DiT forward of DiffSynth's ``model_fn_wan_video`` for image-to-video with a
+clean first frame (latent frame 0 at timestep 0) and adds:
+  * world-coordinate rotary positions (``offset_px``),
+  * the optional known/new input channel (``mask_input``, a zero-initialised Conv3d added to the patch embedding),
+  * per-token timesteps (``timestep_map``) for CNS,
+  * per-latent-frame text (context [b, f, L, C]: every latent frame cross-attends only to its own prompt).
+With ``offset_px=None``, no mask, no timestep map and one prompt it equals the original forward.
+"""
+from __future__ import annotations
+
+import functools
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from diffsynth.core.gradient import gradient_checkpoint_forward
+from diffsynth.models.wan_video_dit import WanModel, modulate, sinusoidal_embedding_1d
+
+ROPE_THETA = 10000.0
+MASK_EMBEDDING = "arope_mask_embedding"
+
+
+# ----------------------------------------------------------------------------- world-coordinate RoPE
+def _inv_freq(dim: int, device=None) -> torch.Tensor:
+    """Inverse frequencies of the DiT's 1-D RoPE (identical to DiffSynth's ``precompute_freqs_cis(dim)``)."""
+    return 1.0 / (ROPE_THETA ** (torch.arange(0, dim, 2, device=device)[: (dim // 2)].double() / dim))
+
+
+def _axis_freqs(pos: torch.Tensor, dim: int) -> torch.Tensor:
+    """1-D RoPE at arbitrary (fractional) positions, complex128 [..., dim // 2]; equals the table at integers."""
+    angle = pos.double().unsqueeze(-1) * _inv_freq(dim, pos.device)
+    return torch.polar(torch.ones_like(angle), angle)
+
+
+def build_world_freqs(dit: WanModel, f: int, h: int, w: int, offset_tok: torch.Tensor | None, device=None) -> torch.Tensor:
+    """Rotary frequencies of all f*h*w tokens at positions (k, i + dy_k, j + dx_k).
+
+    offset_tok: [b, f, 2] = (dx, dy) in token units, or None for the standard screen grid.
+    Returns [f*h*w, 1, d] (standard) or [b, f*h*w, 1, d] (world coordinates), complex like ``dit.freqs``.
+    """
+    f_table, h_table, w_table = dit.freqs
+    if device is None:
+        device = offset_tok.device if offset_tok is not None else f_table.device
+    d_h, d_w = h_table.shape[1] * 2, w_table.shape[1] * 2
+    if offset_tok is None:
+        return torch.cat([
+            f_table[:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+            h_table[:h].view(1, h, 1, -1).expand(f, h, w, -1),
+            w_table[:w].view(1, 1, w, -1).expand(f, h, w, -1),
+        ], dim=-1).reshape(f * h * w, 1, -1).to(device)
+    offset_tok = offset_tok.to(device=device, dtype=torch.float64)
+    b = offset_tok.shape[0]
+    assert offset_tok.shape == (b, f, 2), f"offset_tok must be [b, {f}, 2], got {tuple(offset_tok.shape)}"
+    dx, dy = offset_tok[..., 0], offset_tok[..., 1]
+    pos_h = torch.arange(h, device=device, dtype=torch.float64).view(1, 1, h) + dy.view(b, f, 1)
+    pos_w = torch.arange(w, device=device, dtype=torch.float64).view(1, 1, w) + dx.view(b, f, 1)
+    fh, fw = _axis_freqs(pos_h, d_h), _axis_freqs(pos_w, d_w)
+    freqs = torch.cat([
+        f_table[:f].to(device).view(1, f, 1, 1, -1).expand(b, f, h, w, -1),
+        fh.view(b, f, h, 1, -1).expand(b, f, h, w, -1),
+        fw.view(b, f, 1, w, -1).expand(b, f, h, w, -1),
+    ], dim=-1)
+    return freqs.reshape(b, f * h * w, 1, -1).to(f_table.dtype)
+
+
+# ----------------------------------------------------------------------------- known/new input channel
+def install_mask_channel(dit: WanModel) -> WanModel:
+    """Add the zero-initialised known/new input channel ``dit.arope_mask_embedding`` = Conv3d(1, dim, patch, patch).
+    The model output is unchanged at installation time."""
+    if not hasattr(dit, MASK_EMBEDDING):
+        ref = dit.patch_embedding
+        conv = nn.Conv3d(1, dit.dim, kernel_size=dit.patch_size, stride=dit.patch_size)
+        nn.init.zeros_(conv.weight)
+        nn.init.zeros_(conv.bias)
+        dit.add_module(MASK_EMBEDDING, conv.to(device=ref.weight.device, dtype=ref.weight.dtype))
+    return dit
+
+
+# ----------------------------------------------------------------------------- forward
+def _per_token_timestep(latents, timestep, n_tok: int, first_frame_cond: bool, timestep_map=None) -> torch.Tensor:
+    """[b, f * n_tok] timesteps (0..1000); latent frame 0 is the clean condition frame (t = 0)."""
+    b, f = latents.shape[0], latents.shape[2]
+    if timestep_map is not None:
+        t_tok = timestep_map.to(device=latents.device, dtype=latents.dtype).reshape(b, f, n_tok).clone()
+    else:
+        timestep = torch.as_tensor(timestep, device=latents.device).to(latents.dtype).reshape(-1)
+        timestep = timestep.expand(b) if timestep.numel() == 1 else timestep
+        t_tok = timestep.view(b, 1, 1).expand(b, f, n_tok).clone()
+    if first_frame_cond:
+        t_tok[:, 0] = 0
+    return t_tok.reshape(b, f * n_tok)
+
+
+def _embed_context(dit: WanModel, context: torch.Tensor, context_ids=None) -> torch.Tensor:
+    """[b, L, C] -> [b, L, dim]; per-frame context [b, f, L, C] -> [(b f), L, dim]. Frames with equal
+    ``context_ids`` [b, f] share a prompt and are embedded once."""
+    if context.dim() == 3:
+        return dit.text_embedding(context)
+    b, f, L, c = context.shape
+    rows = context.reshape(b * f, L, c)
+    if context_ids is None:
+        return dit.text_embedding(rows)
+    ids = torch.as_tensor(context_ids).reshape(-1).tolist()
+    first = {}
+    for i, k in enumerate(ids):
+        first.setdefault(k, i)
+    uniq = list(first)
+    emb = dit.text_embedding(rows[torch.tensor([first[k] for k in uniq], device=rows.device)])
+    return emb[torch.tensor([uniq.index(k) for k in ids], device=rows.device)]
+
+
+def _block_forward(block, x, context, t_mod, freqs, *, f: int):
+    """``DiTBlock.forward`` with per-frame text cross-attention when context is [(b f), L, dim]."""
+    has_seq = len(t_mod.shape) == 4
+    chunk_dim = 2 if has_seq else 1
+    shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+        block.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=chunk_dim)
+    if has_seq:
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            shift_msa.squeeze(2), scale_msa.squeeze(2), gate_msa.squeeze(2),
+            shift_mlp.squeeze(2), scale_mlp.squeeze(2), gate_mlp.squeeze(2))
+    x = block.gate(x, gate_msa, block.self_attn(modulate(block.norm1(x), shift_msa, scale_msa), freqs))
+    q_in = block.norm3(x)
+    if context.shape[0] == x.shape[0]:
+        x = x + block.cross_attn(q_in, context)
+    else:
+        q_in = rearrange(q_in, "b (f n) c -> (b f) n c", f=f)
+        x = x + rearrange(block.cross_attn(q_in, context), "(b f) n c -> b (f n) c", f=f)
+    return block.gate(x, gate_mlp, block.ffn(modulate(block.norm2(x), shift_mlp, scale_mlp)))
+
+
+def arope_forward(dit: WanModel, latents: torch.Tensor, timestep, context: torch.Tensor, *,
+                  offset_px: torch.Tensor | None = None, px_per_token: float = 32, mask_input: torch.Tensor | None = None,
+                  timestep_map: torch.Tensor | None = None, context_ids=None, first_frame_cond: bool = True,
+                  t_mod_fn=None, use_gradient_checkpointing: bool = False) -> torch.Tensor:
+    """Velocity prediction for latents [b, c, f, H, W].
+
+    offset_px     [b, f, 2] camera offset of every latent frame in pixels (None: standard RoPE)
+    px_per_token  pixels per DiT token (VAE stride x patch size)
+    mask_input    [b, 1, f, H, W] known = 1 / new = 0 (requires ``install_mask_channel``)
+    timestep_map  [b, f, H / 2, W / 2] per-token timesteps for CNS (None: ``timestep`` for every token)
+    context       [b, L, C] (one prompt) or [b, f, L, C] (one prompt per latent frame, see ``context_ids``)
+    t_mod_fn      optional callable applied to the timestep modulation [b, n, 6, dim]
+    """
+    b, _, f, lat_h, lat_w = latents.shape
+    _, p_h, p_w = dit.patch_size
+    n_tok = (lat_h // p_h) * (lat_w // p_w)
+
+    t_tok = _per_token_timestep(latents, timestep, n_tok, first_frame_cond, timestep_map)
+    t = dit.time_embedding(sinusoidal_embedding_1d(dit.freq_dim, t_tok.reshape(-1)).view(b, t_tok.shape[1], -1))
+    t_mod = dit.time_projection(t).unflatten(2, (6, dit.dim))
+    if t_mod_fn is not None:
+        t_mod = t_mod_fn(t_mod)
+    context = _embed_context(dit, context, context_ids)
+
+    x = dit.patch_embedding(latents)
+    if mask_input is not None:
+        if not hasattr(dit, MASK_EMBEDDING):
+            raise ValueError("mask_input given but the known/new input channel is not installed")
+        x = x + getattr(dit, MASK_EMBEDDING)(mask_input.to(device=x.device, dtype=x.dtype))
+    f, h, w = x.shape[2:]
+    x = rearrange(x, "b c f h w -> b (f h w) c").contiguous()
+
+    offset_tok = None if offset_px is None else offset_px.to(torch.float64) / px_per_token
+    freqs = build_world_freqs(dit, f, h, w, offset_tok, device=x.device)
+    for block in dit.blocks:
+        x = gradient_checkpoint_forward(functools.partial(_block_forward, block, f=f),
+                                        use_gradient_checkpointing, False, x, context, t_mod, freqs)
+    x = dit.head(x, t)
+    return dit.unpatchify(x, (f, h, w))

actionrope/prompts.py

@@ -0,0 +1,34 @@
+"""Text prompts.
+
+ActionRoPE never puts the action into the text: the prompt only describes the scene,
+
+    "In a 2.5D top down view, the player is in a high rocky mountain pass."
+
+and ``NULL_PROMPT`` (no scene) is used for scene dropout during training and for all evaluation clips.
+"""
+from __future__ import annotations
+
+import re
+
+NULL_PROMPT = "In a 2.5D top down view."
+SCENE_TEMPLATE = "In a 2.5D top down view, the player is {scene}."
+
+_WITH_ACTION = re.compile(r"^In a 2\.5D top down view, (?P<scene>.+), the player is (?P<action>.+)\.$")
+_ACTION_ONLY = re.compile(r"^In a 2\.5D top down view, the player is (?P<action>.+)\.$")
+
+
+def scene_prompt(scene: str | None) -> str:
+    """Scene description (e.g. "in a dark cave lit by blue crystal glow") -> prompt; None -> NULL_PROMPT."""
+    return NULL_PROMPT if not scene else SCENE_TEMPLATE.format(scene=scene)
+
+
+def strip_action(prompt: str) -> str:
+    """"In a 2.5D top down view, {scene}, the player is {action}." -> scene_prompt(scene).
+
+    Captions that carry an action but no scene collapse to NULL_PROMPT."""
+    m = _WITH_ACTION.match(prompt)
+    if m:
+        return scene_prompt(m.group("scene"))
+    if _ACTION_ONLY.match(prompt):
+        return NULL_PROMPT
+    raise ValueError(f"prompt does not follow the caption template: {prompt!r}")

app.py

@@ -0,0 +1,110 @@
+"""ActionRoPE: action as coordinates for spatially consistent control in planar-camera game world models.
+
+Gradio demo of the paper's Wan2.2-TI2V-5B backbone: a first frame (832x480) plus an
+action string produce an 81-frame (5 s, 16 fps) video of the game world, with the
+camera displacement injected through ActionRoPE (world-coordinate rotary positions)
+and Coordinate-Conditioned Noise Scheduling (CNS).
+"""
+import os
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # noqa: E402  must come before torch
+import tempfile  # noqa: E402
+import time  # noqa: E402
+
+import gradio as gr  # noqa: E402
+import numpy as np  # noqa: E402
+
+from actionrope.geometry import FRAMES as N_FRAMES  # noqa: E402
+from actionrope.geometry import offsets_from_actions  # noqa: E402
+from actionrope.media import load_first_frame  # noqa: E402
+from actionrope.media import write_video  # noqa: E402
+from actionrope.prompts import NULL_PROMPT  # noqa: E402
+from arope_loader import Generator, download_space_weights  # noqa: E402
+
+CKPT, NULL_PROMPT_PT, VAE_PATH = download_space_weights()
+print("[app] weights cached; building model...")
+GEN = Generator(CKPT, VAE_PATH, NULL_PROMPT_PT, device="cuda", torch_dtype=__import__("torch").bfloat16)
+print("[app] model ready")
+
+
+@spaces.GPU(duration=180)
+def play(first_frame, actions: str, seed: int, steps: int, use_cns: bool = True):
+    """Render an 81-frame game-world video from a first frame and an action plan.
+
+    Args:
+        first_frame: initial observation frame (any size; centre-cropped to 26:15 and resized to 832x480).
+        actions: comma-separated move plan, each segment "<direction>[*speed]:<frames>",
+            e.g. "right:40,left:40". Directions: idle, up, down, left, right, and the four
+            diagonals (aliases u/d/l/r/ul/ur/dl/dr). Speed 1.0 is the game's walking speed
+            (69.2 px/s horizontally, 46.0 px/s vertically at 832x480).
+        seed: RNG seed for the initial noise.
+        steps: number of Euler flow-matching steps (paper setting: 50).
+        use_cns: use Coordinate-Conditioned Noise Scheduling (the paper's method).
+
+    Returns:
+        The rendered 832x480 16 fps mp4 of the 5-second rollout.
+    """
+    t0 = time.perf_counter()
+    offsets = offsets_from_actions(actions)  # validates segment count/spec
+    if np.abs(offsets[N_FRAMES - 1]).max() == 0 and "idle" not in actions:
+        pass  # a valid all-idle plan is fine
+    frames, fps = GEN(load_first_frame(first_frame), offsets, prompt=NULL_PROMPT,
+                      seed=int(seed), steps=int(steps), cns=bool(use_cns))
+    out = os.path.join(tempfile.gettempdir(), f"arope_{int(time.time() * 1000) % 100000}.mp4")
+    write_video(out, frames, fps=fps, quality=9)
+    print(f"[app] generated {len(frames)} frames in {time.perf_counter() - t0:.1f}s -> {out}")
+    return out
+
+
+def build_ui():
+    with gr.Blocks(theme=gr.themes.Citrus(), title="ActionRoPE") as demo:
+        gr.Markdown(
+            """
+            # 🎮 ActionRoPE — action as coordinates for game world models
+            Drive a video-diffusion world model with plain **actions** instead of text.
+            The first frame pins the scene; the action plan moves the camera through the world.
+            Backbones: Wan2.2-TI2V-5B (this demo) — from the paper
+            [ActionRoPE: action as coordinates for spatially consistent control in planar-camera game world models](https://huggingface.co/papers/2610.25861).
+            """
+        )
+        with gr.Row():
+            with gr.Column():
+                first_frame = gr.Image(label="First frame (832x480; other sizes are centre-cropped)", type="filepath")
+                actions = gr.Textbox(
+                    label="Actions",
+                    value="right:40,left:40",
+                    info="comma list of <direction>[*speed]:<frames> over 80 frames; "
+                         "speed 1.0 = walking speed (69.2 px/s horizontal / 46.0 px/s vertical)",
+                    placeholder="right:40,left:40",
+                )
+                btn = gr.Button("Generate rollout", variant="primary")
+                with gr.Accordion("Advanced", open=False):
+                    seed = gr.Number(0, label="Seed", precision=0)
+                    steps = gr.Slider(10, 50, value=50, step=5, label="Sampling steps", info="paper setting = 50")
+                    use_cns = gr.Checkbox(True, label="CNS (coordinate-conditioned noise scheduling)")
+            with gr.Column():
+                video = gr.Video(label="5-second rollout (81 frames, 16 fps)", autoplay=True)
+
+        gr.Examples(
+            examples=[
+                ["examples/back_01.png", "right:40,left:40"],
+                ["examples/back_264.png", "right:60,left:20"],
+                ["examples/bamboo_L.png", "up-right*1.5:32,left:40,down-left:8"],
+                ["examples/watercolor_L.png", "right:80"],
+            ],
+            inputs=[first_frame, actions],
+            fn=play,
+            outputs=[video],
+            cache_examples=True,
+            cache_mode="lazy",
+            label="Example first frames and action plans",
+        )
+
+        btn.click(play, inputs=[first_frame, actions, seed, steps, use_cns], outputs=[video])
+    return demo
+
+
+if __name__ == "__main__":
+    build_ui().launch(mcp_server=True)

arope_loader.py

@@ -0,0 +1,166 @@
+"""Weight plumbing for the ActionRoPE Space.
+
+The upstream ``actionrope.models.wan.load_pipeline`` loads the *base* Wan2.2-TI2V-5B
+DiT shards (3 files, ~20 GB) and then overwrites them with the ActionRoPE
+checkpoint. The ActionRoPE checkpoint is itself a full DiffSynth-native DiT state
+dict (827 tensors: blocks.* / head.* / patch_embedding / text_embedding /
+time_embedding / time_projection / arope_mask_embedding), so the base shards are
+redundant. This module builds ``diffsynth.models.wan_video_dit.WanModel`` directly
+from the registered TI2V-5B config and strict-loads the ActionRoPE checkpoint
+(exactly what upstream ``load_checkpoint`` + ``TrainTask.build_module`` do), loads
+the Wan2.2 VAE the same way, and returns a minimal pipeline object plus a
+``Generator`` whose ``__call__`` is the upstream one — the paper's sampling path,
+unchanged.
+"""
+from __future__ import annotations
+
+import os
+
+
+def download_space_weights() -> tuple[str, str, str]:
+    """Fetch the three runtime weight files (ActionRoPE DiT, its NULL-prompt table, the Wan2.2 VAE).
+
+    Returns absolute (ckpt_path, null_prompt_path, vae_path).
+    """
+    from huggingface_hub import hf_hub_download
+
+    ckpt_path = hf_hub_download(
+        "teawhite/ActionRoPE", "wan2.2-ti2v-5b/actionrope_cns.safetensors", repo_type="model"
+    )
+    null_prompt_path = hf_hub_download(
+        "teawhite/ActionRoPE", "wan2.2-ti2v-5b/umt5_null_prompt.pt", repo_type="model"
+    )
+    vae_path = hf_hub_download("Wan-AI/Wan2.2-TI2V-5B", "Wan2.2_VAE.pth", repo_type="model")
+    return ckpt_path, null_prompt_path, vae_path
+
+
+class _VAEPipeline:
+    """Minimal stand-in for DiffSynth's WanVideoPipeline holding only the VAE.
+
+    Exposes exactly the three pipeline entry points ActionRoPE's ``wan.py`` uses:
+    ``preprocess_image``, ``vae.encode/decode`` and ``vae_output_to_video``.
+    """
+
+    def __init__(self, vae, device, torch_dtype):
+        from diffsynth.diffusion.base_pipeline import BasePipeline
+
+        self.vae = vae
+        self.device = device
+        self.torch_dtype = torch_dtype
+        self.preprocess_image = BasePipeline.preprocess_image.__get__(self)
+        self.vae_output_to_video = BasePipeline.vae_output_to_video.__get__(self)
+
+    def encode(self, videos, device, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)):
+        return self.vae.encode(videos, device=device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)
+
+    def decode(self, hidden_states, device, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)):
+        return self.vae.decode(hidden_states, device=device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)
+
+
+def build_pipe(ckpt_path: str, vae_path: str, device, torch_dtype) -> tuple:
+    """Build (pipeline, dit, meta): DiT from the ActionRoPE ckpt, VAE from Wan2.2_VAE.pth.
+
+    No base-DiT download, no text encoder — the DiT state dict is complete in the
+    ActionRoPE checkpoint, and prompts come from the shipped ``umt5_null_prompt.pt``
+    embedding table.
+    """
+    from safetensors import safe_open
+
+    from diffsynth.core.loader.file import load_state_dict
+    from diffsynth.models.wan_video_dit import WanModel
+    from diffsynth.models.wan_video_vae import WanVideoVAE38
+    from diffsynth.utils.state_dict_converters.wan_video_vae import WanVideoVAEStateDictConverter
+
+    from actionrope.checkpoint import read_metadata
+    from actionrope.models.wan_dit import install_mask_channel
+
+    # Registered DiffSynth config for Wan-AI/Wan2.2-TI2V-5B
+    # (diffsynth/configs/model_configs.py, model_hash 1f5ab7703c6fc803fdded85ff040c316)
+    ti2v_5b_kwargs = {
+        "has_image_input": False, "patch_size": [1, 2, 2], "in_dim": 48, "dim": 3072, "ffn_dim": 14336,
+        "freq_dim": 256, "text_dim": 4096, "out_dim": 48, "num_heads": 24, "num_layers": 30, "eps": 1e-06,
+        "seperated_timestep": True, "require_clip_embedding": False, "require_vae_embedding": False,
+        "fuse_vae_embedding_in_latents": True,
+    }
+
+    meta = read_metadata(ckpt_path)
+    sd = {}
+    with safe_open(ckpt_path, framework="pt", device="cpu", strided=False) as f:
+        for k in f.keys():
+            sd[k] = f.get_tensor(k)
+    has_mask = any(k.startswith("arope_mask_embedding.") for k in sd)
+
+    dit = WanModel(**ti2v_5b_kwargs)
+    if has_mask:
+        install_mask_channel(dit)
+    dit.load_state_dict(sd, strict=True)  # raises if the checkpoint is not the full TI2V-5B DiT
+    del sd
+    dit = dit.to(device=device, dtype=torch_dtype).eval().requires_grad_(False)
+    # RoPE tables are plain tensors (not buffers): move them once (mirrors TrainTask.on_prepared)
+    dit.freqs = tuple(t.to(device) for t in dit.freqs)
+
+    vae_sd = WanVideoVAEStateDictConverter(load_state_dict(vae_path, torch_dtype=torch_dtype, device="cpu"))
+    vae = WanVideoVAE38()
+    vae.load_state_dict(vae_sd, strict=True)
+    del vae_sd
+    vae = vae.to(device=device, dtype=torch_dtype).eval().requires_grad_(False)
+    vae.mean = vae.mean.to(device)
+    vae.std = vae.std.to(device)
+
+    print(f"[arope_loader] DiT {sum(p.numel() for p in dit.parameters()) / 1e9:.2f} B params, "
+          f"mask_channel={has_mask}, metadata={dict(sorted(meta.items()))}")
+    return _VAEPipeline(vae, device, torch_dtype), dit, meta
+
+
+class Generator:
+    """Upstream ``actionrope.models.wan.Generator.__call__`` on a directly-built pipe.
+
+    The upstream ``__init__`` calls ``load_pipeline(model_dir)`` which needs the base
+    DiT shards; everything else in it is replicated here verbatim.
+    """
+
+    def __init__(self, ckpt_path: str, vae_path: str, null_prompt_path: str,
+                 device="cuda", torch_dtype=None):
+        import torch
+
+        from actionrope.checkpoint import meta_bool, meta_float
+        from actionrope.models.wan import SHIFT, PromptEmbeddings
+        from actionrope.models.wan_dit import MASK_EMBEDDING
+
+        torch_dtype = torch_dtype or torch.bfloat16
+        self.device = device
+        self.pipe, self.dit, self.meta = build_pipe(ckpt_path, vae_path, device, torch_dtype)
+        self.mask_channel = hasattr(self.dit, MASK_EMBEDDING)
+        self.shift = meta_float(self.meta, "cns_shift", SHIFT)
+        self.cns = (
+            {
+                "shift_known": meta_float(self.meta, "cns_shift_known"),
+                "alpha": meta_float(self.meta, "cns_alpha", 1.0),
+            }
+            if meta_bool(self.meta, "cns")
+            else None
+        )
+        self.prompts = PromptEmbeddings("", device, null_prompt_path)
+
+    # The full paper sampling path, verbatim from actionrope.models.wan.Generator.
+    from actionrope.models.wan import Generator as _Upstream  # noqa: F401  (readability marker)
+
+    def __call__(self, first_frame, frame_offset_px, prompt=None, seed=0, steps=50, shift=None, cns=None):
+        import numpy as np
+
+        from actionrope.models import wan as wan_module
+        from actionrope.prompts import NULL_PROMPT
+
+        z0 = wan_module.encode_first_frame(self.pipe, first_frame)
+        context = self.prompts(prompt if prompt is not None else NULL_PROMPT)
+        offset = wan_module.cell_offsets(frame_offset_px).unsqueeze(0).to(self.device)
+        mask = wan_module.known_input(offset) if self.mask_channel else None
+        use_cns = self.cns is not None if cns is None else bool(cns)
+        if use_cns and self.cns is None:
+            raise ValueError("cns=True needs a checkpoint trained with CNS")
+        latents = wan_module.sample(
+            self.dit, z0, context, offset, mask,
+            steps=steps, shift=self.shift if shift is None else shift, seed=seed,
+            cns=self.cns if use_cns else None,
+        )
+        return wan_module.decode(self.pipe, latents), float(wan_module.FPS)

checkpoints/PLACEHOLDER.md

@@ -0,0 +1,4 @@
+Weights are downloaded at runtime by app.py from:
+- teawhite/ActionRoPE : wan2.2-ti2v-5b/actionrope_cns.safetensors (10 GB, full ActionRoPE DiT)
+- teawhite/ActionRoPE : wan2.2-ti2v-5b/umt5_null_prompt.pt
+- Wan-AI/Wan2.2-TI2V-5B : Wan2.2_VAE.pth (2.8 GB)

requirements.txt

@@ -0,0 +1,14 @@
+diffsynth @ git+https://github.com/modelscope/DiffSynth-Studio@7686e54d41d25c0e8ed5f1318acc23b6bb832654
+torchvision
+torchaudio
+transformers
+accelerate
+safetensors
+einops
+ftfy
+pandas
+modelscope
+regex
+imageio[ffmpeg]
+av
+numpy