Initial commit

.gitignore

@@ -0,0 +1,27 @@
+__pycache__/
+*.py[cod]
+*$py.class
+
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.ipynb_checkpoints/
+
+.env
+.venv/
+venv/
+
+outputs/
+slurm_logs/
+latest-run
+.wandb_run_id
+.wandb_osh_command_dir/
+wandb/
+
+*.log
+*.ckpt
+*.pt
+*.pth
+*.safetensors
+
+data/

LICENSE.md

@@ -0,0 +1,14 @@
+# S-Lab License 1.0
+
+Copyright 2025 S-Lab
+
+Redistribution and use for non-commercial purpose in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.\
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+4. In the event that redistribution and/or use for commercial purpose in source or binary forms, with or without modification is required, please contact the contributor(s) of the work.
+
+
+---
+For the commercial use of the code, please consult Prof. Chen Change Loy (ccloy@ntu.edu.sg)

README.md

@@ -0,0 +1,37 @@
+# DeMemWM
+
+This repo is forked from [Boyuan Chen](https://boyuan.space/)'s research template [repo](https://github.com/buoyancy99/research-template). By its MIT license, you must keep this sentence in `README.md` and the `LICENSE` file to credit the author.
+
+DeMemWM is a Memory-DiT video prediction project built on the local research template. The primary algorithm entry point is `DeMemWMMinecraft`, registered through the Hydra algorithm config `dememwm_memory_dit`.
+
+## Quick Start
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+python -m pytest tests
+```
+
+Run a local offline experiment after setting the dataset path in `configurations/dataset/video_minecraft.yaml`:
+
+```bash
+python main.py +name=dememwm_debug algorithm=dememwm_memory_dit wandb.mode=offline
+```
+
+Use `resume_ckpt_path=/path/to/checkpoint.ckpt` for deterministic checkpoint resume, or keep `auto_resume=true` to resume from `output_dir/checkpoints` when available.
+
+## Layout
+
+- `algorithms/worldmem/dememwm/`: DeMemWM memory construction, retrieval, scheduling, diagnostics, and injection code.
+- `algorithms/worldmem/dememwm_memory_dit.py`: primary DeMemWM algorithm class.
+- `configurations/algorithm/dememwm_memory_dit.yaml`: consumed DeMemWM training and evaluation contract.
+- `scripts/`: Slurm and inspection scripts using the DeMemWM naming.
+- `tests/`: static and unit coverage for DeMemWM config, retrieval, compression, schedules, and training behavior.
+
+## Reproducibility Notes
+
+- Keep `wandb.mode=offline` for local reproducible runs that do not depend on network access.
+- Set `seed=<int>` on the command line to seed Lightning and dataloader workers.
+- Runtime artifacts such as `outputs/`, `slurm_logs/`, Python caches, checkpoints, and local datasets are ignored by git.
+- The default Hydra training config selects `algorithm: dememwm_memory_dit`.

algorithms/common/base_algo.py

@@ -0,0 +1,21 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from omegaconf import DictConfig
+
+
+class BaseAlgo(ABC):
+    """
+    A base class for generic algorithms.
+    """
+
+    def __init__(self, cfg: DictConfig):
+        super().__init__()
+        self.cfg = cfg
+
+    @abstractmethod
+    def run(*args: Any, **kwargs: Any) -> Any:
+        """
+        Run the algorithm.
+        """
+        raise NotImplementedError

algorithms/common/base_pytorch_algo.py

@@ -0,0 +1,277 @@
+from abc import ABC, abstractmethod
+import warnings
+import random
+from typing import Any, Union, Sequence, Optional
+
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from omegaconf import DictConfig
+import lightning.pytorch as pl
+import torch
+import numpy as np
+from PIL import Image
+import wandb
+import einops
+
+
+class BasePytorchAlgo(pl.LightningModule, ABC):
+    """
+    A base class for Pytorch algorithms using Pytorch Lightning.
+    See https://lightning.ai/docs/pytorch/stable/starter/introduction.html for more details.
+    """
+
+    def __init__(self, cfg: DictConfig):
+        super().__init__()
+        self.cfg = cfg
+        self._build_model()
+
+    @abstractmethod
+    def _build_model(self):
+        """
+        Create all pytorch nn.Modules here.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def training_step(self, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        r"""Here you compute and return the training loss and some additional metrics for e.g. the progress bar or
+        logger.
+
+        Args:
+            batch: The output of your data iterable, normally a :class:`~torch.utils.data.DataLoader`.
+            batch_idx: The index of this batch.
+            dataloader_idx: (only if multiple dataloaders used) The index of the dataloader that produced this batch.
+
+        Return:
+            Any of these options:
+            - :class:`~torch.Tensor` - The loss tensor
+            - ``dict`` - A dictionary. Can include any keys, but must include the key ``'loss'``.
+            - ``None`` - Skip to the next batch. This is only supported for automatic optimization.
+                This is not supported for multi-GPU, TPU, IPU, or DeepSpeed.
+
+        In this step you'd normally do the forward pass and calculate the loss for a batch.
+        You can also do fancier things like multiple forward passes or something model specific.
+
+        Example::
+
+            def training_step(self, batch, batch_idx):
+                x, y, z = batch
+                out = self.encoder(x)
+                loss = self.loss(out, x)
+                return loss
+
+        To use multiple optimizers, you can switch to 'manual optimization' and control their stepping:
+
+        .. code-block:: python
+
+            def __init__(self):
+                super().__init__()
+                self.automatic_optimization = False
+
+
+            # Multiple optimizers (e.g.: GANs)
+            def training_step(self, batch, batch_idx):
+                opt1, opt2 = self.optimizers()
+
+                # do training_step with encoder
+                ...
+                opt1.step()
+                # do training_step with decoder
+                ...
+                opt2.step()
+
+        Note:
+            When ``accumulate_grad_batches`` > 1, the loss returned here will be automatically
+            normalized by ``accumulate_grad_batches`` internally.
+
+        """
+        return super().training_step(*args, **kwargs)
+
+    def configure_optimizers(self):
+        """
+        Return an optimizer. If you need to use more than one optimizer, refer to pytorch lightning documentation:
+        https://lightning.ai/docs/pytorch/stable/common/optimization.html
+        """
+        parameters = self.parameters()
+        return torch.optim.Adam(parameters, lr=self.cfg.lr)
+
+    def on_save_checkpoint(self, checkpoint):
+        checkpoint["rng_states"] = {
+            "python": random.getstate(),
+            "numpy": np.random.get_state(),
+            "torch": torch.get_rng_state(),
+            "cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
+        }
+
+    def on_load_checkpoint(self, checkpoint):
+        rng_states = checkpoint.get("rng_states")
+        if rng_states is None:
+            if getattr(self, "_strict_resume_state", False):
+                raise RuntimeError(
+                    "Cannot deterministically resume because this checkpoint has no rng_states entry. "
+                    "Use a checkpoint created after automatic resume support was added, or start a fresh run."
+                )
+            return
+
+        random.setstate(rng_states["python"])
+        np.random.set_state(rng_states["numpy"])
+        torch.set_rng_state(rng_states["torch"])
+        if torch.cuda.is_available() and rng_states["cuda"] is not None:
+            torch.cuda.set_rng_state_all(rng_states["cuda"])
+
+    def log_video(
+        self,
+        key: str,
+        video: Union[np.ndarray, torch.Tensor],
+        mean: Union[np.ndarray, torch.Tensor, Sequence, float] = None,
+        std: Union[np.ndarray, torch.Tensor, Sequence, float] = None,
+        fps: int = 5,
+        format: str = "mp4",
+    ):
+        """
+        Log video to wandb. WandbLogger in pytorch lightning does not support video logging yet, so we call wandb directly.
+
+        Args:
+            video: a numpy array or tensor, either in form (time, channel, height, width) or in the form
+                (batch, time, channel, height, width). The content must be be in 0-255 if under dtype uint8
+                or [0, 1] otherwise.
+            mean: optional, the mean to unnormalize video tensor, assuming unnormalized data is in [0, 1].
+            std: optional, the std to unnormalize video tensor, assuming unnormalized data is in [0, 1].
+            key: the name of the video.
+            fps: the frame rate of the video.
+            format: the format of the video. Can be either "mp4" or "gif".
+        """
+
+        if isinstance(video, torch.Tensor):
+            video = video.detach().cpu().numpy()
+
+        expand_shape = [1] * (len(video.shape) - 2) + [3, 1, 1]
+        if std is not None:
+            if isinstance(std, (float, int)):
+                std = [std] * 3
+            if isinstance(std, torch.Tensor):
+                std = std.detach().cpu().numpy()
+            std = np.array(std).reshape(*expand_shape)
+            video = video * std
+        if mean is not None:
+            if isinstance(mean, (float, int)):
+                mean = [mean] * 3
+            if isinstance(mean, torch.Tensor):
+                mean = mean.detach().cpu().numpy()
+            mean = np.array(mean).reshape(*expand_shape)
+            video = video + mean
+
+        if video.dtype != np.uint8:
+            video = np.clip(video, a_min=0, a_max=1) * 255
+            video = video.astype(np.uint8)
+
+        self.logger.experiment.log(
+            {
+                key: wandb.Video(video, fps=fps, format=format),
+            },
+            step=self.global_step,
+        )
+
+    def log_image(
+        self,
+        key: str,
+        image: Union[np.ndarray, torch.Tensor, Image.Image, Sequence[Image.Image]],
+        mean: Union[np.ndarray, torch.Tensor, Sequence, float] = None,
+        std: Union[np.ndarray, torch.Tensor, Sequence, float] = None,
+        **kwargs: Any,
+    ):
+        """
+        Log image(s) using WandbLogger.
+        Args:
+            key: the name of the video.
+            image: a single image or a batch of images. If a batch of images, the shape should be (batch, channel, height, width).
+            mean: optional, the mean to unnormalize image tensor, assuming unnormalized data is in [0, 1].
+            std: optional, the std to unnormalize tensor, assuming unnormalized data is in [0, 1].
+            kwargs: optional, WandbLogger log_image kwargs, such as captions=xxx.
+        """
+        if isinstance(image, Image.Image):
+            image = [image]
+        elif len(image) and not isinstance(image[0], Image.Image):
+            if isinstance(image, torch.Tensor):
+                image = image.detach().cpu().numpy()
+
+            if len(image.shape) == 3:
+                image = image[None]
+
+            if image.shape[1] == 3:
+                if image.shape[-1] == 3:
+                    warnings.warn(f"Two channels in shape {image.shape} have size 3, assuming channel first.")
+                image = einops.rearrange(image, "b c h w -> b h w c")
+
+            if std is not None:
+                if isinstance(std, (float, int)):
+                    std = [std] * 3
+                if isinstance(std, torch.Tensor):
+                    std = std.detach().cpu().numpy()
+                std = np.array(std)[None, None, None]
+                image = image * std
+            if mean is not None:
+                if isinstance(mean, (float, int)):
+                    mean = [mean] * 3
+                if isinstance(mean, torch.Tensor):
+                    mean = mean.detach().cpu().numpy()
+                mean = np.array(mean)[None, None, None]
+                image = image + mean
+
+            if image.dtype != np.uint8:
+                image = np.clip(image, a_min=0.0, a_max=1.0) * 255
+                image = image.astype(np.uint8)
+                image = [img for img in image]
+
+        self.logger.log_image(key=key, images=image, **kwargs)
+
+    def log_gradient_stats(self):
+        """Log gradient statistics such as the mean or std of norm."""
+
+        with torch.no_grad():
+            grad_norms = []
+            gpr = []  # gradient-to-parameter ratio
+            for param in self.parameters():
+                if param.grad is not None:
+                    grad_norms.append(torch.norm(param.grad).item())
+                    gpr.append(torch.norm(param.grad) / torch.norm(param))
+            if len(grad_norms) == 0:
+                return
+            grad_norms = torch.tensor(grad_norms)
+            gpr = torch.tensor(gpr)
+            self.log_dict(
+                {
+                    "train/grad_norm/min": grad_norms.min(),
+                    "train/grad_norm/max": grad_norms.max(),
+                    "train/grad_norm/std": grad_norms.std(),
+                    "train/grad_norm/mean": grad_norms.mean(),
+                    "train/grad_norm/median": torch.median(grad_norms),
+                    "train/gpr/min": gpr.min(),
+                    "train/gpr/max": gpr.max(),
+                    "train/gpr/std": gpr.std(),
+                    "train/gpr/mean": gpr.mean(),
+                    "train/gpr/median": torch.median(gpr),
+                }
+            )
+
+    def register_data_mean_std(
+        self, mean: Union[str, float, Sequence], std: Union[str, float, Sequence], namespace: str = "data"
+    ):
+        """
+        Register mean and std of data as tensor buffer.
+
+        Args:
+            mean: the mean of data.
+            std: the std of data.
+            namespace: the namespace of the registered buffer.
+        """
+        for k, v in [("mean", mean), ("std", std)]:
+            if isinstance(v, str):
+                if v.endswith(".npy"):
+                    v = torch.from_numpy(np.load(v))
+                elif v.endswith(".pt"):
+                    v = torch.load(v)
+                else:
+                    raise ValueError(f"Unsupported file type {v.split('.')[-1]}.")
+            else:
+                v = torch.tensor(v)
+            self.register_buffer(f"{namespace}_{k}", v.float().to(self.device))

algorithms/common/metrics/__init__.py

@@ -0,0 +1,3 @@
+from .fid import FrechetInceptionDistance
+from .lpips import LearnedPerceptualImagePatchSimilarity
+from .fvd import FrechetVideoDistance

algorithms/common/metrics/fid.py

@@ -0,0 +1 @@
+from torchmetrics.image.fid import FrechetInceptionDistance

algorithms/common/metrics/fvd.py

@@ -0,0 +1,158 @@
+"""
+Adopted from https://github.com/cvpr2022-stylegan-v/stylegan-v
+Verified to be the same as tf version by https://github.com/universome/fvd-comparison
+"""
+
+import io
+import re
+import requests
+import html
+import hashlib
+import urllib
+import urllib.request
+from typing import Any, List, Tuple, Union, Dict
+import scipy
+
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+def open_url(
+    url: str,
+    num_attempts: int = 10,
+    verbose: bool = True,
+    return_filename: bool = False,
+) -> Any:
+    """Download the given URL and return a binary-mode file object to access the data."""
+    assert num_attempts >= 1
+
+    # Doesn't look like an URL scheme so interpret it as a local filename.
+    if not re.match("^[a-z]+://", url):
+        return url if return_filename else open(url, "rb")
+
+    # Handle file URLs.  This code handles unusual file:// patterns that
+    # arise on Windows:
+    #
+    # file:///c:/foo.txt
+    #
+    # which would translate to a local '/c:/foo.txt' filename that's
+    # invalid.  Drop the forward slash for such pathnames.
+    #
+    # If you touch this code path, you should test it on both Linux and
+    # Windows.
+    #
+    # Some internet resources suggest using urllib.request.url2pathname() but
+    # but that converts forward slashes to backslashes and this causes
+    # its own set of problems.
+    if url.startswith("file://"):
+        filename = urllib.parse.urlparse(url).path
+        if re.match(r"^/[a-zA-Z]:", filename):
+            filename = filename[1:]
+        return filename if return_filename else open(filename, "rb")
+
+    url_md5 = hashlib.md5(url.encode("utf-8")).hexdigest()
+
+    # Download.
+    url_name = None
+    url_data = None
+    with requests.Session() as session:
+        if verbose:
+            print("Downloading %s ..." % url, end="", flush=True)
+        for attempts_left in reversed(range(num_attempts)):
+            try:
+                with session.get(url) as res:
+                    res.raise_for_status()
+                    if len(res.content) == 0:
+                        raise IOError("No data received")
+
+                    if len(res.content) < 8192:
+                        content_str = res.content.decode("utf-8")
+                        if "download_warning" in res.headers.get("Set-Cookie", ""):
+                            links = [
+                                html.unescape(link)
+                                for link in content_str.split('"')
+                                if "export=download" in link
+                            ]
+                            if len(links) == 1:
+                                url = requests.compat.urljoin(url, links[0])
+                                raise IOError("Google Drive virus checker nag")
+                        if "Google Drive - Quota exceeded" in content_str:
+                            raise IOError(
+                                "Google Drive download quota exceeded -- please try again later"
+                            )
+
+                    match = re.search(
+                        r'filename="([^"]*)"',
+                        res.headers.get("Content-Disposition", ""),
+                    )
+                    url_name = match[1] if match else url
+                    url_data = res.content
+                    if verbose:
+                        print(" done")
+                    break
+            except KeyboardInterrupt:
+                raise
+            except:
+                if not attempts_left:
+                    if verbose:
+                        print(" failed")
+                    raise
+                if verbose:
+                    print(".", end="", flush=True)
+
+    # Return data as file object.
+    assert not return_filename
+    return io.BytesIO(url_data)
+
+
+def compute_fvd(feats_fake: np.ndarray, feats_real: np.ndarray) -> float:
+    mu_gen, sigma_gen = compute_stats(feats_fake)
+    mu_real, sigma_real = compute_stats(feats_real)
+
+    m = np.square(mu_gen - mu_real).sum()
+    s, _ = scipy.linalg.sqrtm(
+        np.dot(sigma_gen, sigma_real), disp=False
+    )  # pylint: disable=no-member
+    fid = np.real(m + np.trace(sigma_gen + sigma_real - s * 2))
+
+    return float(fid)
+
+
+def compute_stats(feats: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    mu = feats.mean(axis=0)  # [d]
+    sigma = np.cov(feats, rowvar=False)  # [d, d]
+
+    return mu, sigma
+
+
+class FrechetVideoDistance(nn.Module):
+    def __init__(self):
+        super().__init__()
+        detector_url = (
+            "https://www.dropbox.com/s/ge9e5ujwgetktms/i3d_torchscript.pt?dl=1"
+        )
+        # Return raw features before the softmax layer.
+        self.detector_kwargs = dict(rescale=False, resize=True, return_features=True)
+        with open_url(detector_url, verbose=False) as f:
+            self.detector = torch.jit.load(f).eval()
+
+    @torch.no_grad()
+    def compute(self, videos_fake: torch.Tensor, videos_real: torch.Tensor):
+        """
+        :param videos_fake: predicted video tensor of shape (frame, batch, channel, height, width)
+        :param videos_real: ground-truth observation tensor of shape (frame, batch, channel, height, width)
+        :return:
+        """
+        n_frames, batch_size, c, h, w = videos_fake.shape
+        if n_frames < 2:
+            raise ValueError("Video must have more than 1 frame for FVD")
+
+        videos_fake = videos_fake.permute(1, 2, 0, 3, 4).contiguous()
+        videos_real = videos_real.permute(1, 2, 0, 3, 4).contiguous()
+
+        # detector takes in tensors of shape [batch_size, c, video_len, h, w] with range -1 to 1
+        feats_fake = self.detector(videos_fake, **self.detector_kwargs).cpu().numpy()
+        feats_real = self.detector(videos_real, **self.detector_kwargs).cpu().numpy()
+
+        return compute_fvd(feats_fake, feats_real)

algorithms/common/metrics/lpips.py

@@ -0,0 +1 @@
+from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity

algorithms/worldmem/__init__.py

@@ -0,0 +1 @@
+from .dememwm_memory_dit import DeMemWMMinecraft, DeMemWMMemoryDiTMinecraft

algorithms/worldmem/dememwm/__init__.py

@@ -0,0 +1,18 @@
+
+from .types import MemoryRecord, MemorySourceType, MemoryStreamTensors, RevisitRetrievalResult, StreamGateState
+from .memory import CausalMemoryBank, MemoryBankQuery, stack_record_tokens
+from .compression import CausalConv3DDynamicCompressor, latent_patch_tokens, spatial_pool_tokens
+from .retrieval import deterministic_revisit_retrieval
+from .schedules import compute_stream_gates, CurriculumState, resolve_curriculum, DeMemWMCurriculumState, resolve_dememwm_curriculum
+from .gates import RevisitRawGate
+from .cache import StreamingCache, DeMemWMStreamingCache
+from .injection import InjectionAdapter, DeMemWMInjectionAdapter
+
+__all__ = [
+    "MemoryRecord", "MemorySourceType", "MemoryStreamTensors", "RevisitRetrievalResult", "StreamGateState",
+    "CausalMemoryBank", "MemoryBankQuery", "stack_record_tokens",
+    "CausalConv3DDynamicCompressor", "latent_patch_tokens", "spatial_pool_tokens",
+    "deterministic_revisit_retrieval", "compute_stream_gates", "CurriculumState", "resolve_curriculum",
+    "DeMemWMCurriculumState", "resolve_dememwm_curriculum", "RevisitRawGate",
+    "StreamingCache", "DeMemWMStreamingCache", "InjectionAdapter", "DeMemWMInjectionAdapter",
+]

algorithms/worldmem/dememwm/cache.py

@@ -0,0 +1,513 @@
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Iterable, Optional
+import warnings
+
+import torch
+
+from .memory import CausalMemoryBank
+from .types import MemoryRecord
+
+
+@dataclass
+class _RawLatentSegment:
+    latents: torch.Tensor
+    frame_indices: torch.Tensor
+    source_is_generated: torch.Tensor
+    pose: Optional[torch.Tensor]
+
+
+class StreamingCache:
+    """Per-video DeMemWM streaming cache with strict no-eviction semantics.
+
+    The cache is intentionally allowed to grow for the current video. It stores
+    detached CPU (or pinned CPU) raw latents plus compressed MemoryRecord objects,
+    while DiT readout tensors remain bounded by the caller's manual budgets.
+    """
+
+    def __init__(
+        self,
+        *,
+        enabled: bool = True,
+        device: str = "cpu",
+        keep_raw_latents: str = "all",
+        keep_compressed_records: bool = True,
+        keep_prefix_anchors: bool = True,
+        eviction_policy: str = "none",
+        no_evict: bool = True,
+        clear_between_videos: bool = True,
+        max_records: Optional[int] = None,
+        max_slots: Optional[int] = None,
+        on_capacity_exceeded: str = "warn",
+    ) -> None:
+        self.enabled = bool(enabled)
+        self.device = str(device or "cpu")
+        self.keep_raw_latents = keep_raw_latents
+        self.keep_compressed_records = bool(keep_compressed_records)
+        self.keep_prefix_anchors = bool(keep_prefix_anchors)
+        self.eviction_policy = str(eviction_policy or "none")
+        self.no_evict = bool(no_evict)
+        self.clear_between_videos = bool(clear_between_videos)
+        self.max_records = max_records
+        self.max_slots = max_slots
+        self.on_capacity_exceeded = str(on_capacity_exceeded or "warn")
+        if self.eviction_policy != "none" or not self.no_evict:
+            raise ValueError("DeMemWMStreamingCache only supports eviction_policy='none' with no_evict=true")
+        if self.device not in {"cpu", "pinned_cpu", "cuda"}:
+            raise ValueError("cache.device must be one of: cpu, pinned_cpu, cuda")
+        self.reset_count = 0
+        self.evictions = 0
+        self.capacity_exceeded_count = 0
+        self.current_video_id: Any = None
+        self._raw_segments: list[_RawLatentSegment] = []
+        self._records: dict[str, dict[int, list[MemoryRecord]]] = {"anchor": {}, "revisit": {}}
+        self._raw_keys: set[tuple[int, int]] = set()
+        self._raw_index: dict[tuple[int, int], tuple[int, int]] = {}
+        self._record_keys: set[tuple[str, int, str, int, int, bool]] = set()
+        self._batch_size: Optional[int] = None
+        # Concat cache: avoids repeated torch.cat across DDIM steps within one chunk.
+        # Invalidated whenever new raw segments are added.
+        self._raw_concat_version: int = 0
+        self._raw_concat_built: int = -1
+        self._raw_concat_cache: Optional[tuple] = None  # (latents, frame_indices, generated, pose)
+        # GPU memory-bank cache: avoids repeated CPU→GPU record transfers across DDIM steps.
+        # Invalidated whenever new records are added.
+        self._banks_version: int = 0
+        self._banks_built_cache: dict[tuple, tuple[int, list[CausalMemoryBank]]] = {}
+
+    @classmethod
+    def from_config(cls, cfg: Any, *, enabled_default: bool = True) -> "StreamingCache":
+        def get(name: str, default: Any) -> Any:
+            return getattr(cfg, name, default) if cfg is not None else default
+
+        return cls(
+            enabled=bool(get("enabled", enabled_default)),
+            device=str(get("device", "cpu")),
+            keep_raw_latents=str(get("keep_raw_latents", "all")),
+            keep_compressed_records=bool(get("keep_compressed_records", True)),
+            keep_prefix_anchors=bool(get("keep_prefix_anchors", True)),
+            eviction_policy=str(get("eviction_policy", "none")),
+            no_evict=bool(get("no_evict", True)),
+            clear_between_videos=bool(get("clear_between_videos", True)),
+            max_records=get("max_records", None),
+            max_slots=get("max_slots", None),
+            on_capacity_exceeded=str(get("on_capacity_exceeded", "warn")),
+        )
+
+    @property
+    def batch_size(self) -> int:
+        return int(self._batch_size or 0)
+
+    @property
+    def raw_segment_count(self) -> int:
+        return len(self._raw_segments)
+
+    @property
+    def raw_frame_slots(self) -> int:
+        return sum(int(seg.latents.shape[0] * seg.latents.shape[1]) for seg in self._raw_segments)
+
+    @property
+    def record_count(self) -> int:
+        return sum(len(records) for by_batch in self._records.values() for records in by_batch.values())
+
+    @property
+    def slot_count(self) -> int:
+        return sum(record.valid_slots for by_batch in self._records.values() for records in by_batch.values() for record in records)
+
+    def records_count(self, kind: str | None = None) -> int:
+        if kind is None:
+            return self.record_count
+        return sum(len(records) for records in self._records.get(kind, {}).values())
+
+    def reset(self, video_id: Any = None) -> None:
+        self.current_video_id = video_id
+        self._raw_segments.clear()
+        self._records = {"anchor": {}, "revisit": {}}
+        self._raw_keys.clear()
+        self._raw_index.clear()
+        self._record_keys.clear()
+        self._batch_size = None
+        self.evictions = 0
+        self.capacity_exceeded_count = 0
+        self.reset_count += 1
+        self._raw_concat_version += 1
+        self._raw_concat_built = -1
+        self._raw_concat_cache = None
+        self._banks_version += 1
+        self._banks_built_cache.clear()
+
+    def _store_tensor(self, tensor: Optional[torch.Tensor], *, dtype: torch.dtype | None = None) -> Optional[torch.Tensor]:
+        if tensor is None:
+            return None
+        out = tensor.detach()
+        if dtype is not None and out.is_floating_point():
+            out = out.to(dtype=dtype)
+        if self.device in {"cpu", "pinned_cpu"}:
+            out = out.to(device="cpu", copy=True)
+            if self.device == "pinned_cpu":
+                try:
+                    out = out.pin_memory()
+                except RuntimeError:
+                    # Keep stable CPU behavior if pinning is unavailable in a worker/process.
+                    pass
+        elif self.device == "cuda":
+            out = out.clone()
+        return out
+
+    def _metadata_to_storage(self, metadata: dict) -> dict:
+        out = {}
+        for key, value in dict(metadata or {}).items():
+            if torch.is_tensor(value):
+                out[key] = self._store_tensor(value)
+            elif isinstance(value, dict):
+                out[key] = self._metadata_to_storage(value)
+            else:
+                out[key] = value
+        return out
+
+    def _metadata_to_device(self, metadata: dict, *, device: torch.device, dtype: torch.dtype) -> dict:
+        out = {}
+        for key, value in dict(metadata or {}).items():
+            if torch.is_tensor(value):
+                tensor = value.to(device=device)
+                out[key] = tensor.to(dtype=dtype) if tensor.is_floating_point() else tensor
+            elif isinstance(value, dict):
+                out[key] = self._metadata_to_device(value, device=device, dtype=dtype)
+            else:
+                out[key] = value
+        return out
+
+    def _record_to_storage(self, record: MemoryRecord) -> MemoryRecord:
+        return MemoryRecord(
+            tokens=self._store_tensor(record.tokens),
+            mask=self._store_tensor(record.mask),
+            source_start=int(record.source_start),
+            source_end=int(record.source_end),
+            frame_indices=self._store_tensor(record.frame_indices),
+            pose=self._store_tensor(record.pose),
+            source_type=record.source_type,
+            is_generated=bool(record.is_generated),
+            score=None if record.score is None or not torch.is_tensor(record.score) else self._store_tensor(record.score),
+            chunk_id=record.chunk_id,
+            metadata=self._metadata_to_storage(record.metadata),
+        )
+
+    def _record_to_device(self, record: MemoryRecord, *, device: torch.device, dtype: torch.dtype) -> MemoryRecord:
+        return MemoryRecord(
+            tokens=record.tokens.to(device=device, dtype=dtype),
+            mask=record.mask.to(device=device, dtype=torch.bool),
+            source_start=int(record.source_start),
+            source_end=int(record.source_end),
+            frame_indices=record.frame_indices.to(device=device),
+            pose=None if record.pose is None else record.pose.to(device=device),
+            source_type=record.source_type,
+            is_generated=bool(record.is_generated),
+            score=record.score,
+            chunk_id=record.chunk_id,
+            metadata=self._metadata_to_device(record.metadata, device=device, dtype=dtype),
+        )
+
+    def _check_capacity(self) -> None:
+        exceeded = False
+        if self.max_records is not None and self.record_count > int(self.max_records):
+            exceeded = True
+        if self.max_slots is not None and self.slot_count > int(self.max_slots):
+            exceeded = True
+        if not exceeded:
+            return
+        self.capacity_exceeded_count += 1
+        msg = (
+            "DeMemWMStreamingCache capacity exceeded "
+            f"records={self.record_count}/{self.max_records}, slots={self.slot_count}/{self.max_slots}; "
+            "no eviction performed because no_evict=true"
+        )
+        if self.on_capacity_exceeded == "error":
+            raise RuntimeError(msg)
+        if self.on_capacity_exceeded == "warn":
+            warnings.warn(msg, RuntimeWarning, stacklevel=2)
+
+    def add_raw_latents(
+        self,
+        latents: torch.Tensor,
+        frame_indices: torch.Tensor,
+        source_is_generated: Optional[torch.Tensor] = None,
+        pose: Optional[torch.Tensor] = None,
+    ) -> None:
+        if not self.enabled or self.keep_raw_latents != "all":
+            return
+        if latents.ndim != 5:
+            raise ValueError("cached raw latents must have shape (T,B,C,H,W)")
+        T, B = int(latents.shape[0]), int(latents.shape[1])
+        if frame_indices.shape != (T, B):
+            raise ValueError("cached frame_indices must have shape (T,B)")
+        if self._batch_size is None:
+            self._batch_size = B
+        elif self._batch_size != B:
+            raise ValueError("streaming cache batch size changed within a video")
+        keep_positions: list[int] = []
+        frame_cpu = frame_indices.detach().cpu()
+        for t in range(T):
+            keys = [(b, int(frame_cpu[t, b].item())) for b in range(B)]
+            if any(key not in self._raw_keys for key in keys):
+                keep_positions.append(t)
+                self._raw_keys.update(keys)
+        if not keep_positions:
+            return
+        pos = torch.as_tensor(keep_positions, dtype=torch.long)
+        seg_latents = latents.index_select(0, pos.to(device=latents.device))
+        seg_frames = frame_indices.index_select(0, pos.to(device=frame_indices.device))
+        if source_is_generated is None:
+            seg_generated = torch.zeros(seg_frames.shape, device=seg_frames.device, dtype=torch.bool)
+        else:
+            seg_generated = source_is_generated.index_select(0, pos.to(device=source_is_generated.device)).bool()
+        seg_pose = None if pose is None else pose.index_select(0, pos.to(device=pose.device))
+        segment_idx = len(self._raw_segments)
+        self._raw_segments.append(
+            _RawLatentSegment(
+                latents=self._store_tensor(seg_latents),
+                frame_indices=self._store_tensor(seg_frames),
+                source_is_generated=self._store_tensor(seg_generated),
+                pose=self._store_tensor(seg_pose),
+            )
+        )
+        for local_pos, source_pos in enumerate(keep_positions):
+            for b in range(B):
+                key = (b, int(frame_cpu[source_pos, b].item()))
+                self._raw_index.setdefault(key, (segment_idx, local_pos))
+        # Invalidate the concat cache — new segment was added.
+        self._raw_concat_version += 1
+        self._raw_concat_cache = None
+
+    def add_records(self, kind: str, batch_idx: int, records: Iterable[MemoryRecord]) -> None:
+        if not self.enabled or not self.keep_compressed_records:
+            return
+        if kind not in self._records:
+            raise ValueError(f"unsupported cache record kind: {kind}")
+        batch_idx = int(batch_idx)
+        bucket = self._records[kind].setdefault(batch_idx, [])
+        added_any = False
+        for record in records:
+            if kind == "anchor" and not self.keep_prefix_anchors:
+                continue
+            key = (
+                kind,
+                batch_idx,
+                str(record.chunk_id or ""),
+                int(record.source_start),
+                int(record.source_end),
+                bool(record.is_generated),
+            )
+            if key in self._record_keys:
+                continue
+            self._record_keys.add(key)
+            bucket.append(self._record_to_storage(record))
+            added_any = True
+        if added_any:
+            # Invalidate the GPU banks cache — new records were added.
+            self._banks_version += 1
+            self._banks_built_cache.clear()
+        self._check_capacity()
+
+    def add_memory_banks(self, anchor_banks: list[CausalMemoryBank], revisit_banks: list[CausalMemoryBank]) -> None:
+        for batch_idx, bank in enumerate(anchor_banks):
+            self.add_records("anchor", batch_idx, bank.records)
+        for batch_idx, bank in enumerate(revisit_banks):
+            self.add_records("revisit", batch_idx, bank.records)
+
+    def memory_banks(self, kind: str, *, device: torch.device, dtype: torch.dtype, batch_size: int | None = None) -> list[CausalMemoryBank]:
+        if kind not in self._records:
+            raise ValueError(f"unsupported cache record kind: {kind}")
+        B = int(batch_size or self.batch_size or (max(self._records[kind].keys()) + 1 if self._records[kind] else 0))
+        cache_key = (kind, device, dtype, B)
+        cached = self._banks_built_cache.get(cache_key)
+        if cached is not None and cached[0] == self._banks_version:
+            return cached[1]
+        banks: list[CausalMemoryBank] = []
+        for batch_idx in range(B):
+            bank = CausalMemoryBank()
+            for record in self._records[kind].get(batch_idx, []):
+                bank.add_record(self._record_to_device(record, device=device, dtype=dtype))
+            banks.append(bank)
+        self._banks_built_cache[cache_key] = (self._banks_version, banks)
+        return banks
+
+    def records_for_batch(self, kind: str, batch_idx: int) -> tuple[MemoryRecord, ...]:
+        if kind not in self._records:
+            raise ValueError(f"unsupported cache record kind: {kind}")
+        return tuple(self._records[kind].get(int(batch_idx), ()))
+
+    def raw_latents_for_frames(
+        self,
+        *,
+        batch_idx: int,
+        frame_indices: torch.Tensor,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        frames = frame_indices.detach().cpu().reshape(-1)
+        rows = []
+        batch_idx = int(batch_idx)
+        for frame in frames.tolist():
+            key = (batch_idx, int(frame))
+            location = self._raw_index.get(key)
+            if location is None:
+                raise KeyError(f"raw latent for batch={batch_idx}, frame={int(frame)} is not cached")
+            segment_idx, local_pos = location
+            rows.append(self._raw_segments[segment_idx].latents[local_pos, batch_idx])
+        if not rows:
+            template = self._raw_segments[0].latents
+            return template[:0, batch_idx:batch_idx + 1].to(device=device, dtype=dtype)
+        return torch.stack(rows, dim=0).unsqueeze(1).to(device=device, dtype=dtype)
+
+    def _select_time_positions(
+        self,
+        frame_indices: torch.Tensor,
+        target_frame_indices: Optional[torch.Tensor],
+        max_recent_frames: Optional[int],
+        exclude_latest_local_frames: int = 0,
+    ) -> torch.Tensor:
+        T, B = frame_indices.shape
+        if target_frame_indices is None or max_recent_frames is None or int(max_recent_frames) <= 0:
+            return torch.arange(T, dtype=torch.long)
+        targets = target_frame_indices.detach().cpu()
+        if targets.ndim == 1:
+            targets = targets[:, None].expand(-1, B)
+        frames = frame_indices.detach().cpu()  # (T, B)
+        recent = int(max_recent_frames)
+        exclude = max(0, int(exclude_latest_local_frames))
+        # Vectorized: valid[t_tgt, t_src, b] = True if source position t_src is
+        # causally valid for target t_tgt in batch b.
+        # frames (T, B) → (1, T, B); targets (T_tgt, B) → (T_tgt, 1, B)
+        valid = frames.unsqueeze(0) < (targets.unsqueeze(1) - exclude)  # (T_tgt, T, B)
+        # For each (t_tgt, b), retain only the last `recent` valid positions.
+        # Flip T, cumsum along T (counting from the end), keep where ≤ recent.
+        valid_f = valid.flip(1)
+        keep_f = (valid_f.long().cumsum(1) <= recent) & valid_f
+        # Any position needed by any (t_tgt, b) pair.
+        keep_any = keep_f.flip(1).any(dim=0).any(dim=1)  # (T,)
+        return keep_any.nonzero(as_tuple=False).flatten()
+
+    def materialize_raw_latents(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        max_recent_frames: Optional[int] = None,
+        target_frame_indices: Optional[torch.Tensor] = None,
+        exclude_latest_local_frames: int = 0,
+    ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
+        if not self._raw_segments:
+            return None, None, None, None
+        if target_frame_indices is not None and max_recent_frames is not None and int(max_recent_frames) > 0:
+            return self._materialize_recent_raw_latents(
+                device=device,
+                dtype=dtype,
+                max_recent_frames=int(max_recent_frames),
+                target_frame_indices=target_frame_indices,
+                exclude_latest_local_frames=exclude_latest_local_frames,
+            )
+        # Rebuild the concatenated CPU tensors only when new segments were added.
+        if self._raw_concat_cache is None or self._raw_concat_built != self._raw_concat_version:
+            latents = torch.cat([seg.latents for seg in self._raw_segments], dim=0)
+            frame_indices = torch.cat([seg.frame_indices for seg in self._raw_segments], dim=0)
+            generated = torch.cat([seg.source_is_generated for seg in self._raw_segments], dim=0)
+            pose: Optional[torch.Tensor] = None
+            if all(seg.pose is not None for seg in self._raw_segments):
+                pose = torch.cat([seg.pose for seg in self._raw_segments if seg.pose is not None], dim=0)
+            self._raw_concat_cache = (latents, frame_indices, generated, pose)
+            self._raw_concat_built = self._raw_concat_version
+        else:
+            latents, frame_indices, generated, pose = self._raw_concat_cache
+        pos = self._select_time_positions(frame_indices, target_frame_indices, max_recent_frames, exclude_latest_local_frames)
+        if pos.numel() == 0:
+            empty_latents = latents[:0].to(device=device, dtype=dtype)
+            empty_frames = frame_indices[:0].to(device=device)
+            empty_generated = generated[:0].to(device=device, dtype=torch.bool)
+            empty_pose = None if pose is None else pose[:0].to(device=device)
+            return empty_latents, empty_frames, empty_generated, empty_pose
+        latents = latents.index_select(0, pos.to(device=latents.device)).to(device=device, dtype=dtype)
+        frame_indices = frame_indices.index_select(0, pos.to(device=frame_indices.device)).to(device=device)
+        generated = generated.index_select(0, pos.to(device=generated.device)).to(device=device, dtype=torch.bool)
+        if pose is not None:
+            pose = pose.index_select(0, pos.to(device=pose.device)).to(device=device)
+        return latents, frame_indices, generated, pose
+
+    def _materialize_recent_raw_latents(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        max_recent_frames: int,
+        target_frame_indices: torch.Tensor,
+        exclude_latest_local_frames: int = 0,
+    ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
+        B = self.batch_size
+        targets = target_frame_indices.detach().cpu()
+        if targets.ndim == 1:
+            targets = targets[:, None].expand(-1, B)
+        elif targets.shape[1] == 1 and B > 1:
+            targets = targets.expand(-1, B)
+        if targets.shape[1] != B:
+            raise ValueError("target_frame_indices batch dimension does not match streaming cache")
+
+        recent = max(0, int(max_recent_frames))
+        exclude = max(0, int(exclude_latest_local_frames))
+        counts = torch.zeros(targets.shape, dtype=torch.long)
+        selected: list[tuple[_RawLatentSegment, int]] = []
+
+        for segment in reversed(self._raw_segments):
+            frames = segment.frame_indices.detach().cpu()
+            for local_pos in range(frames.shape[0] - 1, -1, -1):
+                valid = frames[local_pos].unsqueeze(0) < (targets - exclude)
+                needed = valid & (counts < recent)
+                if not needed.any():
+                    continue
+                selected.append((segment, local_pos))
+                counts += needed.long()
+                if bool((counts >= recent).all().item()):
+                    break
+            if bool((counts >= recent).all().item()):
+                break
+
+        if not selected:
+            template = self._raw_segments[0]
+            empty_latents = template.latents[:0].to(device=device, dtype=dtype)
+            empty_frames = template.frame_indices[:0].to(device=device)
+            empty_generated = template.source_is_generated[:0].to(device=device, dtype=torch.bool)
+            empty_pose = None if template.pose is None else template.pose[:0].to(device=device)
+            return empty_latents, empty_frames, empty_generated, empty_pose
+
+        selected.reverse()
+        latents = torch.stack([segment.latents[local_pos] for segment, local_pos in selected], dim=0).to(device=device, dtype=dtype)
+        frame_indices = torch.stack([segment.frame_indices[local_pos] for segment, local_pos in selected], dim=0).to(device=device)
+        generated = torch.stack([segment.source_is_generated[local_pos] for segment, local_pos in selected], dim=0).to(device=device, dtype=torch.bool)
+        pose = None
+        if all(segment.pose is not None for segment, _ in selected):
+            pose = torch.stack(
+                [segment.pose[local_pos] for segment, local_pos in selected if segment.pose is not None],
+                dim=0,
+            ).to(device=device)
+        return latents, frame_indices, generated, pose
+
+    def diagnostics(self, prefix: str = "cache") -> dict[str, Any]:
+        return {
+            f"{prefix}_enabled": bool(self.enabled),
+            f"{prefix}_records": int(self.record_count),
+            f"{prefix}_anchor_records": int(self.records_count("anchor")),
+            f"{prefix}_revisit_records": int(self.records_count("revisit")),
+            f"{prefix}_slots": int(self.slot_count),
+            f"{prefix}_raw_frame_slots": int(self.raw_frame_slots),
+            f"{prefix}_raw_segments": int(self.raw_segment_count),
+            f"{prefix}_evictions": int(self.evictions),
+            f"{prefix}_resets": int(self.reset_count),
+            f"{prefix}_capacity_exceeded": int(self.capacity_exceeded_count),
+            f"{prefix}_device": self.device,
+            f"{prefix}_current_video_id": self.current_video_id,
+            f"{prefix}_clear_between_videos": bool(self.clear_between_videos),
+            f"{prefix}_no_evict": bool(self.no_evict),
+        }
+
+
+DeMemWMStreamingCache = StreamingCache

algorithms/worldmem/dememwm/compression.py

@@ -0,0 +1,260 @@
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+def latent_patch_tokens(latents: torch.Tensor, patch_size: int) -> torch.Tensor:
+    if latents.ndim != 5:
+        raise ValueError("latents must have shape (T,B,C,H,W)")
+    if patch_size <= 0:
+        raise ValueError("patch_size must be positive")
+    T, B, C, H, W = latents.shape
+    if H % patch_size != 0 or W % patch_size != 0:
+        raise ValueError(f"latent H,W=({H},{W}) must be divisible by patch_size={patch_size}")
+    flat = latents.reshape(T * B, C, H, W)
+    patches = F.unfold(flat, kernel_size=patch_size, stride=patch_size).transpose(1, 2).contiguous()
+    return patches.reshape(T, B, patches.shape[1], C * patch_size * patch_size)
+
+
+def spatial_pool_tokens(
+    tokens: torch.Tensor,
+    pool_h: int,
+    pool_w: int,
+    src_h: int,
+    src_w: int,
+) -> torch.Tensor:
+    """2D adaptive average pool on a flattened (src_h*src_w, D) token grid.
+    Preserves 2D spatial layout. Returns (pool_h*pool_w, D)."""
+    if tokens.ndim != 2:
+        raise ValueError("tokens must have shape (N, D)")
+    D = tokens.shape[-1]
+    spatial = tokens.reshape(src_h, src_w, D).permute(2, 0, 1).unsqueeze(0)
+    pooled = F.adaptive_avg_pool2d(spatial, (pool_h, pool_w))
+    return pooled.squeeze(0).permute(1, 2, 0).reshape(-1, D)
+
+
+class SpatialConv2DMemoryProjector(nn.Module):
+    """Project latent maps to DiT hidden tokens while preserving the HxW grid."""
+
+    projects_spatial_latents = True
+
+    def __init__(
+        self,
+        latent_channels: int,
+        dit_hidden_size: int,
+        mid_channels: int,
+        kernel_size: int = 3,
+    ):
+        super().__init__()
+        kernel_size = int(kernel_size)
+        if kernel_size <= 0 or kernel_size % 2 == 0:
+            raise ValueError("kernel_size must be a positive odd integer")
+        self.latent_channels = int(latent_channels)
+        self.dit_hidden_size = int(dit_hidden_size)
+        self.mid_channels = int(mid_channels)
+        self.kernel_size = kernel_size
+        self.out_features = self.dit_hidden_size
+        self.proj_in = nn.Conv2d(self.latent_channels, self.mid_channels, kernel_size=1)
+        self.proj_spatial = nn.Conv2d(
+            self.mid_channels,
+            self.dit_hidden_size,
+            kernel_size=kernel_size,
+            padding=kernel_size // 2,
+        )
+
+    def forward(self, latents: torch.Tensor) -> torch.Tensor:
+        if latents.ndim != 5:
+            raise ValueError("latents must have shape (T,B,C,H,W)")
+        T, B, C, H, W = latents.shape
+        if C != self.latent_channels:
+            raise ValueError(f"expected {self.latent_channels} latent channels, got {C}")
+        x = latents.reshape(T * B, C, H, W)
+        x = self.proj_spatial(self.proj_in(x))
+        x = x.reshape(T, B, self.dit_hidden_size, H, W)
+        return x.permute(1, 0, 3, 4, 2).reshape(B, T, H * W, self.dit_hidden_size).contiguous()
+
+
+class CausalConv3DDynamicCompressor(nn.Module):
+    """Dynamic memory compressor: delta preprocessing + causal Conv3D on raw latents.
+
+    Replaces ShortTermLatentCompressor (slot cross-attention).
+    - Operates directly on (T, C, H, W) raw latents
+    - Delta: inp[0]=latent[0], inp[t]=latent[t]-latent[t-1]
+    - Causal padding prepends temporal zeros and right-aligns fixed outputs
+    - Zero-padded to max_source_frames for fixed output shape
+    - No slot cross-attention, no chunking
+    """
+
+    def __init__(
+        self,
+        latent_channels: int,
+        dit_hidden_size: int,
+        patch_size: int = 2,
+        conv_kernel_t: int = 3,
+        conv_stride_t: int = 2,
+        max_source_frames: int = 8,
+        exclude_latest_local_frames: int = 4,
+    ):
+        super().__init__()
+        self.latent_channels = latent_channels
+        self.dit_hidden_size = dit_hidden_size
+        self.patch_size = patch_size
+        self.conv_kernel_t = conv_kernel_t
+        self.conv_stride_t = conv_stride_t
+        self.max_source_frames = max_source_frames
+        self.exclude_latest_local_frames = int(exclude_latest_local_frames)
+        self.causal_pad = self._temporal_left_pad()
+        self.conv3d = nn.Conv3d(
+            latent_channels, dit_hidden_size,
+            kernel_size=(conv_kernel_t, patch_size, patch_size),
+            stride=(conv_stride_t, patch_size, patch_size),
+            padding=0,
+        )
+        self.out_norm = nn.LayerNorm(dit_hidden_size)
+        self._init_temporal_as_delta()
+
+    def _init_temporal_as_delta(self) -> None:
+        with torch.no_grad():
+            self.conv3d.weight.zero_()
+            k_t, p = self.conv_kernel_t, self.patch_size
+            D_out, D_in = self.conv3d.weight.shape[:2]
+            scale = 1.0 / (p * p)
+            # Delta preprocessing happens in forward. Initialize every output
+            # channel to read a patch-averaged current delta, repeating latent
+            # channels across the wider DiT hidden dimension.
+            for d in range(D_out):
+                self.conv3d.weight[d, d % D_in, k_t - 1, :, :] = scale
+            if self.conv3d.bias is not None:
+                nn.init.zeros_(self.conv3d.bias)
+
+    def _temporal_output_count(self) -> int:
+        return math.ceil(self.max_source_frames / self.conv_stride_t)
+
+    def _temporal_left_pad(self) -> int:
+        t_out = self._temporal_output_count()
+        latest_output_end = (t_out - 1) * self.conv_stride_t + self.conv_kernel_t - 1
+        latest_source = self.max_source_frames - 1
+        return max(0, latest_output_end - latest_source)
+
+    def _output_time_indices(self, device: torch.device) -> torch.Tensor:
+        t_out = self._temporal_output_count()
+        return (
+            torch.arange(t_out, device=device, dtype=torch.long) * self.conv_stride_t
+            + self.conv_kernel_t
+            - 1
+            - self.causal_pad
+        )
+
+    def tokens_per_target(self, H: int, W: int) -> int:
+        p = self.patch_size
+        T_out = self._temporal_output_count()
+        return T_out * (H // p) * (W // p)
+
+    def forward(
+        self,
+        latents: torch.Tensor,
+        frame_indices: torch.Tensor,
+        pose: Optional[torch.Tensor],
+        target_frame_indices: torch.Tensor,
+        source_is_generated: Optional[torch.Tensor] = None,
+        exclude_latest_local_frames: Optional[int] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, dict]:
+        if latents.ndim != 5:
+            raise ValueError("latents must have shape (T_src,B,C,H,W)")
+        exclude_latest_local_frames = (
+            self.exclude_latest_local_frames
+            if exclude_latest_local_frames is None
+            else int(exclude_latest_local_frames)
+        )
+        T_src, B, C, H, W = latents.shape
+        p = self.patch_size
+        if H % p != 0 or W % p != 0:
+            raise ValueError(f"latent H,W=({H},{W}) must be divisible by patch_size={p}")
+        if target_frame_indices.ndim == 1:
+            target_frame_indices = target_frame_indices[:, None].expand(-1, B)
+        T_tgt = target_frame_indices.shape[0]
+        device = latents.device
+        generated_flags = None if source_is_generated is None else source_is_generated.to(device=device, dtype=torch.bool)
+        n_spatial = (H // p) * (W // p)
+        T_out = self._temporal_output_count()
+        num_slots = T_out * n_spatial
+        output_time_idx = self._output_time_indices(device)
+        selected_source_count = torch.zeros((B, T_tgt), dtype=torch.long, device=device)
+        max_source_frame = torch.full((B, T_tgt), -1, dtype=torch.long, device=device)
+        generated_source_fraction = torch.zeros((B, T_tgt), dtype=torch.float32, device=device)
+        min_gap = torch.full((B, T_tgt), -1, dtype=torch.long, device=device)
+        max_gap = torch.full((B, T_tgt), -1, dtype=torch.long, device=device)
+        output_rows, mask_rows = [], []
+        for b in range(B):
+            src_frames_b = frame_indices[:, b]
+            tgt_outputs, tgt_masks = [], []
+            for j in range(T_tgt):
+                target = int(target_frame_indices[j, b].item())
+                valid_idx = (
+                    src_frames_b < target - exclude_latest_local_frames
+                ).nonzero(as_tuple=False).flatten()
+                if valid_idx.numel() == 0:
+                    tgt_outputs.append(latents.new_zeros(num_slots, self.dit_hidden_size))
+                    tgt_masks.append(torch.zeros(num_slots, device=device, dtype=torch.bool))
+                    continue
+                selected_frames = src_frames_b.index_select(0, valid_idx)
+                order = torch.argsort(selected_frames)
+                valid_idx = valid_idx.index_select(0, order)[-self.max_source_frames:]
+                selected_frames = src_frames_b.index_select(0, valid_idx)
+                selected_source_count[b, j] = int(selected_frames.numel())
+                max_source_frame[b, j] = selected_frames.max()
+                gaps = target - selected_frames
+                min_gap[b, j] = gaps.min()
+                max_gap[b, j] = gaps.max()
+                if generated_flags is not None:
+                    generated = generated_flags.index_select(0, valid_idx)[:, b]
+                    generated_source_fraction[b, j] = generated.float().mean()
+                chunk = latents[valid_idx, b]
+                real_mask = torch.ones((chunk.shape[0],), device=device, dtype=torch.bool)
+                if chunk.shape[0] < self.max_source_frames:
+                    pad = chunk.new_zeros(self.max_source_frames - chunk.shape[0], C, H, W)
+                    chunk = torch.cat([pad, chunk], dim=0)
+                    real_mask = torch.cat([
+                        torch.zeros((pad.shape[0],), device=device, dtype=torch.bool),
+                        real_mask,
+                    ])
+                inp = chunk.clone()
+                inp[1:] = chunk[1:] - chunk[:-1]
+                x = inp.permute(1, 0, 2, 3).unsqueeze(0)        # (1,C,T,H,W)
+                x = F.pad(x, (0, 0, 0, 0, self.causal_pad, 0))  # left-pad time
+                x = self.conv3d(x)                               # (1,D,T_out,H//p,W//p)
+                x = x.squeeze(0).permute(1, 2, 3, 0)            # (T_out,H//p,W//p,D)
+                x = self.out_norm(x)
+                tokens = x.reshape(num_slots, self.dit_hidden_size)
+                clamped_time_idx = output_time_idx.clamp(min=0, max=self.max_source_frames - 1)
+                temporal_mask = (
+                    (output_time_idx >= 0)
+                    & (output_time_idx < self.max_source_frames)
+                    & real_mask.index_select(0, clamped_time_idx)
+                )
+                mask = temporal_mask[:, None].expand(T_out, n_spatial).reshape(num_slots)
+                tgt_outputs.append(tokens)
+                tgt_masks.append(mask)
+            output_rows.append(torch.stack(tgt_outputs))
+            mask_rows.append(torch.stack(tgt_masks))
+        out_tokens = torch.stack(output_rows)
+        out_mask = torch.stack(mask_rows)
+        diagnostics = {
+            "num_dynamic_slots": num_slots,
+            "dynamic_T_out": T_out,
+            "dynamic_n_spatial": n_spatial,
+            "dynamic_temporal_left_pad": self.causal_pad,
+            "dynamic_output_time_indices": output_time_idx,
+            "selected_source_count": selected_source_count,
+            "max_source_frame": max_source_frame,
+            "generated_source_fraction": generated_source_fraction,
+            "dynamic_min_gap_to_target_per_target": min_gap,
+            "dynamic_max_gap_to_target_per_target": max_gap,
+            "dynamic_exclude_latest_local_frames": exclude_latest_local_frames,
+        }
+        return out_tokens, out_mask, diagnostics

algorithms/worldmem/dememwm/diagnostics.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from .schedules import EVAL_ABLATION_BRANCH_TO_ID, NOISE_BUCKETS, NOISE_BUCKET_TO_ID, normalize_eval_ablation_branch, normalize_noise_bucket
+
+
+_REVISIT_LABEL_SOURCE = "deterministic_fov_coverage_plucker"
+
+
+def tensor_valid_fraction(mask: torch.Tensor | None) -> float:
+    if mask is None or mask.numel() == 0:
+        return 0.0
+    return float(mask.detach().bool().float().mean().item())
+
+
+def gate_stats(gate: torch.Tensor | float | int | None) -> dict[str, float]:
+    if gate is None:
+        return {"mean": 0.0, "min": 0.0, "max": 0.0}
+    if not torch.is_tensor(gate):
+        value = float(gate)
+        return {"mean": value, "min": value, "max": value}
+    g = gate.detach().float()
+    return {"mean": float(g.mean().item()), "min": float(g.min().item()), "max": float(g.max().item())}
+
+
+def summarize_stream(name: str, tokens: torch.Tensor | None, mask: torch.Tensor | None, gate: torch.Tensor | float | None) -> dict[str, Any]:
+    return {f"{name}_tokens_shape": None if tokens is None else tuple(tokens.shape), f"{name}_valid_fraction": tensor_valid_fraction(mask), f"{name}_valid_tokens": 0 if mask is None else int(mask.detach().bool().sum().item()), f"{name}_gate": gate_stats(gate)}
+
+
+def assert_no_future_sources(target_frame: int, max_source_frame: int | torch.Tensor) -> None:
+    max_src = int(max_source_frame.detach().max().item()) if torch.is_tensor(max_source_frame) else int(max_source_frame)
+    if max_src >= int(target_frame):
+        raise AssertionError(f"DeMemWM memory source {max_src} is not causal for target {target_frame}")
+
+
+def _collect_values(result_diagnostics: list[dict[str, Any]], key: str) -> list[float]:
+    values: list[float] = []
+    for diag in result_diagnostics:
+        for value in diag.get(key, []) or []:
+            values.append(float(value))
+    return values
+
+
+def _value_stats(values: list[float], prefix: str) -> dict[str, float]:
+    if not values:
+        return {f"{prefix}_mean": 0.0, f"{prefix}_min": 0.0, f"{prefix}_max": 0.0}
+    return {
+        f"{prefix}_mean": float(sum(values) / len(values)),
+        f"{prefix}_min": float(min(values)),
+        f"{prefix}_max": float(max(values)),
+    }
+
+
+def summarize_revisit_diagnostics(result_diagnostics: list[dict[str, Any]], valid_revisit_mask: torch.Tensor | None) -> dict[str, Any]:
+    target_count = len(result_diagnostics)
+    candidate_count = sum(int(diag.get("revisit_candidate_frame_count", diag.get("revisit_candidate_count", diag.get("candidate_count", 0)))) for diag in result_diagnostics)
+    candidate_count_mean = float(candidate_count / target_count) if target_count else 0.0
+    valid_candidate_label_count = sum(int(diag.get("valid_candidate_label_count", diag.get("valid_candidate_count", 0))) for diag in result_diagnostics)
+    pose_preselect_input_count = sum(int(diag.get("revisit_pose_preselect_input_count", 0)) for diag in result_diagnostics)
+    pose_preselect_selected_count = sum(int(diag.get("revisit_pose_preselect_selected_count", 0)) for diag in result_diagnostics)
+    exact_fov_candidate_count = sum(int(diag.get("revisit_exact_fov_candidate_count", 0)) for diag in result_diagnostics)
+    valid_count = sum(int(diag.get("valid_revisit_frame_count", diag.get("valid_revisit_count", diag.get("valid_candidate_count", 0)))) for diag in result_diagnostics)
+    valid_count_mean = float(valid_count / target_count) if target_count else 0.0
+    valid_target_count = sum(int(diag.get("valid_revisit_target_count", diag.get("high_quality_selected_revisit", 0))) for diag in result_diagnostics)
+    selected_count = sum(int(diag.get("revisit_selected_frame_count", diag.get("revisit_selected_count", diag.get("selected_count", 0)))) for diag in result_diagnostics)
+    no_valid_count = sum(int(diag.get("no_valid_revisit_count", 0)) for diag in result_diagnostics)
+    abstained_count = sum(int(diag.get("revisit_abstained_count", int(bool(diag.get("abstained", False))))) for diag in result_diagnostics)
+    selected_gaps = [int(diag["revisit_min_gap_to_target"]) for diag in result_diagnostics if int(diag.get("revisit_min_gap_to_target", -1)) >= 0]
+    diagnostics: dict[str, Any] = {
+        "revisit_candidate_frame_count": candidate_count_mean,
+        "revisit_candidate_count": candidate_count_mean,
+        "valid_candidate_label_count": int(valid_candidate_label_count),
+        "revisit_pose_preselect_input_count": float(pose_preselect_input_count / target_count) if target_count else 0.0,
+        "revisit_pose_preselect_selected_count": float(pose_preselect_selected_count / target_count) if target_count else 0.0,
+        "revisit_exact_fov_candidate_count": float(exact_fov_candidate_count / target_count) if target_count else 0.0,
+        "valid_revisit_frame_count": valid_count_mean,
+        "valid_revisit_count": valid_count_mean,
+        "valid_revisit_target_count": int(valid_target_count),
+        "no_valid_revisit_count": int(no_valid_count),
+        "valid_revisit_mask_fraction": tensor_valid_fraction(valid_revisit_mask),
+        "revisit_selected_frame_count": int(selected_count),
+        "revisit_selected_count": int(selected_count),
+        "revisit_abstained_count": int(abstained_count),
+        "revisit_min_gap_to_target": int(min(selected_gaps)) if selected_gaps else -1,
+        "revisit_label_source": _REVISIT_LABEL_SOURCE,
+    }
+    frame_fov_values = _collect_values(result_diagnostics, "frame_fov_overlap_values")
+    if not frame_fov_values:
+        frame_fov_values = _collect_values(result_diagnostics, "fov_overlap_values")
+    diagnostics.update(_value_stats(frame_fov_values, "revisit_frame_fov_overlap"))
+    diagnostics.update(_value_stats(frame_fov_values, "revisit_fov_overlap"))
+    diagnostics.update(_value_stats(_collect_values(result_diagnostics, "plucker_overlap_values"), "revisit_plucker_overlap"))
+    diagnostics.update(_value_stats(_collect_values(result_diagnostics, "best_selected_fov_overlap_values"), "revisit_best_selected_fov_overlap"))
+    diagnostics.update(_value_stats(_collect_values(result_diagnostics, "best_selected_plucker_overlap_values"), "revisit_best_selected_plucker_overlap"))
+    diagnostics.update(_value_stats(_collect_values(result_diagnostics, "best_selected_gap_frame_values"), "revisit_best_selected_gap_frames"))
+    diagnostics.update(_value_stats(_collect_values(result_diagnostics, "best_selected_frame_fov_overlap_values"), "revisit_best_selected_frame_fov_overlap"))
+    diagnostics.update(_value_stats(_collect_values(result_diagnostics, "selected_frame_fov_overlap_values"), "revisit_selected_frame_fov_overlap"))
+    diagnostics.update(_value_stats(_collect_values(result_diagnostics, "selected_incremental_fov_overlap_values"), "revisit_incremental_fov_overlap"))
+    return diagnostics
+
+
+def summarize_noise_bucket_diagnostics(
+    *,
+    noise_bucket: str | None,
+    valid_revisit_mask: torch.Tensor | None,
+    no_valid_revisit_mask: torch.Tensor | None,
+    noise_bucket_ids: torch.Tensor | None = None,
+) -> dict[str, Any]:
+    bucket = normalize_noise_bucket(noise_bucket)
+    diagnostics: dict[str, Any] = {
+        "noise_bucket": bucket,
+        "noise_bucket_id": int(NOISE_BUCKET_TO_ID[bucket]),
+    }
+    for candidate in NOISE_BUCKETS:
+        diagnostics[f"noise_bucket_is_{candidate}"] = int(bucket == candidate)
+
+    valid = torch.zeros(0, dtype=torch.bool) if valid_revisit_mask is None else valid_revisit_mask.detach().bool().reshape(-1).cpu()
+    no_valid = torch.zeros_like(valid) if no_valid_revisit_mask is None else no_valid_revisit_mask.detach().bool().reshape(-1).cpu()
+    target_count = int(valid.numel())
+    diagnostics["noise_bucket_target_count"] = target_count
+    if noise_bucket_ids is None:
+        target_bucket_ids = torch.full((target_count,), int(NOISE_BUCKET_TO_ID[bucket]), dtype=torch.long)
+    else:
+        target_bucket_ids = noise_bucket_ids.detach().long().reshape(-1).cpu()
+        if int(target_bucket_ids.numel()) != target_count:
+            raise ValueError(
+                f"noise_bucket_ids has {int(target_bucket_ids.numel())} targets, expected {target_count}"
+            )
+
+    for bucket_name in NOISE_BUCKETS:
+        bucket_mask = target_bucket_ids == int(NOISE_BUCKET_TO_ID[bucket_name])
+        diagnostics[f"noise_bucket_{bucket_name}_target_count"] = int(bucket_mask.long().sum().item())
+
+    mask_specs = (
+        ("valid_revisit", valid),
+        ("no_valid_revisit", no_valid),
+    )
+    for mask_name, mask in mask_specs:
+        for bucket_name in NOISE_BUCKETS:
+            bucket_mask = target_bucket_ids == int(NOISE_BUCKET_TO_ID[bucket_name])
+            count = int((mask & bucket_mask).long().sum().item()) if mask.numel() else 0
+            diagnostics[f"{mask_name}_noise_bucket_{bucket_name}_count"] = count
+    return diagnostics
+
+
+def summarize_eval_ablation_diagnostics(
+    *,
+    enabled: bool,
+    branch: str | None,
+    valid_revisit_mask: torch.Tensor | None,
+    no_valid_revisit_mask: torch.Tensor | None,
+    eval_corrupted_revisit_mask: torch.Tensor | None,
+) -> dict[str, Any]:
+    branch = normalize_eval_ablation_branch(branch)
+    valid = torch.zeros(0, dtype=torch.bool) if valid_revisit_mask is None else valid_revisit_mask.detach().bool().reshape(-1).cpu()
+    no_valid = torch.zeros_like(valid) if no_valid_revisit_mask is None else no_valid_revisit_mask.detach().bool().reshape(-1).cpu()
+    corrupted = torch.zeros_like(valid) if eval_corrupted_revisit_mask is None else eval_corrupted_revisit_mask.detach().bool().reshape(-1).cpu()
+    true_revisit = valid & (~corrupted)
+    diagnostics: dict[str, Any] = {
+        "eval_ablation_enabled": bool(enabled),
+        "eval_ablation_branch": branch,
+        "eval_ablation_branch_id": int(EVAL_ABLATION_BRANCH_TO_ID[branch]),
+        "eval_bucket_true_revisit_count": int(true_revisit.long().sum().item()),
+        "eval_bucket_no_valid_revisit_count": int(no_valid.long().sum().item()),
+        "eval_bucket_corrupted_memory_count": int(corrupted.long().sum().item()),
+    }
+    total = max(int(valid.numel()), 1)
+    diagnostics["eval_bucket_true_revisit_fraction"] = float(diagnostics["eval_bucket_true_revisit_count"] / total)
+    diagnostics["eval_bucket_no_valid_revisit_fraction"] = float(diagnostics["eval_bucket_no_valid_revisit_count"] / total)
+    diagnostics["eval_bucket_corrupted_memory_fraction"] = float(diagnostics["eval_bucket_corrupted_memory_count"] / total)
+    return diagnostics

algorithms/worldmem/dememwm/gates.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class RevisitRawGate(nn.Module):
+    """Learned per-target revisit gate from selected-revisit quality features.
+
+    The caller applies validity masking and stage/denoise scaling after this
+    module. This module never turns selected revisit validity into a target.
+    """
+
+    def __init__(self, init_logit: float = 1.0):
+        super().__init__()
+        self.net = nn.Linear(3, 1)
+        nn.init.zeros_(self.net.weight)
+        nn.init.constant_(self.net.bias, float(init_logit))
+
+    def forward(
+        self,
+        *,
+        valid_revisit_mask: torch.Tensor,
+        best_selected_fov_overlap: torch.Tensor | None = None,
+        best_selected_plucker_overlap: torch.Tensor | None = None,
+        selected_gap_frames: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if valid_revisit_mask.ndim != 2:
+            raise ValueError("valid_revisit_mask must have shape (B,T)")
+        device = valid_revisit_mask.device
+        dtype = torch.float32
+        shape = valid_revisit_mask.shape
+
+        def _feature(value: torch.Tensor | None) -> torch.Tensor:
+            if value is None:
+                return torch.zeros(shape, device=device, dtype=dtype)
+            tensor = value.to(device=device, dtype=dtype)
+            if tensor.ndim == 0:
+                return torch.full(shape, float(tensor.item()), device=device, dtype=dtype)
+            return tensor.expand(shape)
+
+        fov = _feature(best_selected_fov_overlap).clamp(min=0.0, max=1.0)
+        plucker = _feature(best_selected_plucker_overlap).clamp(min=0.0, max=1.0)
+        log_age = torch.log1p(_feature(selected_gap_frames).clamp_min(0.0)).clamp(max=8.0) / 8.0
+        features = torch.stack([fov, plucker, log_age], dim=-1)
+        return torch.sigmoid(self.net(features).squeeze(-1))

algorithms/worldmem/dememwm/injection.py

@@ -0,0 +1,83 @@
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from .diagnostics import summarize_stream
+from .types import MemoryStreamTensors
+
+
+@dataclass
+class InjectionAdapter:
+    """Convert DeMemWM stream tensors to Diffusion/DiT Option-C kwargs."""
+
+    omit_disabled: bool = False
+
+    def _tokens(self, name: str, tokens: torch.Tensor, device, dtype) -> torch.Tensor:
+        if tokens.ndim != 4:
+            raise ValueError(f"{name} tokens must have shape (B,T,M,D), got {tuple(tokens.shape)}")
+        return tokens.to(device=device, dtype=dtype)
+
+    def _mask(self, name: str, mask: torch.Tensor, tokens: torch.Tensor, device) -> torch.Tensor:
+        if mask.shape != tokens.shape[:3]:
+            raise ValueError(f"{name} mask must have shape {tuple(tokens.shape[:3])}, got {tuple(mask.shape)}")
+        return mask.to(device=device, dtype=torch.bool)
+
+    def _gate(self, gate: torch.Tensor | float | int, tokens: torch.Tensor, device, dtype):
+        if torch.is_tensor(gate):
+            return gate.to(device=device, dtype=dtype)
+        return torch.tensor(float(gate), device=device, dtype=dtype)
+
+    def __call__(self, streams: MemoryStreamTensors, device=None, dtype=None) -> tuple[dict[str, Any], dict[str, Any]]:
+        ref = streams.anchor_tokens
+        device = device or ref.device
+        dtype = dtype or ref.dtype
+        anchor_tokens = self._tokens("anchor", streams.anchor_tokens, device, dtype)
+        dynamic_tokens = self._tokens("dynamic", streams.dynamic_tokens, device, dtype)
+        revisit_tokens = self._tokens("revisit", streams.revisit_tokens, device, dtype)
+        anchor_mask = self._mask("anchor", streams.anchor_mask, anchor_tokens, device)
+        dynamic_mask = self._mask("dynamic", streams.dynamic_mask, dynamic_tokens, device)
+        revisit_mask = self._mask("revisit", streams.revisit_mask, revisit_tokens, device)
+        kwargs = {
+            "memory_tokens": anchor_tokens,
+            "memory_token_mask": anchor_mask,
+            "memory_dynamic_tokens": dynamic_tokens,
+            "memory_dynamic_mask": dynamic_mask,
+            "memory_retrieval_tokens": revisit_tokens,
+            "memory_retrieval_mask": revisit_mask,
+            "memory_anchor_gate": self._gate(streams.anchor_gate, anchor_tokens, device, dtype),
+            "memory_dynamic_gate": self._gate(streams.dynamic_gate, dynamic_tokens, device, dtype),
+            "memory_retrieval_gate": self._gate(streams.revisit_gate, revisit_tokens, device, dtype),
+        }
+        if self.omit_disabled:
+            if not anchor_mask.any():
+                kwargs["memory_tokens"] = None
+                kwargs["memory_token_mask"] = None
+            if not dynamic_mask.any():
+                kwargs["memory_dynamic_tokens"] = None
+                kwargs["memory_dynamic_mask"] = None
+            if not revisit_mask.any():
+                kwargs["memory_retrieval_tokens"] = None
+                kwargs["memory_retrieval_mask"] = None
+        diagnostics = dict(streams.diagnostics)
+        diagnostics.update(summarize_stream("anchor", anchor_tokens, anchor_mask, kwargs["memory_anchor_gate"]))
+        diagnostics.update(summarize_stream("dynamic", dynamic_tokens, dynamic_mask, kwargs["memory_dynamic_gate"]))
+        diagnostics.update(summarize_stream("revisit", revisit_tokens, revisit_mask, kwargs["memory_retrieval_gate"]))
+        if streams.revisit_gate_raw is not None:
+            raw_gate = streams.revisit_gate_raw.to(device=device, dtype=dtype)
+            diagnostics["revisit_gate_raw"] = raw_gate
+            diagnostics["revisit_gate_raw_mean"] = float(raw_gate.detach().float().mean().item()) if raw_gate.numel() else 0.0
+            diagnostics["revisit_gate_raw_min"] = float(raw_gate.detach().float().min().item()) if raw_gate.numel() else 0.0
+            diagnostics["revisit_gate_raw_max"] = float(raw_gate.detach().float().max().item()) if raw_gate.numel() else 0.0
+        if streams.no_valid_revisit_mask is not None:
+            diagnostics["no_valid_revisit_mask"] = streams.no_valid_revisit_mask.to(device=device, dtype=torch.bool)
+        max_sources = [v for k, v in streams.diagnostics.items() if k.endswith("max_source_frame")]
+        if max_sources:
+            diagnostics["max_source_frame"] = max(int(torch.as_tensor(v).max().item()) for v in max_sources)
+        return kwargs, diagnostics
+
+
+DeMemWMInjectionAdapter = InjectionAdapter

algorithms/worldmem/dememwm/labels.py

@@ -0,0 +1,479 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch
+
+from .types import MemoryRecord
+
+
+LABEL_SOURCE = "deterministic_fov_coverage_plucker"
+
+
+@dataclass(frozen=True)
+class RevisitCandidateLabel:
+    record: MemoryRecord
+    valid: bool
+    gap_valid: bool
+    gap_to_target: int
+    fov_overlap: Optional[float]
+    plucker_overlap: Optional[float]
+    primary_overlap: float
+    coverage_mask: Optional[torch.Tensor]
+    reject_reasons: tuple[str, ...]
+    best_frame_index: Optional[int] = None
+    best_frame_fov_overlap: Optional[float] = None
+
+    @property
+    def sort_key(self) -> tuple[float, float, float, int, int, str]:
+        fov = 0.0 if self.fov_overlap is None else float(self.fov_overlap)
+        plucker = 0.0 if self.plucker_overlap is None else float(self.plucker_overlap)
+        return (
+            -self.primary_overlap,
+            -fov,
+            -plucker,
+            self.gap_to_target,
+            int(self.record.source_start),
+            str(self.record.chunk_id or ""),
+        )
+
+
+def _as_float_tensor(value) -> Optional[torch.Tensor]:
+    if value is None:
+        return None
+    if torch.is_tensor(value):
+        if value.numel() == 0:
+            return None
+        return value.detach().float()
+    tensor = torch.as_tensor(value, dtype=torch.float32)
+    if tensor.numel() == 0:
+        return None
+    return tensor
+
+
+def _pose_frames(value) -> Optional[torch.Tensor]:
+    tensor = _as_float_tensor(value)
+    if tensor is None or tensor.shape[-1] < 5:
+        return None
+    return tensor.reshape(-1, tensor.shape[-1])[:, :5]
+
+
+def _angle_diff_degrees(value: torch.Tensor) -> torch.Tensor:
+    diff = value.abs() % 360.0
+    return torch.where(diff > 180.0, 360.0 - diff, diff)
+
+
+def _target_fov_points(
+    target_pose: torch.Tensor,
+    *,
+    fov_half_h: float,
+    fov_half_v: float,
+    yaw_samples: int,
+    pitch_samples: int,
+    depth_samples: int,
+    radius: float,
+) -> torch.Tensor:
+    yaw_samples = max(1, int(yaw_samples))
+    pitch_samples = max(1, int(pitch_samples))
+    depth_samples = max(1, int(depth_samples))
+    device = target_pose.device
+    dtype = target_pose.dtype
+    if yaw_samples == 1:
+        yaw_offsets = torch.zeros((1,), device=device, dtype=dtype)
+    else:
+        yaw_offsets = torch.linspace(-float(fov_half_h), float(fov_half_h), yaw_samples + 2, device=device, dtype=dtype)[1:-1]
+    if pitch_samples == 1:
+        pitch_offsets = torch.zeros((1,), device=device, dtype=dtype)
+    else:
+        pitch_offsets = torch.linspace(-float(fov_half_v), float(fov_half_v), pitch_samples + 2, device=device, dtype=dtype)[1:-1]
+    if depth_samples == 1:
+        depths = torch.full((1,), float(radius), device=device, dtype=dtype)
+    else:
+        depths = torch.linspace(float(radius) / float(depth_samples), float(radius), depth_samples, device=device, dtype=dtype)
+    depth_grid, pitch_grid, yaw_grid = torch.meshgrid(depths, pitch_offsets, yaw_offsets, indexing="ij")
+    pitch = torch.deg2rad(target_pose[3] + pitch_grid.reshape(-1))
+    yaw = torch.deg2rad(target_pose[4] + yaw_grid.reshape(-1))
+    depth = depth_grid.reshape(-1)
+    cos_pitch = torch.cos(pitch)
+    vectors = torch.stack(
+        [
+            depth * cos_pitch * torch.sin(yaw),
+            depth * torch.sin(pitch),
+            depth * cos_pitch * torch.cos(yaw),
+        ],
+        dim=-1,
+    )
+    return target_pose[:3].reshape(1, 3) + vectors
+
+
+def _inside_fov_3d_hv(
+    points: torch.Tensor,
+    poses: torch.Tensor,
+    *,
+    fov_half_h: float,
+    fov_half_v: float,
+) -> torch.Tensor:
+    vectors = points.unsqueeze(0) - poses[:, None, :3]
+    x = vectors[..., 0]
+    y = vectors[..., 1]
+    z = vectors[..., 2]
+    azimuth = torch.atan2(x, z) * (180.0 / math.pi)
+    elevation = torch.atan2(y, torch.sqrt(x.square() + z.square()).clamp_min(1e-8)) * (180.0 / math.pi)
+    diff_azimuth = _angle_diff_degrees(azimuth - poses[:, None, 4])
+    diff_elevation = _angle_diff_degrees(elevation - poses[:, None, 3])
+    return (diff_azimuth < float(fov_half_h)) & (diff_elevation < float(fov_half_v))
+
+
+def fov_coverage_overlap(
+    source_pose,
+    target_pose,
+    *,
+    fov_half_h: float = 105.0 / 2.0,
+    fov_half_v: float = 75.0 / 2.0,
+    yaw_samples: int = 25,
+    pitch_samples: int = 20,
+    depth_samples: int = 20,
+    radius: float = 30.0,
+) -> tuple[Optional[float], Optional[torch.Tensor]]:
+    source_poses = _pose_frames(source_pose)
+    target_poses = _pose_frames(target_pose)
+    if source_poses is None or target_poses is None:
+        return None, None
+    target = target_poses[-1].to(device=source_poses.device, dtype=source_poses.dtype)
+    points = _target_fov_points(
+        target,
+        fov_half_h=fov_half_h,
+        fov_half_v=fov_half_v,
+        yaw_samples=yaw_samples,
+        pitch_samples=pitch_samples,
+        depth_samples=depth_samples,
+        radius=radius,
+    )
+    coverage_mask = _inside_fov_3d_hv(points, source_poses, fov_half_h=fov_half_h, fov_half_v=fov_half_v).any(dim=0)
+    return float(coverage_mask.float().mean().item()), coverage_mask.detach()
+
+
+def _rotation_from_pose(poses: torch.Tensor) -> torch.Tensor:
+    pitch = torch.deg2rad(poses[:, 3])
+    yaw = torch.deg2rad(poses[:, 4])
+    cos_pitch, sin_pitch = torch.cos(pitch), torch.sin(pitch)
+    cos_yaw, sin_yaw = torch.cos(yaw), torch.sin(yaw)
+    zeros = torch.zeros_like(pitch)
+    ones = torch.ones_like(pitch)
+    r_pitch = torch.stack(
+        [
+            ones, zeros, zeros,
+            zeros, cos_pitch, -sin_pitch,
+            zeros, sin_pitch, cos_pitch,
+        ],
+        dim=-1,
+    ).reshape(-1, 3, 3)
+    r_yaw = torch.stack(
+        [
+            cos_yaw, zeros, sin_yaw,
+            zeros, ones, zeros,
+            -sin_yaw, zeros, cos_yaw,
+        ],
+        dim=-1,
+    ).reshape(-1, 3, 3)
+    return torch.matmul(r_yaw, r_pitch)
+
+
+def _plucker_descriptor(
+    poses: torch.Tensor,
+    *,
+    grid_h: int,
+    grid_w: int,
+    focal_length: float,
+) -> torch.Tensor:
+    grid_h = max(1, int(grid_h))
+    grid_w = max(1, int(grid_w))
+    poses = poses.float()
+    device = poses.device
+    dtype = torch.float32
+    ys, xs = torch.meshgrid(
+        torch.linspace(0, grid_h - 1, grid_h, device=device, dtype=dtype),
+        torch.linspace(0, grid_w - 1, grid_w, device=device, dtype=dtype),
+        indexing="ij",
+    )
+    fx = float(focal_length) * float(grid_w)
+    fy = float(focal_length) * float(grid_h)
+    cx = 0.5 * float(grid_w)
+    cy = 0.5 * float(grid_h)
+    zs = torch.ones_like(xs)
+    dirs = torch.stack([-(xs + 0.5 - cx) / fx, -(ys + 0.5 - cy) / fy, zs], dim=-1)
+    dirs = dirs.reshape(-1, 3)
+    dirs = dirs / dirs.norm(dim=-1, keepdim=True).clamp_min(1e-8)
+    rotation = _rotation_from_pose(poses)
+    rays_d = torch.matmul(dirs.unsqueeze(0), rotation.transpose(1, 2)).float()
+    rays_o = poses[:, None, :3].expand_as(rays_d).float()
+    moments = torch.linalg.cross(rays_o, rays_d, dim=-1)
+    return torch.cat([moments, rays_d], dim=-1).reshape(poses.shape[0], -1)
+
+
+def plucker_overlap(
+    source_pose,
+    target_pose,
+    *,
+    grid_h: int = 4,
+    grid_w: int = 4,
+    focal_length: float = 0.35,
+) -> Optional[float]:
+    source_poses = _pose_frames(source_pose)
+    target_poses = _pose_frames(target_pose)
+    if source_poses is None or target_poses is None:
+        return None
+    target = target_poses[-1:].to(device=source_poses.device, dtype=source_poses.dtype)
+    source_desc = _plucker_descriptor(source_poses, grid_h=grid_h, grid_w=grid_w, focal_length=focal_length)
+    target_desc = _plucker_descriptor(target, grid_h=grid_h, grid_w=grid_w, focal_length=focal_length)
+    diff = source_desc - target_desc
+    distance = torch.linalg.vector_norm(diff, dim=-1) / math.sqrt(float(diff.shape[-1]))
+    best_distance = float(distance.min().item())
+    return float(1.0 / (1.0 + max(best_distance, 0.0)))
+
+
+def _passes_threshold(overlap: Optional[float], threshold: Optional[float]) -> bool:
+    if threshold is None or overlap is None:
+        return True
+    return float(overlap) >= float(threshold)
+
+
+def _batched_pose_overlaps(
+    records: list[MemoryRecord],
+    *,
+    target_pose=None,
+    fov_half_h: float = 105.0 / 2.0,
+    fov_half_v: float = 75.0 / 2.0,
+    fov_yaw_samples: int = 25,
+    fov_pitch_samples: int = 20,
+    fov_depth_samples: int = 20,
+    fov_radius: float = 30.0,
+    plucker_grid_h: int = 4,
+    plucker_grid_w: int = 4,
+    plucker_focal_length: float = 0.35,
+) -> tuple[
+    list[Optional[float]],
+    list[Optional[float]],
+    list[Optional[torch.Tensor]],
+    list[Optional[int]],
+    list[Optional[float]],
+]:
+    fov_overlaps: list[Optional[float]] = [None] * len(records)
+    plucker_overlaps: list[Optional[float]] = [None] * len(records)
+    coverage_masks: list[Optional[torch.Tensor]] = [None] * len(records)
+    best_frame_indices: list[Optional[int]] = [None] * len(records)
+    best_frame_fov_overlaps: list[Optional[float]] = [None] * len(records)
+    target_poses = _pose_frames(target_pose)
+    if not records or target_poses is None:
+        return fov_overlaps, plucker_overlaps, coverage_masks, best_frame_indices, best_frame_fov_overlaps
+
+    pose_records: list[tuple[int, torch.Tensor, torch.Tensor | None]] = []
+    device = None
+    for record_idx, record in enumerate(records):
+        source_poses = _pose_frames(record.pose)
+        if source_poses is None:
+            continue
+        if device is None:
+            device = source_poses.device
+        frame_values = record.frame_indices.detach().reshape(-1)
+        frame_values = frame_values if int(frame_values.numel()) == int(source_poses.shape[0]) else None
+        pose_records.append((record_idx, source_poses, frame_values))
+    if not pose_records or device is None:
+        return fov_overlaps, plucker_overlaps, coverage_masks, best_frame_indices, best_frame_fov_overlaps
+
+    source_pose_blocks = [poses.to(device=device, dtype=torch.float32) for _, poses, _ in pose_records]
+    source_poses = torch.cat(source_pose_blocks, dim=0)
+    record_ids = torch.cat(
+        [
+            torch.full((poses.shape[0],), int(record_idx), device=device, dtype=torch.long)
+            for (record_idx, _, _), poses in zip(pose_records, source_pose_blocks)
+        ],
+        dim=0,
+    )
+    source_frame_blocks = [
+        torch.full((poses.shape[0],), -1, device=device, dtype=torch.long)
+        if frame_values is None
+        else frame_values.to(device=device, dtype=torch.long)
+        for (_, _, frame_values), poses in zip(pose_records, source_pose_blocks)
+    ]
+    source_frame_values = torch.cat(source_frame_blocks, dim=0)
+    pose_record_indices = [record_idx for record_idx, _, _ in pose_records]
+    target = target_poses[-1].to(device=device, dtype=source_poses.dtype)
+
+    points = _target_fov_points(
+        target,
+        fov_half_h=fov_half_h,
+        fov_half_v=fov_half_v,
+        yaw_samples=fov_yaw_samples,
+        pitch_samples=fov_pitch_samples,
+        depth_samples=fov_depth_samples,
+        radius=fov_radius,
+    )
+    inside = _inside_fov_3d_hv(points, source_poses, fov_half_h=fov_half_h, fov_half_v=fov_half_v)
+    per_frame_fov = inside.float().mean(dim=1)
+
+    source_desc = _plucker_descriptor(source_poses, grid_h=plucker_grid_h, grid_w=plucker_grid_w, focal_length=plucker_focal_length)
+    target_desc = _plucker_descriptor(
+        target.reshape(1, -1),
+        grid_h=plucker_grid_h,
+        grid_w=plucker_grid_w,
+        focal_length=plucker_focal_length,
+    )
+    diff = source_desc - target_desc
+    distance = torch.linalg.vector_norm(diff, dim=-1) / math.sqrt(float(diff.shape[-1]))
+    best_distance = torch.full((len(records),), float("inf"), device=device, dtype=distance.dtype)
+    best_distance.scatter_reduce_(0, record_ids, distance, reduce="amin", include_self=True)
+    plucker_values = (1.0 / (1.0 + best_distance.clamp_min(0.0))).detach().cpu().tolist()
+
+    for record_idx in pose_record_indices:
+        rows = record_ids == int(record_idx)
+        if not rows.any():
+            continue
+        record_fov = per_frame_fov[rows]
+        record_inside = inside[rows]
+        best_pose_row = int(torch.argmax(record_fov).item())
+        best_score_value = float(record_fov[best_pose_row].item())
+        fov_overlaps[record_idx] = best_score_value
+        plucker_overlaps[record_idx] = float(plucker_values[record_idx])
+        coverage_masks[record_idx] = record_inside[best_pose_row].detach()
+        valid_rows = rows & (source_frame_values >= 0)
+        if valid_rows.any():
+            frame_scores = per_frame_fov[valid_rows].detach().cpu().tolist()
+            frame_values = source_frame_values[valid_rows].detach().cpu().tolist()
+            best_score, best_frame = max(
+                ((float(score), int(frame)) for score, frame in zip(frame_scores, frame_values)),
+                key=lambda item: (item[0], item[1]),
+            )
+            best_frame_indices[record_idx] = int(best_frame)
+            best_frame_fov_overlaps[record_idx] = float(best_score)
+    return fov_overlaps, plucker_overlaps, coverage_masks, best_frame_indices, best_frame_fov_overlaps
+
+
+def make_revisit_candidate_labels(
+    records: list[MemoryRecord],
+    *,
+    target_frame: int,
+    exclude_local_context_frames: int,
+    target_pose=None,
+    fov_overlap_threshold: Optional[float] = 0.0,
+    plucker_weight: float = 0.1,
+    target_video_id: Any = None,
+    fov_half_h: float = 105.0 / 2.0,
+    fov_half_v: float = 75.0 / 2.0,
+    fov_yaw_samples: int = 25,
+    fov_pitch_samples: int = 20,
+    fov_depth_samples: int = 20,
+    fov_radius: float = 30.0,
+    plucker_grid_h: int = 4,
+    plucker_grid_w: int = 4,
+    plucker_focal_length: float = 0.35,
+) -> list[RevisitCandidateLabel]:
+    del plucker_weight, target_video_id
+    target_frame = int(target_frame)
+    exclude_local_context_frames = int(exclude_local_context_frames)
+
+    score_indices: list[int] = []
+    max_source_frames: list[int] = []
+    gap_flags: list[bool] = []
+    for record_idx, record in enumerate(records):
+        max_source_frame = int(record.source_end) - 1
+        gap_ok = max_source_frame < target_frame - exclude_local_context_frames
+        max_source_frames.append(max_source_frame)
+        gap_flags.append(gap_ok)
+        if gap_ok:
+            score_indices.append(record_idx)
+
+    fov_overlaps: list[Optional[float]] = [None] * len(records)
+    plucker_overlaps: list[Optional[float]] = [None] * len(records)
+    coverage_masks: list[Optional[torch.Tensor]] = [None] * len(records)
+    best_frame_indices: list[Optional[int]] = [None] * len(records)
+    best_frame_fov_overlaps: list[Optional[float]] = [None] * len(records)
+    scored_records = [records[record_idx] for record_idx in score_indices]
+    scored_fov, scored_plucker, scored_masks, scored_best_frames, scored_best_frame_fov = _batched_pose_overlaps(
+        scored_records,
+        target_pose=target_pose,
+        fov_half_h=fov_half_h,
+        fov_half_v=fov_half_v,
+        fov_yaw_samples=fov_yaw_samples,
+        fov_pitch_samples=fov_pitch_samples,
+        fov_depth_samples=fov_depth_samples,
+        fov_radius=fov_radius,
+        plucker_grid_h=plucker_grid_h,
+        plucker_grid_w=plucker_grid_w,
+        plucker_focal_length=plucker_focal_length,
+    )
+    for scored_idx, record_idx in enumerate(score_indices):
+        fov_overlaps[record_idx] = scored_fov[scored_idx]
+        plucker_overlaps[record_idx] = scored_plucker[scored_idx]
+        coverage_masks[record_idx] = scored_masks[scored_idx]
+        best_frame_indices[record_idx] = scored_best_frames[scored_idx]
+        best_frame_fov_overlaps[record_idx] = scored_best_frame_fov[scored_idx]
+
+    labels: list[RevisitCandidateLabel] = []
+    for record_idx, record in enumerate(records):
+        gap = target_frame - max_source_frames[record_idx]
+        fov_overlap = fov_overlaps[record_idx]
+        reasons: list[str] = []
+        if not gap_flags[record_idx]:
+            reasons.append("inside_c_short")
+        if not _passes_threshold(fov_overlap, fov_overlap_threshold):
+            reasons.append("fov_overlap_below_threshold")
+
+        fov_score = 0.0 if fov_overlap is None else float(fov_overlap)
+        labels.append(
+            RevisitCandidateLabel(
+                record=record,
+                valid=not reasons,
+                gap_valid=gap_flags[record_idx],
+                gap_to_target=gap,
+                fov_overlap=fov_overlap,
+                plucker_overlap=plucker_overlaps[record_idx],
+                primary_overlap=fov_score,
+                coverage_mask=coverage_masks[record_idx],
+                reject_reasons=tuple(reasons),
+                best_frame_index=best_frame_indices[record_idx],
+                best_frame_fov_overlap=best_frame_fov_overlaps[record_idx],
+            )
+        )
+    return labels
+
+
+def make_revisit_candidate_label(
+    record: MemoryRecord,
+    *,
+    target_frame: int,
+    exclude_local_context_frames: int,
+    target_pose=None,
+    fov_overlap_threshold: Optional[float] = 0.0,
+    plucker_weight: float = 0.1,
+    target_video_id: Any = None,
+    fov_half_h: float = 105.0 / 2.0,
+    fov_half_v: float = 75.0 / 2.0,
+    fov_yaw_samples: int = 25,
+    fov_pitch_samples: int = 20,
+    fov_depth_samples: int = 20,
+    fov_radius: float = 30.0,
+    plucker_grid_h: int = 4,
+    plucker_grid_w: int = 4,
+    plucker_focal_length: float = 0.35,
+) -> RevisitCandidateLabel:
+    return make_revisit_candidate_labels(
+        [record],
+        target_frame=target_frame,
+        exclude_local_context_frames=exclude_local_context_frames,
+        target_pose=target_pose,
+        fov_overlap_threshold=fov_overlap_threshold,
+        plucker_weight=plucker_weight,
+        target_video_id=target_video_id,
+        fov_half_h=fov_half_h,
+        fov_half_v=fov_half_v,
+        fov_yaw_samples=fov_yaw_samples,
+        fov_pitch_samples=fov_pitch_samples,
+        fov_depth_samples=fov_depth_samples,
+        fov_radius=fov_radius,
+        plucker_grid_h=plucker_grid_h,
+        plucker_grid_w=plucker_grid_w,
+        plucker_focal_length=plucker_focal_length,
+    )[0]

algorithms/worldmem/dememwm/memory.py

@@ -0,0 +1,208 @@
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable, Optional
+
+import torch
+
+from .types import MemoryRecord, MemorySourceType
+
+
+@dataclass
+class MemoryBankQuery:
+    target_frame: int
+    source_type: Optional[MemorySourceType] = None
+    include_generated: bool = True
+    max_records: Optional[int] = None
+    max_slots: Optional[int] = None
+
+
+class CausalMemoryBank:
+    """Small causal memory bank for DeMemWM records."""
+
+    def __init__(self, max_records: Optional[int] = None, max_slots: Optional[int] = None):
+        self.max_records = max_records
+        self.max_slots = max_slots
+        self._records: list[MemoryRecord] = []
+
+    def __len__(self) -> int:
+        return len(self._records)
+
+    @property
+    def records(self) -> tuple[MemoryRecord, ...]:
+        return tuple(self._records)
+
+    def add_record(self, record: MemoryRecord) -> None:
+        if record.source_type == MemorySourceType.PREFIX_GT and record.is_generated:
+            raise ValueError("generated records cannot be high-trust prefix anchors")
+        self._records.append(record)
+        if self.max_records is not None and len(self._records) > self.max_records:
+            self._records = self._records[-self.max_records:]
+
+    def add_prefix_anchors(
+        self,
+        tokens: torch.Tensor,
+        mask: torch.Tensor,
+        frame_indices: torch.Tensor,
+        pose: Optional[torch.Tensor] = None,
+        slots_per_anchor: Optional[int] = None,
+    ) -> None:
+        if tokens.ndim == 2:
+            tokens = tokens.unsqueeze(0)
+        if mask.ndim == 1:
+            mask = mask.unsqueeze(0)
+        flat_frames = frame_indices.detach().reshape(-1)
+        if tokens.shape[0] != flat_frames.numel():
+            raise ValueError("tokens first dimension must match number of frame indices")
+        for i, frame in enumerate(flat_frames.tolist()):
+            rec_tokens = tokens[i]
+            rec_mask = mask[i].bool()
+            if slots_per_anchor is not None:
+                rec_tokens = rec_tokens[:slots_per_anchor]
+                rec_mask = rec_mask[:slots_per_anchor]
+            self.add_record(
+                MemoryRecord(
+                    tokens=rec_tokens,
+                    mask=rec_mask,
+                    source_start=int(frame),
+                    source_end=int(frame) + 1,
+                    frame_indices=torch.as_tensor([frame], device=rec_tokens.device),
+                    pose=None if pose is None else pose[i],
+                    source_type=MemorySourceType.PREFIX_GT,
+                    is_generated=False,
+                    chunk_id=f"prefix_{int(frame)}",
+                )
+            )
+
+    def add_chunk_record(
+        self,
+        tokens: torch.Tensor,
+        mask: torch.Tensor,
+        frame_indices: torch.Tensor,
+        pose: Optional[torch.Tensor] = None,
+        source_type: MemorySourceType = MemorySourceType.PREFIX_GT,
+        is_generated: bool = False,
+        chunk_id: Optional[str] = None,
+        metadata: Optional[dict] = None,
+    ) -> None:
+        flat_frames = frame_indices.detach().reshape(-1)
+        if flat_frames.numel() == 0:
+            raise ValueError("chunk frame_indices must be non-empty")
+        if tokens.ndim != 2:
+            raise ValueError("chunk tokens must have shape (M,D)")
+        if mask.ndim != 1 or mask.shape[0] != tokens.shape[0]:
+            raise ValueError("chunk mask must have shape (M,)")
+        start = int(flat_frames.min().item())
+        end = int(flat_frames.max().item()) + 1
+        self.add_record(
+            MemoryRecord(
+                tokens=tokens,
+                mask=mask.bool(),
+                source_start=start,
+                source_end=end,
+                frame_indices=flat_frames.to(device=tokens.device),
+                pose=pose,
+                source_type=source_type,
+                is_generated=bool(is_generated),
+                chunk_id=chunk_id or f"{source_type.value}_chunk_{start}_{end}",
+                metadata=dict(metadata or {}),
+            )
+        )
+
+    def add_frame_record(
+        self,
+        tokens: torch.Tensor,
+        mask: torch.Tensor,
+        frame_index: torch.Tensor | int,
+        pose: Optional[torch.Tensor] = None,
+        source_type: MemorySourceType = MemorySourceType.REVISIT,
+        is_generated: bool = False,
+        record_id: Optional[str] = None,
+        metadata: Optional[dict] = None,
+    ) -> None:
+        frame_tensor = torch.as_tensor([int(torch.as_tensor(frame_index).reshape(-1)[0].item())], device=tokens.device)
+        frame = int(frame_tensor.item())
+        self.add_record(
+            MemoryRecord(
+                tokens=tokens,
+                mask=mask.bool(),
+                source_start=frame,
+                source_end=frame + 1,
+                frame_indices=frame_tensor,
+                pose=pose,
+                source_type=source_type,
+                is_generated=bool(is_generated),
+                chunk_id=record_id or f"{source_type.value}_frame_{frame}",
+                metadata=dict(metadata or {}),
+            )
+        )
+
+    def add_generated_records(
+        self,
+        tokens: torch.Tensor,
+        mask: torch.Tensor,
+        frame_indices: torch.Tensor,
+        pose: Optional[torch.Tensor] = None,
+        source_type: MemorySourceType = MemorySourceType.GENERATED,
+    ) -> None:
+        if source_type == MemorySourceType.PREFIX_GT:
+            raise ValueError("generated frames cannot be added as PREFIX_GT anchors by default")
+        if tokens.ndim == 2:
+            tokens = tokens.unsqueeze(0)
+        if mask.ndim == 1:
+            mask = mask.unsqueeze(0)
+        flat_frames = frame_indices.detach().reshape(-1)
+        for i, frame in enumerate(flat_frames.tolist()):
+            self.add_record(
+                MemoryRecord(
+                    tokens=tokens[i],
+                    mask=mask[i].bool(),
+                    source_start=int(frame),
+                    source_end=int(frame) + 1,
+                    frame_indices=torch.as_tensor([frame], device=tokens.device),
+                    pose=None if pose is None else pose[i],
+                    source_type=source_type,
+                    is_generated=True,
+                    chunk_id=f"generated_{int(frame)}",
+                )
+            )
+
+    def query(self, query: MemoryBankQuery | int, **kwargs) -> list[MemoryRecord]:
+        if isinstance(query, int):
+            query = MemoryBankQuery(target_frame=query, **kwargs)
+        out: list[MemoryRecord] = []
+        used_slots = 0
+        for record in self._records:
+            if int(record.source_end) > int(query.target_frame):
+                continue
+            if query.source_type is not None and record.source_type != query.source_type:
+                continue
+            if not query.include_generated and record.is_generated:
+                continue
+            if query.max_slots is not None and used_slots >= query.max_slots:
+                break
+            out.append(record)
+            if query.max_slots is not None:
+                used_slots += record.valid_slots
+            if query.max_records is not None and len(out) >= query.max_records:
+                break
+            if query.max_slots is not None and used_slots >= query.max_slots:
+                break
+        return out
+
+    def assert_causal(self, target_frame: int, records: Iterable[MemoryRecord]) -> None:
+        offenders = [r.chunk_id or f"[{r.source_start},{r.source_end})" for r in records if int(r.source_end) > int(target_frame)]
+        if offenders:
+            raise AssertionError(f"future/non-causal memory selected for target {target_frame}: {offenders}")
+
+
+def stack_record_tokens(records: list[MemoryRecord], max_slots: int | None = None):
+    if not records:
+        return None, None
+    tokens = torch.cat([r.tokens for r in records], dim=0)
+    mask = torch.cat([r.mask.bool() for r in records], dim=0)
+    if max_slots is not None:
+        tokens = tokens[:max_slots]
+        mask = mask[:max_slots]
+    return tokens, mask

algorithms/worldmem/dememwm/negatives.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import torch
+
+from .schedules import EVAL_CORRUPTION_BRANCHES
+
+
+def _deterministic_noise_like(tokens: torch.Tensor, seed: int) -> torch.Tensor:
+    flat = torch.arange(tokens.numel(), device=tokens.device, dtype=torch.float32).reshape(tokens.shape)
+    noise = torch.sin(flat + float(seed) * 0.137).to(dtype=tokens.dtype)
+    scale = tokens.detach().float().std().to(device=tokens.device).clamp_min(0.05).to(dtype=tokens.dtype)
+    return noise * scale
+
+
+def apply_revisit_eval_corruption(
+    *,
+    tokens: torch.Tensor,
+    mask: torch.Tensor,
+    branch: str,
+    target_frame: int,
+) -> tuple[torch.Tensor, bool]:
+    if branch not in EVAL_CORRUPTION_BRANCHES or not mask.any():
+        return tokens, False
+
+    corrupted = tokens.clone()
+    if branch == "wrong_pose":
+        corrupted = -corrupted
+    elif branch == "time_shuffle":
+        corrupted = torch.flip(corrupted, dims=(0,))
+    elif branch == "source_matched_random":
+        corrupted = _deterministic_noise_like(corrupted, seed=int(target_frame))
+    elif branch == "local_context_overlap_fake_revisit":
+        corrupted = torch.roll(corrupted, shifts=1, dims=0)
+    elif branch == "pose_shuffle":
+        corrupted = torch.roll(corrupted, shifts=1, dims=-1)
+    elif branch == "wrong_video":
+        corrupted = corrupted.detach().mean().to(dtype=corrupted.dtype).expand_as(corrupted).clone()
+    else:
+        return tokens, False
+
+    return corrupted, True

algorithms/worldmem/dememwm/retrieval.py

@@ -0,0 +1,476 @@
+from __future__ import annotations
+
+import math
+from dataclasses import replace
+from typing import Any, Optional
+
+import torch
+
+from .labels import (
+    LABEL_SOURCE,
+    RevisitCandidateLabel,
+    _inside_fov_3d_hv,
+    _plucker_descriptor,
+    _target_fov_points,
+)
+from .types import MemoryRecord, RevisitRetrievalResult
+
+
+def _overlap_values(labels, name: str) -> list[float]:
+    values: list[float] = []
+    for label in labels:
+        value = getattr(label, name)
+        if value is not None:
+            values.append(float(value))
+    return values
+
+
+def _overlap_stats(values: list[float], prefix: str) -> dict[str, float]:
+    if not values:
+        return {f"{prefix}_mean": 0.0, f"{prefix}_min": 0.0, f"{prefix}_max": 0.0}
+    return {
+        f"{prefix}_mean": float(sum(values) / len(values)),
+        f"{prefix}_min": float(min(values)),
+        f"{prefix}_max": float(max(values)),
+    }
+
+
+def _pose_rows(pose) -> torch.Tensor | None:
+    if pose is None:
+        return None
+    pose_tensor = pose if torch.is_tensor(pose) else torch.as_tensor(pose, dtype=torch.float32)
+    if pose_tensor.ndim == 0 or pose_tensor.numel() == 0 or pose_tensor.shape[-1] < 5:
+        return None
+    return pose_tensor.detach().reshape(-1, pose_tensor.shape[-1])[:, :5].to(dtype=torch.float32)
+
+
+def _pose_forward(poses: torch.Tensor) -> torch.Tensor:
+    pitch = torch.deg2rad(poses[:, 3])
+    yaw = torch.deg2rad(poses[:, 4])
+    cos_pitch = torch.cos(pitch)
+    return torch.stack(
+        [
+            cos_pitch * torch.sin(yaw),
+            torch.sin(pitch),
+            cos_pitch * torch.cos(yaw),
+        ],
+        dim=-1,
+    )
+
+
+def _single_frame_pose(record: MemoryRecord) -> torch.Tensor | None:
+    if int(record.frame_indices.numel()) != 1:
+        return None
+    pose_rows = _pose_rows(record.pose)
+    if pose_rows is None or int(pose_rows.shape[0]) != 1:
+        return None
+    return pose_rows[0]
+
+
+def _vectorized_frame_candidate_labels(
+    records: list[MemoryRecord],
+    *,
+    target_frame: int,
+    target_pose,
+    fov_overlap_threshold: Optional[float],
+    fov_half_h: float,
+    fov_half_v: float,
+    fov_yaw_samples: int,
+    fov_pitch_samples: int,
+    fov_depth_samples: int,
+    fov_radius: float,
+    plucker_grid_h: int,
+    plucker_grid_w: int,
+    plucker_focal_length: float,
+    pose_preselect_topk: Optional[int],
+) -> tuple[list[RevisitCandidateLabel], dict[str, float | int]]:
+    diagnostics: dict[str, float | int] = {
+        "revisit_pose_preselect_input_count": len(records),
+        "revisit_pose_preselect_scored_count": len(records),
+        "revisit_pose_preselect_unscored_count": 0,
+        "revisit_pose_preselect_selected_count": len(records),
+        "revisit_pose_preselect_min_distance": 0.0,
+        "revisit_pose_preselect_max_distance": 0.0,
+        "revisit_exact_fov_candidate_count": len(records),
+        "revisit_vectorized_frame_scorer_used": 1,
+    }
+    if not records:
+        return [], diagnostics
+
+    target_poses = _pose_rows(target_pose)
+    if target_poses is None:
+        raise ValueError("DeMemWM revisit retrieval requires target_pose for frame-level FoV scoring")
+
+    pose_rows: list[torch.Tensor] = []
+    for record in records:
+        pose_row = _single_frame_pose(record)
+        if pose_row is None:
+            raise ValueError(
+                "DeMemWM revisit retrieval requires frame-level records with exactly one frame index and one pose row"
+            )
+        pose_rows.append(pose_row)
+
+    device = pose_rows[0].device
+    if target_poses.is_cuda:
+        device = target_poses.device
+    source_poses = torch.stack([row.to(device=device, dtype=torch.float32) for row in pose_rows], dim=0)
+    target = target_poses[-1].to(device=device, dtype=torch.float32)
+
+    selected_indices = list(range(len(records)))
+    topk = None if pose_preselect_topk is None else int(pose_preselect_topk)
+    if topk is not None and topk > 0 and len(records) > topk:
+        translation_norm = torch.linalg.vector_norm(source_poses[:, :3] - target[:3], dim=-1) / max(float(fov_radius), 1e-6)
+        source_forward = _pose_forward(source_poses)
+        target_forward = _pose_forward(target.reshape(1, -1)).squeeze(0)
+        dot = (source_forward * target_forward.reshape(1, 3)).sum(dim=-1).clamp(-1.0, 1.0)
+        distances = translation_norm + (torch.acos(dot) / math.pi)
+        distance_values = [float(value) for value in distances.detach().cpu().tolist()]
+        ranked = [
+            (
+                distance_values[idx],
+                -int(record.max_source_frame),
+                int(record.source_start),
+                str(record.chunk_id or ""),
+                idx,
+            )
+            for idx, record in enumerate(records)
+        ]
+        ranked.sort()
+        selected_indices = [idx for *_, idx in ranked[:topk]]
+        diagnostics["revisit_pose_preselect_selected_count"] = len(selected_indices)
+        diagnostics["revisit_pose_preselect_min_distance"] = float(min(distance_values))
+        diagnostics["revisit_pose_preselect_max_distance"] = float(max(distance_values))
+
+    selected_tensor = torch.tensor(selected_indices, device=device, dtype=torch.long)
+    selected_records = [records[idx] for idx in selected_indices]
+    selected_poses = source_poses.index_select(0, selected_tensor)
+    points = _target_fov_points(
+        target,
+        fov_half_h=fov_half_h,
+        fov_half_v=fov_half_v,
+        yaw_samples=fov_yaw_samples,
+        pitch_samples=fov_pitch_samples,
+        depth_samples=fov_depth_samples,
+        radius=fov_radius,
+    )
+    inside = _inside_fov_3d_hv(points, selected_poses, fov_half_h=fov_half_h, fov_half_v=fov_half_v)
+    fov_values = inside.float().mean(dim=1)
+
+    source_desc = _plucker_descriptor(
+        selected_poses,
+        grid_h=plucker_grid_h,
+        grid_w=plucker_grid_w,
+        focal_length=plucker_focal_length,
+    )
+    target_desc = _plucker_descriptor(
+        target.reshape(1, -1),
+        grid_h=plucker_grid_h,
+        grid_w=plucker_grid_w,
+        focal_length=plucker_focal_length,
+    )
+    diff = source_desc - target_desc
+    distances = torch.linalg.vector_norm(diff, dim=-1) / math.sqrt(float(diff.shape[-1]))
+    plucker_values = 1.0 / (1.0 + distances.clamp_min(0.0))
+    valid_mask = torch.ones_like(fov_values, dtype=torch.bool)
+    if fov_overlap_threshold is not None:
+        valid_mask = fov_values >= float(fov_overlap_threshold)
+
+    diagnostics["revisit_exact_fov_candidate_count"] = len(selected_records)
+    fov_list = [float(value) for value in fov_values.detach().cpu().tolist()]
+    plucker_list = [float(value) for value in plucker_values.detach().cpu().tolist()]
+    valid_list = [bool(value) for value in valid_mask.detach().cpu().tolist()]
+
+    labels: list[RevisitCandidateLabel] = []
+    for row_idx, record in enumerate(selected_records):
+        fov_overlap = fov_list[row_idx]
+        reasons = () if valid_list[row_idx] else ("fov_overlap_below_threshold",)
+        gap_to_target = int(target_frame) - (int(record.source_end) - 1)
+        labels.append(
+            RevisitCandidateLabel(
+                record=record,
+                valid=valid_list[row_idx],
+                gap_valid=True,
+                gap_to_target=gap_to_target,
+                fov_overlap=fov_overlap,
+                plucker_overlap=plucker_list[row_idx],
+                primary_overlap=fov_overlap,
+                coverage_mask=inside[row_idx].detach(),
+                reject_reasons=reasons,
+                best_frame_index=int(record.max_source_frame),
+                best_frame_fov_overlap=fov_overlap,
+            )
+        )
+    return labels, diagnostics
+
+
+def _coverage_gain(label: RevisitCandidateLabel, covered_mask: torch.Tensor | None) -> float:
+    mask = label.coverage_mask
+    if mask is None or mask.numel() == 0:
+        return 0.0 if label.fov_overlap is None else float(label.fov_overlap)
+    mask = mask.detach().bool()
+    if covered_mask is None or covered_mask.shape != mask.shape:
+        return float(mask.float().mean().item())
+    return float((mask & ~covered_mask.to(device=mask.device, dtype=torch.bool)).float().mean().item())
+
+
+def _coverage_gains(labels: list[RevisitCandidateLabel], covered_mask: torch.Tensor | None) -> list[float]:
+    masks = [label.coverage_mask for label in labels]
+    valid_masks = [mask for mask in masks if mask is not None and mask.numel() > 0]
+    if not valid_masks:
+        return [0.0 if label.fov_overlap is None else float(label.fov_overlap) for label in labels]
+
+    shape = valid_masks[0].shape
+    device = valid_masks[0].device
+    if any(mask.shape != shape for mask in valid_masks):
+        return [_coverage_gain(label, covered_mask) for label in labels]
+
+    stacked = torch.stack([
+        torch.zeros(shape, device=device, dtype=torch.bool)
+        if mask is None or mask.numel() == 0
+        else mask.detach().to(device=device, dtype=torch.bool)
+        for mask in masks
+    ])
+    if covered_mask is None or covered_mask.shape != shape:
+        gains = stacked.float().mean(dim=1)
+    else:
+        covered = covered_mask.to(device=device, dtype=torch.bool)
+        gains = (stacked & ~covered).float().mean(dim=1)
+    return [float(value) for value in gains.detach().cpu().tolist()]
+
+
+def _select_greedy_coverage(
+    labels: list[RevisitCandidateLabel],
+    *,
+    topk: int,
+    plucker_weight: float,
+) -> tuple[list[RevisitCandidateLabel], list[float], list[float]]:
+    remaining = list(labels)
+    selected: list[RevisitCandidateLabel] = []
+    selected_scores: list[float] = []
+    selected_gains: list[float] = []
+    covered_mask: torch.Tensor | None = None
+    for _ in range(max(0, int(topk))):
+        if not remaining:
+            break
+        gains = _coverage_gains(remaining, covered_mask)
+        ranked = []
+        for idx, (label, gain) in enumerate(zip(remaining, gains)):
+            plucker = 0.0 if label.plucker_overlap is None else float(label.plucker_overlap)
+            fov = 0.0 if label.fov_overlap is None else float(label.fov_overlap)
+            plucker_secondary = float(plucker_weight) * plucker
+            ranked.append((
+                -gain,
+                -fov,
+                -plucker_secondary,
+                label.gap_to_target,
+                int(label.record.source_start),
+                str(label.record.chunk_id or ""),
+                idx,
+                gain,
+            ))
+        ranked.sort()
+        _, _, _, _, _, _, best_idx, best_gain = ranked[0]
+        label = remaining.pop(best_idx)
+        selected.append(label)
+        selected_scores.append(float(best_gain))
+        selected_gains.append(float(best_gain))
+        if label.coverage_mask is not None and label.coverage_mask.numel() > 0:
+            mask = label.coverage_mask.detach().bool()
+            if covered_mask is None or covered_mask.shape != mask.shape:
+                covered_mask = torch.zeros_like(mask, dtype=torch.bool)
+            covered_mask = covered_mask.to(device=mask.device, dtype=torch.bool) | mask
+    return selected, selected_scores, selected_gains
+
+
+def _best_selected_label(labels: list[RevisitCandidateLabel]) -> RevisitCandidateLabel | None:
+    if not labels:
+        return None
+    return max(
+        labels,
+        key=lambda label: (
+            0.0 if label.fov_overlap is None else float(label.fov_overlap),
+            0.0 if label.plucker_overlap is None else float(label.plucker_overlap),
+            -int(label.gap_to_target),
+            -int(label.record.source_start),
+            str(label.record.chunk_id or ""),
+        ),
+    )
+
+
+def _best_selected_frame_label(labels: list[RevisitCandidateLabel]) -> RevisitCandidateLabel | None:
+    frame_labels = [label for label in labels if label.best_frame_fov_overlap is not None]
+    if not frame_labels:
+        return None
+    return max(
+        frame_labels,
+        key=lambda label: (
+            float(label.best_frame_fov_overlap),
+            0.0 if label.fov_overlap is None else float(label.fov_overlap),
+            0.0 if label.plucker_overlap is None else float(label.plucker_overlap),
+            -int(label.gap_to_target),
+            -int(label.record.source_start),
+            str(label.record.chunk_id or ""),
+        ),
+    )
+
+
+def _record_with_selected_frame_metadata(
+    label: RevisitCandidateLabel,
+    *,
+    high_quality_fov_threshold: float,
+) -> MemoryRecord:
+    metadata = dict(label.record.metadata or {})
+    if label.fov_overlap is not None:
+        metadata["dememwm_selected_revisit_fov_overlap"] = float(label.fov_overlap)
+    if label.plucker_overlap is not None:
+        metadata["dememwm_selected_revisit_plucker_overlap"] = float(label.plucker_overlap)
+    if label.best_frame_index is not None:
+        metadata["dememwm_selected_frame_index"] = int(label.best_frame_index)
+    if label.best_frame_fov_overlap is not None:
+        frame_fov = float(label.best_frame_fov_overlap)
+        metadata["dememwm_selected_frame_fov_overlap"] = frame_fov
+        metadata["dememwm_selected_frame_fov_threshold"] = float(high_quality_fov_threshold)
+        metadata["dememwm_selected_frame_passes_high_quality"] = bool(frame_fov >= float(high_quality_fov_threshold))
+    return replace(label.record, metadata=metadata)
+
+
+def deterministic_revisit_retrieval(
+    records: list[MemoryRecord],
+    target_frame: int,
+    target_pose: Optional[torch.Tensor] = None,
+    target_summary: Optional[torch.Tensor] = None,
+    topk: int = 2,
+    exclude_local_context_frames: int = 0,
+    fov_overlap_threshold: Optional[float] = 0.30,
+    high_quality_fov_threshold: float = 0.70,
+    plucker_weight: float = 0.1,
+    target_video_id: Any = None,
+    fov_half_h: float = 105.0 / 2.0,
+    fov_half_v: float = 75.0 / 2.0,
+    fov_yaw_samples: int = 25,
+    fov_pitch_samples: int = 20,
+    fov_depth_samples: int = 20,
+    fov_radius: float = 30.0,
+    plucker_grid_h: int = 4,
+    plucker_grid_w: int = 4,
+    plucker_focal_length: float = 0.35,
+    pose_preselect_topk: Optional[int] = 64,
+    **_legacy_scoring_kwargs,
+) -> RevisitRetrievalResult:
+    del target_summary, target_video_id
+    topk = max(0, int(topk))
+    target_frame = int(target_frame)
+    exclude_local_context_frames = int(exclude_local_context_frames)
+    causal_records = [record for record in records if int(record.source_end) <= target_frame]
+    score_records = [
+        record
+        for record in causal_records
+        if int(record.source_end) <= target_frame - exclude_local_context_frames
+    ]
+    labels, pose_preselect_diagnostics = _vectorized_frame_candidate_labels(
+        score_records,
+        target_frame=target_frame,
+        target_pose=target_pose,
+        fov_overlap_threshold=fov_overlap_threshold,
+        fov_half_h=fov_half_h,
+        fov_half_v=fov_half_v,
+        fov_yaw_samples=fov_yaw_samples,
+        fov_pitch_samples=fov_pitch_samples,
+        fov_depth_samples=fov_depth_samples,
+        fov_radius=fov_radius,
+        plucker_grid_h=plucker_grid_h,
+        plucker_grid_w=plucker_grid_w,
+        plucker_focal_length=plucker_focal_length,
+        pose_preselect_topk=pose_preselect_topk,
+    )
+    exact_fov_candidate_count = int(pose_preselect_diagnostics["revisit_exact_fov_candidate_count"])
+    valid_labels = [label for label in labels if label.valid]
+    selected_labels, selected_scores, selected_gains = _select_greedy_coverage(
+        valid_labels,
+        topk=topk,
+        plucker_weight=float(plucker_weight),
+    )
+    best_selected = _best_selected_label(selected_labels)
+    best_selected_frame = _best_selected_frame_label(selected_labels)
+    best_selected_fov = 0.0 if best_selected is None or best_selected.fov_overlap is None else float(best_selected.fov_overlap)
+    best_selected_plucker = 0.0 if best_selected is None or best_selected.plucker_overlap is None else float(best_selected.plucker_overlap)
+    best_selected_gap = -1 if best_selected is None else int(best_selected.gap_to_target)
+    best_selected_frame_fov = 0.0 if best_selected_frame is None else float(best_selected_frame.best_frame_fov_overlap)
+    best_selected_frame_index = -1 if best_selected_frame is None or best_selected_frame.best_frame_index is None else int(best_selected_frame.best_frame_index)
+    high_quality_selected = int(best_selected_frame is not None and best_selected_frame_fov >= float(high_quality_fov_threshold))
+    selected_records = [
+        _record_with_selected_frame_metadata(label, high_quality_fov_threshold=float(high_quality_fov_threshold))
+        for label in selected_labels
+    ]
+    score_device = selected_records[0].tokens.device if selected_records else torch.device("cpu")
+    scores = torch.tensor(selected_scores, dtype=torch.float32, device=score_device)
+
+    fov_values = _overlap_values(valid_labels, "fov_overlap")
+    plucker_values = _overlap_values(valid_labels, "plucker_overlap")
+    selected_gaps = [label.gap_to_target for label in selected_labels]
+    selected_frame_fov_values = [
+        float(label.best_frame_fov_overlap)
+        for label in selected_labels
+        if label.best_frame_fov_overlap is not None
+    ]
+    diagnostics = {
+        "target_frame": int(target_frame),
+        "candidate_count": len(causal_records),
+        "candidate_frame_count": len(causal_records),
+        "valid_candidate_count": len(valid_labels),
+        "revisit_exact_fov_candidate_count": exact_fov_candidate_count,
+        "valid_candidate_frame_count": len(valid_labels),
+        "valid_candidate_label_count": len(valid_labels),
+        "selected_count": len(selected_records),
+        "selected_frame_count": len(selected_records),
+        "revisit_candidate_frame_count": len(causal_records),
+        "revisit_candidate_count": len(causal_records),
+        "valid_revisit_frame_count": len(valid_labels),
+        "valid_revisit_count": len(valid_labels),
+        "valid_revisit_target_count": high_quality_selected,
+        "no_valid_revisit_count": int(len(valid_labels) == 0),
+        "valid_revisit_mask": int(len(valid_labels) > 0),
+        "revisit_abstained_count": int(len(selected_records) == 0),
+        "abstained": bool(len(selected_records) == 0),
+        "revisit_selected_frame_count": len(selected_records),
+        "revisit_selected_count": len(selected_records),
+        "revisit_min_gap_to_target": int(min(selected_gaps)) if selected_gaps else -1,
+        "best_selected_fov_overlap": best_selected_fov,
+        "best_selected_plucker_overlap": best_selected_plucker,
+        "best_selected_gap_frames": best_selected_gap,
+        "best_selected_frame_index": best_selected_frame_index,
+        "best_selected_frame_fov_overlap": best_selected_frame_fov,
+        "best_selected_frame_passes_high_quality": high_quality_selected,
+        "high_quality_selected_revisit": high_quality_selected,
+        "high_quality_fov_threshold": float(high_quality_fov_threshold),
+        "revisit_label_source": LABEL_SOURCE,
+        "selected_frame_ids": [int(record.max_source_frame) for record in selected_records],
+        "selected_frame_record_ids": [record.chunk_id for record in selected_records],
+        "selected_ranges": [(record.source_start, record.source_end) for record in selected_records],
+        "frame_fov_overlap_values": fov_values,
+        "fov_overlap_values": fov_values,
+        "plucker_overlap_values": plucker_values,
+        "best_selected_fov_overlap_values": [] if best_selected is None else [best_selected_fov],
+        "best_selected_plucker_overlap_values": [] if best_selected is None else [best_selected_plucker],
+        "best_selected_gap_frame_values": [] if best_selected is None else [best_selected_gap],
+        "best_selected_frame_fov_overlap_values": [] if best_selected_frame is None else [best_selected_frame_fov],
+        "selected_frame_fov_overlap_values": selected_frame_fov_values,
+        "selected_incremental_fov_overlap_values": selected_gains,
+        "selected_revisit_scores": selected_scores,
+        **pose_preselect_diagnostics,
+    }
+    diagnostics.update(_overlap_stats(fov_values, "revisit_frame_fov_overlap"))
+    diagnostics.update(_overlap_stats(fov_values, "revisit_fov_overlap"))
+    diagnostics.update(_overlap_stats(plucker_values, "revisit_plucker_overlap"))
+    diagnostics.update(_overlap_stats(diagnostics["best_selected_fov_overlap_values"], "revisit_best_selected_fov_overlap"))
+    diagnostics.update(_overlap_stats(diagnostics["best_selected_plucker_overlap_values"], "revisit_best_selected_plucker_overlap"))
+    diagnostics.update(_overlap_stats(diagnostics["best_selected_gap_frame_values"], "revisit_best_selected_gap_frames"))
+    diagnostics.update(_overlap_stats(diagnostics["best_selected_frame_fov_overlap_values"], "revisit_best_selected_frame_fov_overlap"))
+    diagnostics.update(_overlap_stats(selected_frame_fov_values, "revisit_selected_frame_fov_overlap"))
+    diagnostics.update(_overlap_stats(selected_gains, "revisit_incremental_fov_overlap"))
+    return RevisitRetrievalResult(
+        records=selected_records,
+        scores=scores,
+        selected_frame_ids=[int(record.max_source_frame) for record in selected_records],
+        diagnostics=diagnostics,
+    )

algorithms/worldmem/dememwm/schedules.py

@@ -0,0 +1,223 @@
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from .types import StreamGateState
+
+NOISE_BUCKETS = ("high", "mid", "low")
+NOISE_BUCKET_TO_ID = {name: idx for idx, name in enumerate(NOISE_BUCKETS)}
+EVAL_ABLATION_BRANCHES = (
+    "memory_off",
+    "A_only",
+    "D_only",
+    "A_plus_D",
+    "A_plus_D_plus_R_normal",
+    "R_forced_off",
+    "R_forced_on",
+    "wrong_pose",
+    "time_shuffle",
+    "source_matched_random",
+    "pose_shuffle",
+    "wrong_video",
+    "local_context_overlap_fake_revisit",
+)
+EVAL_ABLATION_BRANCH_TO_ID = {name: idx for idx, name in enumerate(EVAL_ABLATION_BRANCHES)}
+EVAL_CORRUPTION_BRANCHES = (
+    "wrong_pose",
+    "time_shuffle",
+    "source_matched_random",
+    "pose_shuffle",
+    "wrong_video",
+    "local_context_overlap_fake_revisit",
+)
+
+
+
+def _clamp01(value: float) -> float:
+    return max(0.0, min(1.0, float(value)))
+
+
+def noise_bucket_from_denoising_fraction(denoising_fraction: float | None) -> str:
+    if denoising_fraction is None:
+        return "mid"
+    frac = _clamp01(float(denoising_fraction))
+    if frac < (1.0 / 3.0):
+        return "high"
+    if frac < (2.0 / 3.0):
+        return "mid"
+    return "low"
+
+
+def noise_bucket_from_noise_levels(noise_levels: torch.Tensor | None, timesteps: int | None) -> str:
+    if noise_levels is None or timesteps is None or int(timesteps) <= 1:
+        return "mid"
+    noise_fraction = _clamp01(float(noise_levels.detach().float().mean().item()) / float(int(timesteps) - 1))
+    if noise_fraction >= (2.0 / 3.0):
+        return "high"
+    if noise_fraction >= (1.0 / 3.0):
+        return "mid"
+    return "low"
+
+
+def noise_bucket_ids_from_noise_levels(noise_levels: torch.Tensor | None, timesteps: int | None) -> torch.Tensor | None:
+    if noise_levels is None or timesteps is None or int(timesteps) <= 1:
+        return None
+    noise_fraction = noise_levels.detach().float() / float(int(timesteps) - 1)
+    bucket_ids = torch.full_like(noise_levels, NOISE_BUCKET_TO_ID["mid"], dtype=torch.long)
+    bucket_ids = torch.where(
+        noise_fraction >= (2.0 / 3.0),
+        torch.full_like(bucket_ids, NOISE_BUCKET_TO_ID["high"]),
+        bucket_ids,
+    )
+    bucket_ids = torch.where(
+        noise_fraction < (1.0 / 3.0),
+        torch.full_like(bucket_ids, NOISE_BUCKET_TO_ID["low"]),
+        bucket_ids,
+    )
+    return bucket_ids
+
+
+def denoising_fraction_from_noise_levels(noise_levels: torch.Tensor | None, timesteps: int | None) -> float | None:
+    if noise_levels is None or timesteps is None or int(timesteps) <= 1:
+        return None
+    noise_fraction = _clamp01(float(noise_levels.detach().float().mean().item()) / float(int(timesteps) - 1))
+    return _clamp01(1.0 - noise_fraction)
+
+
+def normalize_eval_ablation_branch(branch: str | None) -> str:
+    if branch is None:
+        return "A_plus_D_plus_R_normal"
+    branch = str(branch)
+    if branch not in EVAL_ABLATION_BRANCH_TO_ID:
+        raise ValueError(f"unknown DeMemWM eval ablation branch: {branch}")
+    return branch
+
+
+def normalize_noise_bucket(noise_bucket: str | None) -> str:
+    if noise_bucket in NOISE_BUCKET_TO_ID:
+        return str(noise_bucket)
+    return "mid"
+
+
+_STAGE_ENABLES = {
+    'stage_1': (True, True, True),
+    'stage_2': (True, True, True),
+}
+
+
+
+@dataclass(frozen=True)
+class CurriculumState:
+    """Step-resolved DeMemWM curriculum/freezing state for one continuous run."""
+
+    global_step: int
+    enabled: bool
+    stage: str
+    anchor_enabled: bool
+    dynamic_enabled: bool
+    revisit_enabled: bool
+    dit_train_state: str
+    freeze_vae: bool
+    dememwm_lr: float
+    memory_adapter_lr: float
+    full_dit_lr: float
+
+
+    @property
+    def dit_full_trainable(self) -> bool:
+        return self.dit_train_state == "full"
+
+    def diagnostics(self) -> dict[str, Any]:
+        return {
+            "dememwm_global_step": self.global_step,
+            "dememwm_curriculum_enabled": self.enabled,
+            "dememwm_stage": self.stage,
+            "curriculum_anchor_enabled": self.anchor_enabled,
+            "curriculum_dynamic_enabled": self.dynamic_enabled,
+            "curriculum_revisit_enabled": self.revisit_enabled,
+            "dit_train_state": self.dit_train_state,
+            "dit_full_trainable": self.dit_full_trainable,
+            "freeze_vae": self.freeze_vae,
+            "lr_dememwm_modules": self.dememwm_lr,
+            "lr_memory_adapters": self.memory_adapter_lr,
+            "lr_full_dit": self.full_dit_lr,
+        }
+
+
+def _cfg_get(obj: Any, name: str, default: Any) -> Any:
+    return getattr(obj, name, default) if obj is not None else default
+
+
+def _stage_for_step(curriculum_cfg: Any, step: int) -> str:
+    full_start = int(_cfg_get(curriculum_cfg, 'full_stage_start_step', 60000))
+    return 'stage_2' if step >= full_start else 'stage_1'
+
+
+def _dit_train_state(curriculum_cfg: Any, step: int) -> str:
+    freeze_cfg = _cfg_get(curriculum_cfg, 'dit_freeze', None)
+    freeze_enabled = bool(_cfg_get(freeze_cfg, 'enabled', True))
+    full_step = int(_cfg_get(curriculum_cfg, 'full_stage_start_step', 60000))
+    if freeze_enabled and step < full_step:
+        return 'frozen'
+    return 'full'
+
+
+def resolve_curriculum(memory_cfg: Any, global_step: int | None = None) -> CurriculumState:
+    """Resolve internal DeMemWM curriculum phase from Lightning global_step.
+
+    This intentionally supports one continuous training run; stage names are internal
+    gates only and do not imply separate jobs/checkpoints.
+    """
+
+    step = max(0, int(global_step or 0))
+    curriculum_cfg = _cfg_get(memory_cfg, "curriculum", None)
+    lr_cfg = _cfg_get(curriculum_cfg, "lr", None)
+    enabled = bool(_cfg_get(curriculum_cfg, "enabled", False))
+
+    if enabled:
+        stage = _stage_for_step(curriculum_cfg, step)
+        dit_state = _dit_train_state(curriculum_cfg, step)
+        freeze_vae = bool(_cfg_get(curriculum_cfg, "freeze_vae", True))
+    else:
+        stage = str(_cfg_get(memory_cfg, "training_stage", "stage_1"))
+        dit_state = "full"
+        freeze_vae = True
+
+    if stage not in _STAGE_ENABLES:
+        raise ValueError(f"unknown DeMemWM stage: {stage}")
+    anchor_on, dynamic_on, revisit_on = _STAGE_ENABLES[stage]
+    return CurriculumState(
+        global_step=step,
+        enabled=enabled,
+        stage=stage,
+        anchor_enabled=anchor_on,
+        dynamic_enabled=dynamic_on,
+        revisit_enabled=revisit_on,
+        dit_train_state=dit_state,
+        freeze_vae=freeze_vae,
+        dememwm_lr=float(_cfg_get(lr_cfg, "dememwm_modules", 1.0e-4)),
+        memory_adapter_lr=float(_cfg_get(lr_cfg, "memory_adapters", 1.0e-4)),
+        full_dit_lr=float(_cfg_get(lr_cfg, "full_dit", 1.0e-5)),
+    )
+
+
+DeMemWMCurriculumState = CurriculumState
+resolve_dememwm_curriculum = resolve_curriculum
+
+
+def compute_stream_gates(stage: str, denoising_fraction: float | None = None, debug_force_all_streams: bool = False, anchor_gate: float = 1.0, dynamic_gate: float = 1.0, revisit_gate: float = 1.0) -> StreamGateState:
+    if debug_force_all_streams:
+        return StreamGateState(True, True, True, float(anchor_gate), float(dynamic_gate), float(revisit_gate), stage, "debug_force_all_streams")
+    if stage not in _STAGE_ENABLES:
+        raise ValueError(f"unknown DeMemWM stage: {stage}")
+    a_on, d_on, r_on = _STAGE_ENABLES[stage]
+    if denoising_fraction is not None:
+        denoising_fraction = max(0.0, min(1.0, float(denoising_fraction)))
+        stage_scale = 0.25 + 0.75 * denoising_fraction
+    else:
+        stage_scale = 1.0
+    return StreamGateState(a_on, d_on, r_on, float(anchor_gate) if a_on else 0.0, float(dynamic_gate) * stage_scale if d_on else 0.0, float(revisit_gate) * stage_scale if r_on else 0.0, stage, "stage_schedule")

algorithms/worldmem/dememwm/types.py

@@ -0,0 +1,98 @@
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Optional
+
+import torch
+
+
+class MemorySourceType(str, Enum):
+    PREFIX_GT = "prefix_gt"
+    GENERATED = "generated"
+    DYNAMIC = "dynamic"
+    REVISIT = "revisit"
+
+
+@dataclass
+class MemoryRecord:
+    """One causal DeMemWM memory item.
+
+    Frame ranges use an inclusive start and exclusive end: [source_start, source_end).
+    All source frame indices represented by this record must be strictly smaller
+    than a queried target frame unless the caller is explicitly querying an
+    already-committed prefix frame.
+    """
+
+    tokens: torch.Tensor
+    mask: torch.Tensor
+    source_start: int
+    source_end: int
+    frame_indices: torch.Tensor
+    pose: Optional[torch.Tensor]
+    source_type: MemorySourceType
+    is_generated: bool
+    score: Optional[float | torch.Tensor] = None
+    chunk_id: Optional[str] = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if self.source_end <= self.source_start:
+            raise ValueError("source_end must be greater than source_start")
+        if self.tokens.ndim < 2:
+            raise ValueError("tokens must include slot and channel dimensions")
+        self.mask = self.mask.bool()
+        if self.mask.ndim != 1:
+            raise ValueError("mask must have shape (M,)")
+        if self.mask.shape[0] != self.tokens.shape[0]:
+            raise ValueError("mask length must match token slots")
+        if self.source_type == MemorySourceType.PREFIX_GT and self.is_generated:
+            raise ValueError("generated records cannot be PREFIX_GT anchors")
+        if self.frame_indices.numel() == 0:
+            raise ValueError("frame_indices cannot be empty")
+
+    @property
+    def max_source_frame(self) -> int:
+        return int(self.frame_indices.detach().max().item())
+
+    @property
+    def valid_slots(self) -> int:
+        return int(self.mask.detach().sum().item())
+
+
+@dataclass
+class MemoryStreamTensors:
+    anchor_tokens: torch.Tensor
+    anchor_mask: torch.Tensor
+    dynamic_tokens: torch.Tensor
+    dynamic_mask: torch.Tensor
+    revisit_tokens: torch.Tensor
+    revisit_mask: torch.Tensor
+    anchor_gate: torch.Tensor | float
+    dynamic_gate: torch.Tensor | float
+    revisit_gate: torch.Tensor | float
+    revisit_gate_raw: torch.Tensor | None = None
+    valid_revisit_mask: torch.Tensor | None = None
+    no_valid_revisit_mask: torch.Tensor | None = None
+    diagnostics: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class StreamGateState:
+    anchor_enabled: bool
+    dynamic_enabled: bool
+    revisit_enabled: bool
+    anchor_gate: float
+    dynamic_gate: float
+    revisit_gate: float
+    stage: str
+    reason: str = ""
+
+
+@dataclass
+class RevisitRetrievalResult:
+    records: list[MemoryRecord]
+    scores: torch.Tensor
+    selected_frame_ids: list[int]
+    diagnostics: dict[str, Any]

algorithms/worldmem/dememwm_memory_dit.py

@@ -0,0 +1,18 @@
+
+from __future__ import annotations
+
+from .dememwm.algorithm import MemoryDiTMixin
+from .df_video import BaseVideoDiTMinecraft
+
+
+class DeMemWMMinecraft(MemoryDiTMixin, BaseVideoDiTMinecraft):
+    """Standalone DeMemWM / Memory-DiT algorithm.
+
+    Reuses the base video-DiT VAE/diffusion/training infrastructure,
+    but owns memory construction/injection. Does not route through the legacy memory method.
+    """
+
+    pass
+
+
+DeMemWMMemoryDiTMinecraft = DeMemWMMinecraft

algorithms/worldmem/df_base.py

@@ -0,0 +1,307 @@
+"""
+This repo is forked from [Boyuan Chen](https://boyuan.space/)'s research 
+template [repo](https://github.com/buoyancy99/research-template). 
+By its MIT license, you must keep the above sentence in `README.md` 
+and the `LICENSE` file to credit the author.
+"""
+
+from typing import Optional
+from tqdm import tqdm
+from omegaconf import DictConfig
+import numpy as np
+import torch
+import torch.nn.functional as F
+from typing import Any
+from einops import rearrange
+
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+
+from algorithms.common.base_pytorch_algo import BasePytorchAlgo
+from .models.diffusion import Diffusion
+
+
+class DiffusionForcingBase(BasePytorchAlgo):
+    def __init__(self, cfg: DictConfig):
+        self.cfg = cfg
+        self.x_shape = cfg.x_shape
+        self.frame_stack = cfg.frame_stack
+        self.x_stacked_shape = list(self.x_shape)
+        self.x_stacked_shape[0] *= cfg.frame_stack
+        self.guidance_scale = cfg.guidance_scale
+        self.context_frames = cfg.context_frames
+        self.chunk_size = cfg.chunk_size
+        self.action_cond_dim = cfg.action_cond_dim
+        self.causal = cfg.causal
+
+        self.uncertainty_scale = cfg.uncertainty_scale
+        self.timesteps = cfg.diffusion.timesteps
+        self.sampling_timesteps = cfg.diffusion.sampling_timesteps
+        self.clip_noise = cfg.diffusion.clip_noise
+
+        self.cfg.diffusion.cum_snr_decay = self.cfg.diffusion.cum_snr_decay ** (self.frame_stack * cfg.frame_skip)
+
+        self.validation_step_outputs = []
+        super().__init__(cfg)
+
+    def _build_model(self):
+        self.diffusion_model = Diffusion(
+            x_shape=self.x_stacked_shape,
+            action_cond_dim=self.action_cond_dim,
+            is_causal=self.causal,
+            cfg=self.cfg.diffusion,
+        )
+        self.register_data_mean_std(self.cfg.data_mean, self.cfg.data_std)
+
+    def configure_optimizers(self):
+        params = tuple(self.diffusion_model.parameters())
+        optimizer_dynamics = torch.optim.AdamW(
+            params, lr=self.cfg.lr, weight_decay=self.cfg.weight_decay, betas=self.cfg.optimizer_beta
+        )
+        return optimizer_dynamics
+
+    def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
+        # update params
+        optimizer.step(closure=optimizer_closure)
+
+        # manually warm up lr without a scheduler
+        if self.trainer.global_step < self.cfg.warmup_steps:
+            lr_scale = min(1.0, float(self.trainer.global_step + 1) / self.cfg.warmup_steps)
+            for pg in optimizer.param_groups:
+                pg["lr"] = lr_scale * self.cfg.lr
+
+    def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
+        xs, conditions, masks = self._preprocess_batch(batch)
+
+        rand_length = torch.randint(3,xs.shape[0]-2, (1,))[0].item()
+        xs = torch.cat([xs[:rand_length], xs[rand_length-3:rand_length-1]])
+        conditions = torch.cat([conditions[:rand_length], conditions[rand_length-3:rand_length-1]])
+        masks = torch.cat([masks[:rand_length], masks[rand_length-3:rand_length-1]])
+        noise_levels=self._generate_noise_levels(xs)
+        noise_levels[:rand_length] = 15 # stable_noise_levels
+        noise_levels[rand_length+1:] = 15 # stable_noise_levels
+
+        xs_pred, loss = self.diffusion_model(xs, conditions, noise_levels=noise_levels)
+        loss = self.reweight_loss(loss, masks)
+
+        # log the loss
+        if batch_idx % 20 == 0:
+            self.log("training/loss", loss)
+
+        xs = self._unstack_and_unnormalize(xs)
+        xs_pred = self._unstack_and_unnormalize(xs_pred)
+
+        output_dict = {
+            "loss": loss,
+            "xs_pred": xs_pred,
+            "xs": xs,
+        }
+
+        return output_dict
+
+    @torch.no_grad()
+    def validation_step(self, batch, batch_idx, namespace="validation") -> STEP_OUTPUT:
+        xs, conditions, masks = self._preprocess_batch(batch)
+        n_frames, batch_size, *_ = xs.shape
+        xs_pred = []
+        curr_frame = 0
+
+        # context
+        n_context_frames = self.context_frames // self.frame_stack
+        xs_pred = xs[:n_context_frames].clone()
+        curr_frame += n_context_frames
+
+        if self.condtion_similar_length:
+            n_frames -= self.condtion_similar_length
+
+        pbar = tqdm(total=n_frames, initial=curr_frame, desc="Sampling")
+        while curr_frame < n_frames:
+            if self.chunk_size > 0:
+                horizon = min(n_frames - curr_frame, self.chunk_size)
+            else:
+                horizon = n_frames - curr_frame
+            assert horizon <= self.n_tokens, "horizon exceeds the number of tokens."
+            scheduling_matrix = self._generate_scheduling_matrix(horizon)
+
+            chunk = torch.randn((horizon, batch_size, *self.x_stacked_shape), device=self.device)
+            chunk = torch.clamp(chunk, -self.clip_noise, self.clip_noise)
+            xs_pred = torch.cat([xs_pred, chunk], 0)
+
+            # sliding window: only input the last n_tokens frames
+            start_frame = max(0, curr_frame + horizon - self.n_tokens)
+
+            pbar.set_postfix(
+                {
+                    "start": start_frame,
+                    "end": curr_frame + horizon,
+                }
+            )
+
+            if self.condtion_similar_length:
+                xs_pred = torch.cat([xs_pred, xs[curr_frame-self.condtion_similar_length:curr_frame].clone()], 0)
+
+            for m in range(scheduling_matrix.shape[0] - 1):
+
+                from_noise_levels = np.concatenate((np.zeros((curr_frame,), dtype=np.int64), scheduling_matrix[m]))[
+                    :, None
+                ].repeat(batch_size, axis=1)
+                to_noise_levels = np.concatenate(
+                    (
+                        np.zeros((curr_frame,), dtype=np.int64),
+                        scheduling_matrix[m + 1],
+                    )
+                )[
+                    :, None
+                ].repeat(batch_size, axis=1)
+
+                if self.condtion_similar_length:
+                    from_noise_levels = np.concatenate([from_noise_levels, np.array([[0,0,0,0]*self.condtion_similar_length])], axis=0)
+                    to_noise_levels = np.concatenate([to_noise_levels, np.array([[0,0,0,0]*self.condtion_similar_length])], axis=0)
+
+                from_noise_levels = torch.from_numpy(from_noise_levels).to(self.device)
+                to_noise_levels = torch.from_numpy(to_noise_levels).to(self.device)
+
+                # update xs_pred by DDIM or DDPM sampling
+                # input frames within the sliding window
+
+                try:
+                    input_condition = conditions[start_frame : curr_frame + horizon].clone()
+                except:
+                    import pdb;pdb.set_trace()
+                if self.condtion_similar_length:
+                    input_condition = torch.cat([conditions[start_frame : curr_frame + horizon], conditions[-self.condtion_similar_length:]], dim=0)
+                xs_pred[start_frame:] = self.diffusion_model.sample_step(
+                    xs_pred[start_frame:],
+                    input_condition,
+                    from_noise_levels[start_frame:],
+                    to_noise_levels[start_frame:],
+                )
+
+            if self.condtion_similar_length:
+                xs_pred = xs_pred[:-self.condtion_similar_length]
+
+            curr_frame += horizon
+            pbar.update(horizon)
+
+        if self.condtion_similar_length:
+            xs = xs[:-self.condtion_similar_length]
+        # FIXME: loss
+        loss = F.mse_loss(xs_pred, xs, reduction="none")
+        loss = self.reweight_loss(loss, masks)
+        self.validation_step_outputs.append((xs_pred.detach().cpu(), xs.detach().cpu()))
+
+        return loss
+
+    def test_step(self, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        return self.validation_step(*args, **kwargs, namespace="test")
+
+    def on_test_epoch_end(self) -> None:
+        self.on_validation_epoch_end(namespace="test")
+
+    def _generate_noise_levels(self, xs: torch.Tensor, masks: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Generate noise levels for training.
+        """
+        num_frames, batch_size, *_ = xs.shape
+        match self.cfg.noise_level:
+            case "random_all":  # entirely random noise levels
+                noise_levels = torch.randint(0, self.timesteps, (num_frames, batch_size), device=xs.device)
+            case "same":
+                noise_levels = torch.randint(0, self.timesteps, (num_frames, batch_size), device=xs.device)
+                noise_levels[1:] = noise_levels[0]
+
+        if masks is not None:
+            # for frames that are not available, treat as full noise
+            discard = torch.all(~rearrange(masks.bool(), "(t fs) b -> t b fs", fs=self.frame_stack), -1)
+            noise_levels = torch.where(discard, torch.full_like(noise_levels, self.timesteps - 1), noise_levels)
+
+        return noise_levels
+
+    def _generate_scheduling_matrix(self, horizon: int):
+        match self.cfg.scheduling_matrix:
+            case "pyramid":
+                return self._generate_pyramid_scheduling_matrix(horizon, self.uncertainty_scale)
+            case "full_sequence":
+                return np.arange(self.sampling_timesteps, -1, -1)[:, None].repeat(horizon, axis=1)
+            case "autoregressive":
+                return self._generate_pyramid_scheduling_matrix(horizon, self.sampling_timesteps)
+            case "trapezoid":
+                return self._generate_trapezoid_scheduling_matrix(horizon, self.uncertainty_scale)
+
+    def _generate_pyramid_scheduling_matrix(self, horizon: int, uncertainty_scale: float):
+        height = self.sampling_timesteps + int((horizon - 1) * uncertainty_scale) + 1
+        scheduling_matrix = np.zeros((height, horizon), dtype=np.int64)
+        for m in range(height):
+            for t in range(horizon):
+                scheduling_matrix[m, t] = self.sampling_timesteps + int(t * uncertainty_scale) - m
+
+        return np.clip(scheduling_matrix, 0, self.sampling_timesteps)
+
+    def _generate_trapezoid_scheduling_matrix(self, horizon: int, uncertainty_scale: float):
+        height = self.sampling_timesteps + int((horizon + 1) // 2 * uncertainty_scale)
+        scheduling_matrix = np.zeros((height, horizon), dtype=np.int64)
+        for m in range(height):
+            for t in range((horizon + 1) // 2):
+                scheduling_matrix[m, t] = self.sampling_timesteps + int(t * uncertainty_scale) - m
+                scheduling_matrix[m, -t] = self.sampling_timesteps + int(t * uncertainty_scale) - m
+
+        return np.clip(scheduling_matrix, 0, self.sampling_timesteps)
+
+    def reweight_loss(self, loss, weight=None):
+        # Note there is another part of loss reweighting (fused_snr) inside the Diffusion class!
+        loss = rearrange(loss, "t b (fs c) ... -> t b fs c ...", fs=self.frame_stack)
+        if weight is not None:
+            expand_dim = len(loss.shape) - len(weight.shape) - 1
+            weight = rearrange(
+                weight,
+                "(t fs) b ... -> t b fs ..." + " 1" * expand_dim,
+                fs=self.frame_stack,
+            )
+            loss = loss * weight
+
+        return loss.mean()
+
+    def _preprocess_batch(self, batch):
+        xs = batch[0]
+        batch_size, n_frames = xs.shape[:2]
+
+        if n_frames % self.frame_stack != 0:
+            raise ValueError("Number of frames must be divisible by frame stack size")
+        if self.context_frames % self.frame_stack != 0:
+            raise ValueError("Number of context frames must be divisible by frame stack size")
+
+        masks = torch.ones(n_frames, batch_size).to(xs.device)
+        n_frames = n_frames // self.frame_stack
+
+        if self.action_cond_dim:
+            conditions = batch[1]
+            conditions = torch.cat([torch.zeros_like(conditions[:, :1]), conditions[:, 1:]], 1)
+            conditions = rearrange(conditions, "b (t fs) d -> t b (fs d)", fs=self.frame_stack).contiguous()
+
+            # f, _, _ = conditions.shape
+            # predefined_1 = torch.tensor([0,0,0,1]).to(conditions.device)
+            # predefined_2 = torch.tensor([0,0,1,0]).to(conditions.device)
+            # conditions[:f//2] = predefined_1
+            # conditions[f//2:] = predefined_2
+        else:
+            conditions = [None for _ in range(n_frames)]
+
+        xs = self._normalize_x(xs)
+        xs = rearrange(xs, "b (t fs) c ... -> t b (fs c) ...", fs=self.frame_stack).contiguous()
+
+        return xs, conditions, masks
+
+    def _normalize_x(self, xs):
+        shape = [1] * (xs.ndim - self.data_mean.ndim) + list(self.data_mean.shape)
+        mean = self.data_mean.reshape(shape)
+        std = self.data_std.reshape(shape)
+        return (xs - mean) / std
+
+    def _unnormalize_x(self, xs):
+        shape = [1] * (xs.ndim - self.data_mean.ndim) + list(self.data_mean.shape)
+        mean = self.data_mean.reshape(shape)
+        std = self.data_std.reshape(shape)
+        return xs * std + mean
+
+    def _unstack_and_unnormalize(self, xs):
+        xs = rearrange(xs, "t b (fs c) ... -> (t fs) b c ...", fs=self.frame_stack)
+        return self._unnormalize_x(xs)

algorithms/worldmem/df_video.py

@@ -0,0 +1,926 @@
+import os
+import random
+import math
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from torchvision.transforms import InterpolationMode
+from PIL import Image
+from packaging import version as pver
+from einops import rearrange
+from tqdm import tqdm
+from omegaconf import DictConfig
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from algorithms.common.metrics import (
+    LearnedPerceptualImagePatchSimilarity,
+)
+from utils.logging_utils import log_video, get_validation_metrics_for_videos
+from .df_base import DiffusionForcingBase
+from .models.vae import VAE_models
+from .models.diffusion import Diffusion
+from .models.pose_prediction import PosePredictionNet
+import glob
+
+# Utility Functions
+def euler_to_rotation_matrix(pitch, yaw):
+    """
+    Convert pitch and yaw angles (in radians) to a 3x3 rotation matrix.
+    Supports batch input.
+
+    Args:
+        pitch (torch.Tensor): Pitch angles in radians.
+        yaw (torch.Tensor): Yaw angles in radians.
+
+    Returns:
+        torch.Tensor: Rotation matrix of shape (batch_size, 3, 3).
+    """
+    cos_pitch, sin_pitch = torch.cos(pitch), torch.sin(pitch)
+    cos_yaw, sin_yaw = torch.cos(yaw), torch.sin(yaw)
+
+    R_pitch = torch.stack([
+        torch.ones_like(pitch), torch.zeros_like(pitch), torch.zeros_like(pitch),
+        torch.zeros_like(pitch), cos_pitch, -sin_pitch,
+        torch.zeros_like(pitch), sin_pitch, cos_pitch
+    ], dim=-1).reshape(-1, 3, 3)
+
+    R_yaw = torch.stack([
+        cos_yaw, torch.zeros_like(yaw), sin_yaw,
+        torch.zeros_like(yaw), torch.ones_like(yaw), torch.zeros_like(yaw),
+        -sin_yaw, torch.zeros_like(yaw), cos_yaw
+    ], dim=-1).reshape(-1, 3, 3)
+
+    return torch.matmul(R_yaw, R_pitch)
+
+
+def euler_to_camera_to_world_matrix(pose):
+    """
+    Convert (x, y, z, pitch, yaw) to a 4x4 camera-to-world transformation matrix using torch.
+    Supports both (5,) and (f, b, 5) shaped inputs.
+
+    Args:
+        pose (torch.Tensor): Pose tensor of shape (5,) or (f, b, 5).
+
+    Returns:
+        torch.Tensor: Camera-to-world transformation matrix of shape (4, 4).
+    """
+
+    origin_dim = pose.ndim
+    if origin_dim == 1:
+        pose = pose.unsqueeze(0).unsqueeze(0)  # Convert (5,) -> (1, 1, 5)
+    elif origin_dim == 2:
+        pose = pose.unsqueeze(0)
+
+    x, y, z, pitch, yaw = pose[..., 0], pose[..., 1], pose[..., 2], pose[..., 3], pose[..., 4]
+    pitch, yaw = torch.deg2rad(pitch), torch.deg2rad(yaw)
+
+    # Compute rotation matrix (batch mode)
+    R = euler_to_rotation_matrix(pitch, yaw)  # Shape (f*b, 3, 3)
+
+    # Create the 4x4 transformation matrix
+    eye = torch.eye(4, dtype=torch.float32, device=pose.device)
+    camera_to_world = eye.repeat(R.shape[0], 1, 1)  # Shape (f*b, 4, 4)
+
+    # Assign rotation
+    camera_to_world[:, :3, :3] = R
+
+    # Assign translation
+    camera_to_world[:, :3, 3] = torch.stack([x.reshape(-1), y.reshape(-1), z.reshape(-1)], dim=-1)
+
+    # Reshape back to (f, b, 4, 4) if needed
+    if origin_dim == 3:
+        return camera_to_world.view(pose.shape[0], pose.shape[1], 4, 4)
+    elif origin_dim == 2:
+        return camera_to_world.view(pose.shape[0], 4, 4)
+    else:
+        return camera_to_world.squeeze(0).squeeze(0)  # Convert (1,1,4,4) -> (4,4)
+
+def is_inside_fov_3d_hv(points, center, center_pitch, center_yaw, fov_half_h, fov_half_v):
+    """
+    Check whether points are within a given 3D field of view (FOV) 
+    with separately defined horizontal and vertical ranges.
+
+    The center view direction is specified by pitch and yaw (in degrees).
+
+    :param points: (N, B, 3) Sample point coordinates
+    :param center: (3,) Center coordinates of the FOV
+    :param center_pitch: Pitch angle of the center view (in degrees)
+    :param center_yaw: Yaw angle of the center view (in degrees)
+    :param fov_half_h: Horizontal half-FOV angle (in degrees)
+    :param fov_half_v: Vertical half-FOV angle (in degrees)
+    :return: Boolean tensor (N, B), indicating whether each point is inside the FOV
+    """
+    # Compute vectors relative to the center
+    vectors = points - center  # shape (N, B, 3)
+    x = vectors[..., 0]
+    y = vectors[..., 1]
+    z = vectors[..., 2]
+    
+    # Compute horizontal angle (yaw): measured with respect to the z-axis as the forward direction,
+    # and the x-axis as left-right, resulting in a range of -180 to 180 degrees.
+    azimuth = torch.atan2(x, z) * (180 / math.pi)
+    
+    # Compute vertical angle (pitch): measured with respect to the horizontal plane,
+    # resulting in a range of -90 to 90 degrees.
+    elevation = torch.atan2(y, torch.sqrt(x**2 + z**2)) * (180 / math.pi)
+    
+    # Compute the angular difference from the center view (handling circular angle wrap-around)
+    diff_azimuth = (azimuth - center_yaw).abs() % 360
+    diff_elevation = (elevation - center_pitch).abs() % 360
+    
+    # Adjust values greater than 180 degrees to the shorter angular difference
+    diff_azimuth = torch.where(diff_azimuth > 180, 360 - diff_azimuth, diff_azimuth)
+    diff_elevation = torch.where(diff_elevation > 180, 360 - diff_elevation, diff_elevation)
+    
+    # Check if both horizontal and vertical angles are within their respective FOV limits
+    return (diff_azimuth < fov_half_h) & (diff_elevation < fov_half_v)
+    
+def generate_points_in_sphere(n_points, radius):
+    # Sample three independent uniform distributions
+    samples_r = torch.rand(n_points)       # For radius distribution
+    samples_phi = torch.rand(n_points)     # For azimuthal angle phi
+    samples_u = torch.rand(n_points)       # For polar angle theta
+
+    # Apply cube root to ensure uniform volumetric distribution
+    r = radius * torch.pow(samples_r, 1/3)
+    # Azimuthal angle phi uniformly distributed in [0, 2π]
+    phi = 2 * math.pi * samples_phi
+    # Convert u to theta to ensure cos(theta) is uniformly distributed
+    theta = torch.acos(1 - 2 * samples_u)
+
+    # Convert spherical coordinates to Cartesian coordinates
+    x = r * torch.sin(theta) * torch.cos(phi)
+    y = r * torch.sin(theta) * torch.sin(phi)
+    z = r * torch.cos(theta)
+
+    points = torch.stack((x, y, z), dim=1)
+    return points
+
+def tensor_max_with_number(tensor, number):
+    number_tensor = torch.tensor(number, dtype=tensor.dtype, device=tensor.device)
+    result = torch.max(tensor, number_tensor)
+    return result
+
+def custom_meshgrid(*args):
+    # ref: https://pytorch.org/docs/stable/generated/torch.meshgrid.html?highlight=meshgrid#torch.meshgrid
+    if pver.parse(torch.__version__) < pver.parse('1.10'):
+        return torch.meshgrid(*args)
+    else:
+        return torch.meshgrid(*args, indexing='ij')
+    
+def camera_to_world_to_world_to_camera(camera_to_world: torch.Tensor) -> torch.Tensor:
+    """
+    Convert Camera-to-World matrices to World-to-Camera matrices for a tensor with shape (f, b, 4, 4).
+
+    Args:
+        camera_to_world (torch.Tensor): A tensor of shape (f, b, 4, 4), where:
+            f = number of frames,
+            b = batch size.
+
+    Returns:
+        torch.Tensor: A tensor of shape (f, b, 4, 4) representing the World-to-Camera matrices.
+    """
+    # Ensure input is a 4D tensor
+    assert camera_to_world.ndim == 4 and camera_to_world.shape[2:] == (4, 4), \
+        "Input must be of shape (f, b, 4, 4)"
+    
+    # Extract the rotation (R) and translation (T) parts
+    R = camera_to_world[:, :, :3, :3]  # Shape: (f, b, 3, 3)
+    T = camera_to_world[:, :, :3, 3]   # Shape: (f, b, 3)
+    
+    # Initialize an identity matrix for the output
+    world_to_camera = torch.eye(4, device=camera_to_world.device).unsqueeze(0).unsqueeze(0)
+    world_to_camera = world_to_camera.repeat(camera_to_world.size(0), camera_to_world.size(1), 1, 1)  # Shape: (f, b, 4, 4)
+    
+    # Compute the rotation (transpose of R)
+    world_to_camera[:, :, :3, :3] = R.transpose(2, 3)
+    
+    # Compute the translation (-R^T * T)
+    world_to_camera[:, :, :3, 3] = -torch.matmul(R.transpose(2, 3), T.unsqueeze(-1)).squeeze(-1)
+    
+    return world_to_camera.to(camera_to_world.dtype)
+
+def convert_to_plucker(poses, curr_frame, focal_length, image_width, image_height):
+
+    intrinsic = np.asarray([focal_length * image_width,
+                                focal_length * image_height,
+                                0.5 * image_width,
+                                0.5 * image_height], dtype=np.float32)
+
+    c2ws = get_relative_pose(poses, zero_first_frame_scale=curr_frame)
+    c2ws = rearrange(c2ws, "t b m n -> b t m n")
+
+    K = torch.as_tensor(intrinsic, device=poses.device, dtype=poses.dtype).repeat(c2ws.shape[0],c2ws.shape[1],1)  # [B, F, 4]
+    plucker_embedding = ray_condition(K, c2ws, image_height, image_width, device=c2ws.device)
+    plucker_embedding = rearrange(plucker_embedding, "b t h w d -> t b h w d").contiguous()
+
+    return plucker_embedding
+
+
+def get_relative_pose(abs_c2ws, zero_first_frame_scale):
+    abs_w2cs = camera_to_world_to_world_to_camera(abs_c2ws)
+    target_cam_c2w = torch.tensor([
+        [1, 0, 0, 0],
+        [0, 1, 0, 0],
+        [0, 0, 1, 0],
+        [0, 0, 0, 1]
+    ]).to(abs_c2ws.device).to(abs_c2ws.dtype)
+    abs2rel = target_cam_c2w @ abs_w2cs[zero_first_frame_scale]
+    ret_poses = [abs2rel @ abs_c2w for abs_c2w in abs_c2ws]
+    ret_poses = torch.stack(ret_poses)
+    return ret_poses
+
+def ray_condition(K, c2w, H, W, device):
+    # c2w: B, V, 4, 4
+    # K: B, V, 4
+
+    B = K.shape[0]
+
+    j, i = custom_meshgrid(
+        torch.linspace(0, H - 1, H, device=device, dtype=c2w.dtype),
+        torch.linspace(0, W - 1, W, device=device, dtype=c2w.dtype),
+    )
+    i = i.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5  # [B, HxW]
+    j = j.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5  # [B, HxW]
+
+    fx, fy, cx, cy = K.chunk(4, dim=-1)  # B,V, 1
+
+    zs = torch.ones_like(i, device=device, dtype=c2w.dtype)  # [B, HxW]
+    xs = -(i - cx) / fx * zs
+    ys = -(j - cy) / fy * zs 
+
+    zs = zs.expand_as(ys)
+
+    directions = torch.stack((xs, ys, zs), dim=-1)  # B, V, HW, 3
+    directions = directions / directions.norm(dim=-1, keepdim=True)  # B, V, HW, 3
+
+    rays_d = directions @ c2w[..., :3, :3].transpose(-1, -2)  # B, V, 3, HW
+    rays_o = c2w[..., :3, 3]  # B, V, 3
+    rays_o = rays_o[:, :, None].expand_as(rays_d)  # B, V, 3, HW
+    # c2w @ dirctions
+    rays_dxo = torch.linalg.cross(rays_o, rays_d)
+    plucker = torch.cat([rays_dxo, rays_d], dim=-1)
+    plucker = plucker.reshape(B, c2w.shape[1], H, W, 6)  # B, V, H, W, 6
+
+    return plucker
+
+def random_transform(tensor):
+    """
+    Apply the same random translation, rotation, and scaling to all frames in the batch.
+
+    Args:
+        tensor (torch.Tensor): Input tensor of shape (F, B, 3, H, W).
+
+    Returns:
+        torch.Tensor: Transformed tensor of shape (F, B, 3, H, W).
+    """
+    if tensor.ndim != 5:
+        raise ValueError("Input tensor must have shape (F, B, 3, H, W)")
+
+    F, B, C, H, W = tensor.shape
+
+    # Generate random transformation parameters
+    max_translate = 0.2  # Translate up to 20% of width/height
+    max_rotate = 30      # Rotate up to 30 degrees
+    max_scale = 0.2      # Scale change by up to +/- 20%
+
+    translate_x = random.uniform(-max_translate, max_translate) * W
+    translate_y = random.uniform(-max_translate, max_translate) * H
+    rotate_angle = random.uniform(-max_rotate, max_rotate)
+    scale_factor = 1 + random.uniform(-max_scale, max_scale)
+
+    # Apply the same transformation to all frames and batches
+
+    tensor = tensor.reshape(F*B, C, H, W)
+    transformed_tensor = TF.affine(
+        tensor,
+        angle=rotate_angle,
+        translate=(translate_x, translate_y),
+        scale=scale_factor,
+        shear=(0, 0),
+        interpolation=InterpolationMode.BILINEAR,
+        fill=0
+    )
+
+    transformed_tensor = transformed_tensor.reshape(F, B, C, H, W)
+    return transformed_tensor
+
+def save_tensor_as_png(tensor, file_path):
+    """
+    Save a 3*H*W tensor as a PNG image.
+
+    Args:
+        tensor (torch.Tensor): Input tensor of shape (3, H, W).
+        file_path (str): Path to save the PNG file.
+    """
+    if tensor.ndim != 3 or tensor.shape[0] != 3:
+        raise ValueError("Input tensor must have shape (3, H, W)")
+
+    # Convert tensor to PIL Image
+    image = TF.to_pil_image(tensor)
+
+    # Save image
+    image.save(file_path)
+
+class BaseVideoDiTMinecraft(DiffusionForcingBase):
+    """
+    Video generation for MineCraft with memory.
+    """
+
+    def __init__(self, cfg: DictConfig):
+        """
+        Initialize the base video-DiT Minecraft class with the given configuration.
+
+        Args:
+            cfg (DictConfig): Configuration object.
+        """
+        self.n_tokens = cfg.n_frames // cfg.frame_stack # number of max tokens for the model
+        self.n_frames = cfg.n_frames
+        if hasattr(cfg, "n_tokens"):
+            self.n_tokens = cfg.n_tokens // cfg.frame_stack
+        self.memory_condition_length = cfg.memory_condition_length
+        self.pose_cond_dim = getattr(cfg, "pose_cond_dim", 5)
+
+        self.use_plucker = getattr(cfg, "use_plucker", True)
+        self.relative_embedding = getattr(cfg, "relative_embedding", True)
+        self.state_embed_only_on_qk = getattr(cfg, "state_embed_only_on_qk", True)
+        self.use_memory_attention = getattr(cfg, "use_memory_attention", True)
+        self.add_timestamp_embedding = getattr(cfg, "add_timestamp_embedding", True)
+        self.ref_mode = getattr(cfg, "ref_mode", 'sequential')
+        self.log_curve = getattr(cfg, "log_curve", False)
+        self.focal_length =  getattr(cfg, "focal_length", 0.35)
+        self.log_video = cfg.log_video
+        self.save_local = getattr(cfg, "save_local", True)
+        self.local_save_dir = getattr(cfg, "local_save_dir", None)
+        self.lpips_batch_size = getattr(cfg, "lpips_batch_size", 16)
+        self.next_frame_length = getattr(cfg, "next_frame_length", 1)
+        self.require_pose_prediction = getattr(cfg, "require_pose_prediction", False)
+
+        super().__init__(cfg)
+            
+    def _build_model(self):
+
+        self.diffusion_model = Diffusion(
+            reference_length=self.memory_condition_length,
+            x_shape=self.x_stacked_shape,
+            action_cond_dim=self.action_cond_dim,
+            pose_cond_dim=self.pose_cond_dim,
+            is_causal=self.causal,
+            cfg=self.cfg.diffusion,
+            is_dit=True,
+            use_plucker=self.use_plucker,
+            relative_embedding=self.relative_embedding,
+            state_embed_only_on_qk=self.state_embed_only_on_qk,
+            use_memory_attention=self.use_memory_attention,
+            add_timestamp_embedding=self.add_timestamp_embedding,
+            ref_mode=self.ref_mode
+        )
+
+        self.validation_lpips_model = LearnedPerceptualImagePatchSimilarity()
+        vae = VAE_models["vit-l-20-shallow-encoder"]()
+        self.vae = vae.eval()
+
+        if self.require_pose_prediction:
+            self.pose_prediction_model = PosePredictionNet()
+
+    def _generate_noise_levels(self, xs: torch.Tensor, masks = None) -> torch.Tensor:
+        """
+        Generate noise levels for training.
+        """
+        num_frames, batch_size, *_ = xs.shape
+        match self.cfg.noise_level:
+            case "random_all":  # entirely random noise levels
+                noise_levels = torch.randint(0, self.timesteps, (num_frames, batch_size), device=xs.device)
+            case "same":
+                noise_levels = torch.randint(0, self.timesteps, (num_frames, batch_size), device=xs.device)
+                noise_levels[1:] = noise_levels[0]
+
+        if masks is not None:
+            # for frames that are not available, treat as full noise
+            discard = torch.all(~rearrange(masks.bool(), "(t fs) b -> t b fs", fs=self.frame_stack), -1)
+            noise_levels = torch.where(discard, torch.full_like(noise_levels, self.timesteps - 1), noise_levels)
+
+        return noise_levels
+
+    def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
+        """
+        Perform a single training step.
+
+        This function processes the input batch,
+        encodes the input frames, generates noise levels, and computes the loss using the diffusion model.
+
+        Args:
+            batch: Input batch of data containing frames, conditions, poses, etc.
+            batch_idx: Index of the current batch.
+
+        Returns:
+            dict: A dictionary containing the training loss.
+        """
+        xs, conditions, pose_conditions, c2w_mat, frame_idx = self._preprocess_batch(batch)
+
+        if self.use_plucker:
+            if self.relative_embedding:
+                input_pose_condition = []
+                frame_idx_list = []
+                for i in range(self.n_frames):
+                    input_pose_condition.append(
+                        convert_to_plucker(
+                            torch.cat([c2w_mat[i:i + 1], c2w_mat[-self.memory_condition_length:]]).clone(),
+                            0,
+                            focal_length=self.focal_length,
+                            image_height=xs.shape[-2],image_width=xs.shape[-1]
+                        ).to(xs.dtype)
+                    )
+                    frame_idx_list.append(
+                        torch.cat([
+                            frame_idx[i:i + 1] - frame_idx[i:i + 1],
+                            frame_idx[-self.memory_condition_length:] - frame_idx[i:i + 1]
+                        ]).clone()
+                    )
+                input_pose_condition = torch.cat(input_pose_condition)
+                frame_idx_list = torch.cat(frame_idx_list)
+            else:
+                input_pose_condition = convert_to_plucker(
+                    c2w_mat, 0, focal_length=self.focal_length
+                ).to(xs.dtype)
+                frame_idx_list = frame_idx
+        else:
+            input_pose_condition = pose_conditions.to(xs.dtype)
+            frame_idx_list = None
+
+        xs = self.encode(xs)
+
+        noise_levels = self._generate_noise_levels(xs)
+
+        if self.memory_condition_length:
+            noise_levels[-self.memory_condition_length:] = self.diffusion_model.stabilization_level
+            conditions[-self.memory_condition_length:] *= 0
+
+        _, loss = self.diffusion_model(
+            xs,
+            conditions,
+            input_pose_condition,
+            noise_levels=noise_levels,
+            reference_length=self.memory_condition_length,
+            frame_idx=frame_idx_list
+        )
+
+        if self.memory_condition_length:
+            loss = loss[:-self.memory_condition_length]
+
+        loss = self.reweight_loss(loss, None)
+
+        if batch_idx % 20 == 0:
+            self.log("training/loss", loss.cpu())
+
+        return {"loss": loss}
+    
+    def on_validation_epoch_end(self, namespace="validation") -> None:
+        if not self.validation_step_outputs:
+            return
+        
+        xs_pred = []
+        xs = []
+        for pred, gt in self.validation_step_outputs:
+            xs_pred.append(pred)
+            xs.append(gt)
+
+        xs_pred = torch.cat(xs_pred, 1)
+        if gt is not None:
+            xs = torch.cat(xs, 1)
+        else:
+            xs = None
+
+        if self.logger and self.log_video:
+            log_video(
+                xs_pred,
+                xs,
+                step=None if namespace == "test" else self.global_step,
+                namespace=namespace + "_vis",
+                context_frames=self.context_frames,
+                logger=self.logger.experiment,
+                save_local=self.save_local,
+                local_save_dir=self.local_save_dir,
+            )
+
+        if xs is not None:
+            # Move data to the same device as LPIPS model for metric calculation
+            device = next(self.validation_lpips_model.parameters()).device
+            xs_pred_device = xs_pred.to(device)
+            xs_device = xs.to(device)
+            
+            metric_dict = get_validation_metrics_for_videos(
+                xs_pred_device, xs_device, 
+                lpips_model=self.validation_lpips_model,
+                lpips_batch_size=self.lpips_batch_size)
+            
+            self.log_dict(
+                {"mse": metric_dict['mse'],
+                "psnr": metric_dict['psnr'],
+                "lpips": metric_dict['lpips']},
+                sync_dist=True
+            )
+
+            if self.log_curve:
+                psnr_values = metric_dict['frame_wise_psnr'].cpu().tolist()
+                frames = list(range(len(psnr_values)))
+                line_plot = wandb.plot.line_series(
+                    xs = frames,
+                    ys = [psnr_values],
+                    keys = ["PSNR"],
+                    title = "Frame-wise PSNR",
+                    xname = "Frame index"
+                )
+
+                self.logger.experiment.log({"frame_wise_psnr_plot": line_plot})
+      
+        self.validation_step_outputs.clear()
+
+    def _preprocess_batch(self, batch):
+
+        xs, conditions, pose_conditions, frame_index = batch
+
+        if self.action_cond_dim:
+            conditions = torch.cat([torch.zeros_like(conditions[:, :1]), conditions[:, 1:]], 1)
+            conditions = rearrange(conditions, "b t d -> t b d").contiguous()
+        else:
+            raise NotImplementedError("Only support external cond.")
+
+        pose_conditions = rearrange(pose_conditions, "b t d -> t b d").contiguous()
+        c2w_mat = euler_to_camera_to_world_matrix(pose_conditions)
+        xs = rearrange(xs, "b t c ... -> t b c ...").contiguous()
+        frame_index = rearrange(frame_index, "b t -> t b").contiguous()
+
+        return xs, conditions, pose_conditions, c2w_mat, frame_index
+    
+    def encode(self, x):
+        # vae encoding
+        T = x.shape[0]
+        H, W = x.shape[-2:]
+        scaling_factor = 0.07843137255
+
+        x = rearrange(x, "t b c h w -> (t b) c h w")
+        with torch.no_grad():
+            x = self.vae.encode(x * 2 - 1).mean * scaling_factor
+        x = rearrange(x, "(t b) (h w) c -> t b c h w", t=T, h=H // self.vae.patch_size, w=W // self.vae.patch_size)
+        return x
+
+    def decode(self, x):
+        total_frames = x.shape[0]
+        scaling_factor = 0.07843137255
+        x = rearrange(x, "t b c h w -> (t b) (h w) c")
+        with torch.no_grad():
+            x = (self.vae.decode(x / scaling_factor) + 1) / 2
+        x = rearrange(x, "(t b) c h w-> t b c h w", t=total_frames)
+        return x
+
+    def _generate_condition_indices(self, curr_frame, memory_condition_length, xs_pred, pose_conditions, frame_idx, horizon):
+        """
+        Generate indices for condition similarity based on the current frame and pose conditions.
+        """
+        if curr_frame < memory_condition_length:
+            random_idx = [i for i in range(curr_frame)] + [0] * (memory_condition_length - curr_frame)
+            random_idx = np.repeat(np.array(random_idx)[:, None], xs_pred.shape[1], -1)
+        else:
+            # Generate points in a sphere and filter based on field of view
+            num_samples = 10000
+            radius = 30
+            points = generate_points_in_sphere(num_samples, radius).to(pose_conditions.device)
+            points = points[:, None].repeat(1, pose_conditions.shape[1], 1)
+            points += pose_conditions[curr_frame, :, :3][None]
+            fov_half_h = torch.tensor(105 / 2, device=pose_conditions.device)
+            fov_half_v = torch.tensor(75 / 2, device=pose_conditions.device)
+
+            # in_fov1 = is_inside_fov_3d_hv(
+            #     points, pose_conditions[curr_frame, :, :3],
+            #     pose_conditions[curr_frame, :, -2], pose_conditions[curr_frame, :, -1],
+            #     fov_half_h, fov_half_v
+            # )
+
+            in_fov1 = torch.stack([
+                is_inside_fov_3d_hv(points, pc[:, :3], pc[:, -2], pc[:, -1], fov_half_h, fov_half_v)
+                for pc in pose_conditions[curr_frame:curr_frame+horizon]
+            ])
+
+            in_fov1 = torch.sum(in_fov1, 0) > 0
+
+            # Compute overlap ratios and select indices
+            in_fov_list = torch.stack([
+                is_inside_fov_3d_hv(points, pc[:, :3], pc[:, -2], pc[:, -1], fov_half_h, fov_half_v)
+                for pc in pose_conditions[:curr_frame]
+            ])
+
+            random_idx = []
+            for _ in range(memory_condition_length):
+                overlap_ratio = ((in_fov1.bool() & in_fov_list).sum(1)) / in_fov1.sum()
+                
+                confidence = overlap_ratio + (curr_frame - frame_idx[:curr_frame]) / curr_frame * (-0.2)
+
+                if len(random_idx) > 0:
+                    confidence[torch.cat(random_idx)] = -1e10
+                _, r_idx = torch.topk(confidence, k=1, dim=0)
+                random_idx.append(r_idx[0])
+
+                # choice 1: directly remove overlapping region
+                occupied_mask = in_fov_list[r_idx[0, range(in_fov1.shape[-1])], :, range(in_fov1.shape[-1])].permute(1,0)
+                in_fov1 = in_fov1 & ~occupied_mask
+
+                # choice 2: apply similarity filter 
+                # cos_sim = F.cosine_similarity(xs_pred.to(r_idx.device)[r_idx[:, range(in_fov1.shape[1])], 
+                #     range(in_fov1.shape[1])], xs_pred.to(r_idx.device)[:curr_frame], dim=2)
+                # cos_sim = cos_sim.mean((-2,-1))
+
+                # mask_sim = cos_sim>0.9
+                # in_fov_list = in_fov_list & ~mask_sim[:,None].to(in_fov_list.device)
+
+            random_idx = torch.stack(random_idx).cpu()
+
+        return random_idx
+
+    def _prepare_conditions(self, 
+                            start_frame, curr_frame, horizon, conditions, 
+                            pose_conditions, c2w_mat, frame_idx, random_idx,
+                            image_width, image_height):
+        """
+        Prepare input conditions and pose conditions for sampling.
+        """
+
+        padding = torch.zeros((len(random_idx),) + conditions.shape[1:], device=conditions.device, dtype=conditions.dtype)
+        input_condition = torch.cat([conditions[start_frame:curr_frame + horizon], padding], dim=0)
+
+        batch_size = conditions.shape[1]
+
+        if self.use_plucker:
+            if self.relative_embedding:
+                frame_idx_list = []
+                input_pose_condition = []
+                for i in range(start_frame, curr_frame + horizon):
+                    input_pose_condition.append(convert_to_plucker(torch.cat([c2w_mat[i:i+1],c2w_mat[random_idx[:,range(batch_size)], range(batch_size)]]).clone(), 0, focal_length=self.focal_length,
+                                                image_width=image_width, image_height=image_height).to(conditions.dtype))
+                    frame_idx_list.append(torch.cat([frame_idx[i:i+1]-frame_idx[i:i+1], frame_idx[random_idx[:,range(batch_size)], range(batch_size)]-frame_idx[i:i+1]]))
+                input_pose_condition = torch.cat(input_pose_condition)
+                frame_idx_list = torch.cat(frame_idx_list)
+
+            else:
+                input_pose_condition = torch.cat([c2w_mat[start_frame : curr_frame + horizon], c2w_mat[random_idx[:,range(batch_size)], range(batch_size)]], dim=0).clone()
+                input_pose_condition = convert_to_plucker(input_pose_condition, 0, focal_length=self.focal_length)
+                frame_idx_list = None
+        else:
+            input_pose_condition = torch.cat([pose_conditions[start_frame : curr_frame + horizon], pose_conditions[random_idx[:,range(batch_size)], range(batch_size)]], dim=0).clone()
+            frame_idx_list = None
+
+        return input_condition, input_pose_condition, frame_idx_list
+
+    def _prepare_noise_levels(self, scheduling_matrix, m, curr_frame, batch_size, memory_condition_length):
+        """
+        Prepare noise levels for the current sampling step.
+        """
+        from_noise_levels = np.concatenate((np.zeros((curr_frame,), dtype=np.int64), scheduling_matrix[m]))[:, None].repeat(batch_size, axis=1)
+        to_noise_levels = np.concatenate((np.zeros((curr_frame,), dtype=np.int64), scheduling_matrix[m + 1]))[:, None].repeat(batch_size, axis=1)
+        if memory_condition_length:
+            from_noise_levels = np.concatenate([from_noise_levels, np.zeros((memory_condition_length, from_noise_levels.shape[-1]), dtype=np.int32)], axis=0)
+            to_noise_levels = np.concatenate([to_noise_levels, np.zeros((memory_condition_length, from_noise_levels.shape[-1]), dtype=np.int32)], axis=0)
+        from_noise_levels = torch.from_numpy(from_noise_levels).to(self.device)
+        to_noise_levels = torch.from_numpy(to_noise_levels).to(self.device)
+        return from_noise_levels, to_noise_levels
+
+    def validation_step(self, batch, batch_idx, namespace="validation") -> STEP_OUTPUT:
+        """
+        Perform a single validation step.
+
+        This function processes the input batch, encodes frames, generates predictions using a sliding window approach,
+        and handles condition similarity logic for sampling. The results are decoded and stored for evaluation.
+
+        Args:
+            batch: Input batch of data containing frames, conditions, poses, etc.
+            batch_idx: Index of the current batch.
+            namespace: Namespace for logging (default: "validation").
+
+        Returns:
+            None: Appends the predicted and ground truth frames to `self.validation_step_outputs`.
+        """
+        # Preprocess the input batch
+        memory_condition_length = self.memory_condition_length
+        xs_raw, conditions, pose_conditions, c2w_mat, frame_idx = self._preprocess_batch(batch)
+
+
+        # Encode frames in chunks if necessary
+        total_frame = xs_raw.shape[0]
+        if total_frame > 10:
+            xs = torch.cat([
+                self.encode(xs_raw[int(total_frame * i / 10):int(total_frame * (i + 1) / 10)]).cpu()
+                for i in range(10)
+            ])
+        else:
+            xs = self.encode(xs_raw).cpu()
+
+        n_frames, batch_size, *_ = xs.shape
+        curr_frame = 0
+
+        # Initialize context frames
+        n_context_frames = self.context_frames // self.frame_stack
+        xs_pred = xs[:n_context_frames].clone()
+        curr_frame += n_context_frames
+
+        # Progress bar for sampling
+        pbar = tqdm(total=n_frames, initial=curr_frame, desc="Sampling")
+
+        while curr_frame < n_frames:
+            # Determine the horizon for the current chunk
+            horizon = min(n_frames - curr_frame, self.chunk_size) if self.chunk_size > 0 else n_frames - curr_frame
+            assert horizon <= self.n_tokens, "Horizon exceeds the number of tokens."
+
+            # Generate scheduling matrix and initialize noise
+            scheduling_matrix = self._generate_scheduling_matrix(horizon)
+            chunk = torch.randn((horizon, batch_size, *xs_pred.shape[2:]))
+            chunk = torch.clamp(chunk, -self.clip_noise, self.clip_noise).to(xs_pred.device)
+            xs_pred = torch.cat([xs_pred, chunk], 0)
+
+            # Sliding window: only input the last `n_tokens` frames
+            start_frame = max(0, curr_frame + horizon - self.n_tokens)
+            pbar.set_postfix({"start": start_frame, "end": curr_frame + horizon})
+
+            # Handle condition similarity logic
+            if memory_condition_length:
+                random_idx = self._generate_condition_indices(
+                    curr_frame, memory_condition_length, xs_pred, pose_conditions, frame_idx, horizon
+                )
+
+                xs_pred = torch.cat([xs_pred, xs_pred[random_idx[:, range(xs_pred.shape[1])], range(xs_pred.shape[1])].clone()], 0)
+
+            # Prepare input conditions and pose conditions
+            input_condition, input_pose_condition, frame_idx_list = self._prepare_conditions(
+                start_frame, curr_frame, horizon, conditions, pose_conditions, c2w_mat, frame_idx, random_idx,
+                image_width=xs_raw.shape[-1], image_height=xs_raw.shape[-2]
+            )
+
+            # Perform sampling for each step in the scheduling matrix
+            for m in range(scheduling_matrix.shape[0] - 1):
+                from_noise_levels, to_noise_levels = self._prepare_noise_levels(
+                    scheduling_matrix, m, curr_frame, batch_size, memory_condition_length
+                )
+
+                xs_pred[start_frame:] = self.diffusion_model.sample_step(
+                    xs_pred[start_frame:].to(input_condition.device),
+                    input_condition,
+                    input_pose_condition,
+                    from_noise_levels[start_frame:],
+                    to_noise_levels[start_frame:],
+                    current_frame=curr_frame,
+                    mode="validation",
+                    reference_length=memory_condition_length,
+                    frame_idx=frame_idx_list
+                ).cpu()
+
+            # Remove condition similarity frames if applicable
+            if memory_condition_length:
+                xs_pred = xs_pred[:-memory_condition_length]
+
+            curr_frame += horizon
+            pbar.update(horizon)
+
+        # Decode predictions and ground truth
+        xs_pred = self.decode(xs_pred[n_context_frames:].to(conditions.device))
+        xs_decode = self.decode(xs[n_context_frames:].to(conditions.device))
+
+        # Store results for evaluation (move to CPU to save GPU memory)
+        self.validation_step_outputs.append((xs_pred.detach().cpu(), xs_decode.detach().cpu()))
+        return
+
+    @torch.no_grad()
+    def interactive(self, first_frame, new_actions, first_pose, device,
+                    memory_latent_frames, memory_actions, memory_poses, memory_c2w, memory_frame_idx):
+    
+        memory_condition_length = self.memory_condition_length
+
+        if memory_latent_frames is None:
+            first_frame = torch.from_numpy(first_frame)
+            new_actions = torch.from_numpy(new_actions)
+            first_pose = torch.from_numpy(first_pose)
+            first_frame_encode = self.encode(first_frame[None, None].to(device))
+            memory_latent_frames = first_frame_encode.cpu()
+            memory_actions = new_actions[None, None].to(device)
+            memory_poses = first_pose[None, None].to(device)
+            new_c2w_mat = euler_to_camera_to_world_matrix(first_pose)
+            memory_c2w = new_c2w_mat[None, None].to(device)
+            memory_frame_idx = torch.tensor([[0]]).to(device)
+            return first_frame.cpu().numpy(), memory_latent_frames.cpu().numpy(), memory_actions.cpu().numpy(), memory_poses.cpu().numpy(), memory_c2w.cpu().numpy(), memory_frame_idx.cpu().numpy()
+        else:
+            memory_latent_frames = torch.from_numpy(memory_latent_frames)
+            memory_actions = torch.from_numpy(memory_actions).to(device)
+            memory_poses = torch.from_numpy(memory_poses).to(device)
+            memory_c2w = torch.from_numpy(memory_c2w).to(device)
+            memory_frame_idx = torch.from_numpy(memory_frame_idx).to(device)
+            new_actions = new_actions.to(device)
+
+        curr_frame = 0
+        batch_size = 1
+        horizon = self.next_frame_length
+        n_frames = curr_frame + horizon
+        # context
+        n_context_frames = len(memory_latent_frames)
+        xs_pred = memory_latent_frames[:n_context_frames].clone()
+        curr_frame += n_context_frames
+
+        pbar = tqdm(total=n_frames, initial=curr_frame, desc="Sampling")
+
+        new_pose_condition_list = []
+        last_frame = xs_pred[-1].clone()
+        last_pose_condition = memory_poses[-1].clone()
+        curr_actions = new_actions.clone()
+        for hi in range(len(new_actions)):
+            last_pose_condition[:,3:] = last_pose_condition[:,3:] // 15
+            new_pose_condition_offset = self.pose_prediction_model(last_frame.to(device), curr_actions[None, hi], last_pose_condition)
+            new_pose_condition_offset[:,3:] = torch.round(new_pose_condition_offset[:,3:])
+            new_pose_condition = last_pose_condition + new_pose_condition_offset
+            new_pose_condition[:,3:] = new_pose_condition[:,3:] * 15
+            new_pose_condition[:,3:] %= 360
+            last_pose_condition = new_pose_condition.clone()
+            new_pose_condition_list.append(new_pose_condition[None])
+        new_pose_condition_list = torch.cat(new_pose_condition_list, 0)
+        
+        ai = 0
+        while ai < len(new_actions):
+            next_horizon = min(horizon, len(new_actions) - ai)
+            last_frame = xs_pred[-1].clone()
+            curr_actions = new_actions[ai:ai+next_horizon].clone()
+
+            new_pose_condition = new_pose_condition_list[ai:ai+next_horizon].clone()
+
+            new_c2w_mat = euler_to_camera_to_world_matrix(new_pose_condition)
+            memory_poses = torch.cat([memory_poses, new_pose_condition])
+            memory_actions = torch.cat([memory_actions, curr_actions[:, None]])
+            memory_c2w = torch.cat([memory_c2w, new_c2w_mat])
+            new_indices = memory_frame_idx[-1,0] + torch.arange(next_horizon, device=memory_frame_idx.device) + 1
+
+            memory_frame_idx = torch.cat([memory_frame_idx, new_indices[:, None]])
+
+            conditions = memory_actions.clone()
+            pose_conditions = memory_poses.clone()
+            c2w_mat = memory_c2w .clone()
+            frame_idx = memory_frame_idx.clone()
+
+            # generation on frame
+            scheduling_matrix = self._generate_scheduling_matrix(next_horizon)
+            chunk = torch.randn((next_horizon, batch_size, *xs_pred.shape[2:])).to(xs_pred.device)
+            chunk = torch.clamp(chunk, -self.clip_noise, self.clip_noise)
+
+            xs_pred = torch.cat([xs_pred, chunk], 0)
+
+            # sliding window: only input the last n_tokens frames
+            start_frame = max(0, curr_frame - self.n_tokens)
+
+            pbar.set_postfix(
+                {
+                    "start": start_frame,
+                    "end": curr_frame + next_horizon,
+                }
+            )
+
+            # Handle condition similarity logic
+            if memory_condition_length:
+                random_idx = self._generate_condition_indices(
+                    curr_frame, memory_condition_length, xs_pred, pose_conditions, frame_idx, next_horizon
+                )
+                
+                # random_idx = np.unique(random_idx)[:, None]
+                # memory_condition_length = len(random_idx)
+                xs_pred = torch.cat([xs_pred, xs_pred[random_idx[:, range(xs_pred.shape[1])], range(xs_pred.shape[1])].clone()], 0)
+
+            # Prepare input conditions and pose conditions
+            input_condition, input_pose_condition, frame_idx_list = self._prepare_conditions(
+                start_frame, curr_frame, next_horizon, conditions, pose_conditions, c2w_mat, frame_idx, random_idx,
+                image_width=first_frame.shape[-1], image_height=first_frame.shape[-2]
+            )
+
+            # Perform sampling for each step in the scheduling matrix
+            for m in range(scheduling_matrix.shape[0] - 1):
+                from_noise_levels, to_noise_levels = self._prepare_noise_levels(
+                    scheduling_matrix, m, curr_frame, batch_size, memory_condition_length
+                )
+
+                xs_pred[start_frame:] = self.diffusion_model.sample_step(
+                    xs_pred[start_frame:].to(input_condition.device),
+                    input_condition,
+                    input_pose_condition,
+                    from_noise_levels[start_frame:],
+                    to_noise_levels[start_frame:],
+                    current_frame=curr_frame,
+                    mode="validation",
+                    reference_length=memory_condition_length,
+                    frame_idx=frame_idx_list
+                ).cpu()
+
+
+            if memory_condition_length:
+                xs_pred = xs_pred[:-memory_condition_length]
+
+            curr_frame += next_horizon
+            pbar.update(next_horizon)
+            ai += next_horizon
+
+        memory_latent_frames = torch.cat([memory_latent_frames, xs_pred[n_context_frames:]])
+        xs_pred = self.decode(xs_pred[n_context_frames:].to(device)).cpu()
+
+        return xs_pred.cpu().numpy(), memory_latent_frames.cpu().numpy(), memory_actions.cpu().numpy(), \
+            memory_poses.cpu().numpy(), memory_c2w.cpu().numpy(), memory_frame_idx.cpu().numpy()

algorithms/worldmem/models/attention.py

@@ -0,0 +1,342 @@
+"""
+Based on https://github.com/buoyancy99/diffusion-forcing/blob/main/algorithms/diffusion_forcing/models/attention.py
+"""
+
+from typing import Optional
+from collections import namedtuple
+import torch
+from torch import nn
+from torch.nn import functional as F
+from einops import rearrange
+from .rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb
+import numpy as np
+
+class TemporalAxialAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        heads: int,
+        dim_head: int,
+        reference_length: int,
+        rotary_emb: RotaryEmbedding,
+        is_causal: bool = True,
+        is_temporal_independent: bool = False,
+        use_domain_adapter = False
+    ):
+        super().__init__()
+        self.inner_dim = dim_head * heads
+        self.heads = heads
+        self.head_dim = dim_head
+        self.inner_dim = dim_head * heads
+        self.to_qkv = nn.Linear(dim, self.inner_dim * 3, bias=False)
+
+        self.use_domain_adapter = use_domain_adapter
+        if self.use_domain_adapter:
+            lora_rank = 8
+            self.lora_A = nn.Linear(dim, lora_rank, bias=False)
+            self.lora_B = nn.Linear(lora_rank, self.inner_dim * 3, bias=False)
+
+        self.to_out = nn.Linear(self.inner_dim, dim)
+
+        self.rotary_emb = rotary_emb
+        self.is_causal = is_causal
+        self.is_temporal_independent = is_temporal_independent
+
+        self.reference_length = reference_length
+
+    def forward(self, x: torch.Tensor):
+        B, T, H, W, D = x.shape
+
+        # if T>=9:
+        #     try:
+        #         # x = torch.cat([x[:,:-1],x[:,16-T:17-T],x[:,-1:]], dim=1)
+        #         x = torch.cat([x[:,16-T:17-T],x], dim=1)
+        #     except:
+        #         import pdb;pdb.set_trace()
+        #     print("="*50)
+        #     print(x.shape)
+
+        B, T, H, W, D = x.shape
+
+        q, k, v = self.to_qkv(x).chunk(3, dim=-1)
+
+        if self.use_domain_adapter:
+            q_lora, k_lora, v_lora = self.lora_B(self.lora_A(x)).chunk(3, dim=-1)
+            q = q+q_lora
+            k = k+k_lora
+            v = v+v_lora
+
+        q = rearrange(q, "B T H W (h d) -> (B H W) h T d", h=self.heads)
+        k = rearrange(k, "B T H W (h d) -> (B H W) h T d", h=self.heads)
+        v = rearrange(v, "B T H W (h d) -> (B H W) h T d", h=self.heads)
+
+        q = self.rotary_emb.rotate_queries_or_keys(q, self.rotary_emb.freqs)
+        k = self.rotary_emb.rotate_queries_or_keys(k, self.rotary_emb.freqs)
+
+        q, k, v = map(lambda t: t.contiguous(), (q, k, v))
+
+        if self.is_temporal_independent:
+            attn_bias = torch.ones((T, T), dtype=q.dtype, device=q.device)
+            attn_bias = attn_bias.masked_fill(attn_bias == 1, float('-inf'))
+            attn_bias[range(T), range(T)] = 0
+        elif self.is_causal:
+            attn_bias = torch.triu(torch.ones((T, T), dtype=q.dtype, device=q.device), diagonal=1)
+            attn_bias = attn_bias.masked_fill(attn_bias == 1, float('-inf'))
+            attn_bias[(T-self.reference_length):] = float('-inf')
+            attn_bias[range(T), range(T)] = 0
+        else:
+            attn_bias = None
+
+        try:
+            x = F.scaled_dot_product_attention(query=q, key=k, value=v, attn_mask=attn_bias)
+        except:
+            import pdb;pdb.set_trace()
+
+        x = rearrange(x, "(B H W) h T d -> B T H W (h d)", B=B, H=H, W=W)
+        x = x.to(q.dtype)
+
+        # linear proj
+        x = self.to_out(x)
+
+        # if T>=10:
+        #     try:
+        #         # x = torch.cat([x[:,:-2],x[:,-1:]], dim=1)
+        #         x = x[:,1:]
+        #     except:
+        #         import pdb;pdb.set_trace()
+        #     print(x.shape)
+        return x
+
+class SpatialAxialAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        heads: int,
+        dim_head: int,
+        rotary_emb: RotaryEmbedding,
+        use_domain_adapter = False
+    ):
+        super().__init__()
+        self.inner_dim = dim_head * heads
+        self.heads = heads
+        self.head_dim = dim_head
+        self.inner_dim = dim_head * heads
+        self.to_qkv = nn.Linear(dim, self.inner_dim * 3, bias=False)
+        self.use_domain_adapter = use_domain_adapter
+        if self.use_domain_adapter:
+            lora_rank = 8
+            self.lora_A = nn.Linear(dim, lora_rank, bias=False)
+            self.lora_B = nn.Linear(lora_rank, self.inner_dim * 3, bias=False)
+
+        self.to_out = nn.Linear(self.inner_dim, dim)
+
+        self.rotary_emb = rotary_emb
+
+    def forward(self, x: torch.Tensor):
+        B, T, H, W, D = x.shape
+
+        q, k, v = self.to_qkv(x).chunk(3, dim=-1)
+
+        if self.use_domain_adapter:
+            q_lora, k_lora, v_lora = self.lora_B(self.lora_A(x)).chunk(3, dim=-1)
+            q = q+q_lora
+            k = k+k_lora
+            v = v+v_lora
+
+        q = rearrange(q, "B T H W (h d) -> (B T) h H W d", h=self.heads)
+        k = rearrange(k, "B T H W (h d) -> (B T) h H W d", h=self.heads)
+        v = rearrange(v, "B T H W (h d) -> (B T) h H W d", h=self.heads)
+
+        freqs = self.rotary_emb.get_axial_freqs(H, W)
+        q = apply_rotary_emb(freqs, q)
+        k = apply_rotary_emb(freqs, k)
+
+        # prepare for attn
+        q = rearrange(q, "(B T) h H W d -> (B T) h (H W) d", B=B, T=T, h=self.heads)
+        k = rearrange(k, "(B T) h H W d -> (B T) h (H W) d", B=B, T=T, h=self.heads)
+        v = rearrange(v, "(B T) h H W d -> (B T) h (H W) d", B=B, T=T, h=self.heads)
+
+        x = F.scaled_dot_product_attention(query=q, key=k, value=v, is_causal=False)
+
+        x = rearrange(x, "(B T) h (H W) d -> B T H W (h d)", B=B, H=H, W=W)
+        x = x.to(q.dtype)
+
+        # linear proj
+        x = self.to_out(x)
+        return x
+
+class MemTemporalAxialAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        heads: int,
+        dim_head: int,
+        rotary_emb: RotaryEmbedding,
+        is_causal: bool = True,
+    ):
+        super().__init__()
+        self.inner_dim = dim_head * heads
+        self.heads = heads
+        self.head_dim = dim_head
+        self.inner_dim = dim_head * heads
+        self.to_qkv = nn.Linear(dim, self.inner_dim * 3, bias=False)
+        self.to_out = nn.Linear(self.inner_dim, dim)
+
+        self.rotary_emb = rotary_emb
+        self.is_causal = is_causal
+
+        self.reference_length = 3
+
+    def forward(self, x: torch.Tensor):
+        B, T, H, W, D = x.shape
+
+        q, k, v = self.to_qkv(x).chunk(3, dim=-1)
+
+
+        q = rearrange(q, "B T H W (h d) -> (B H W) h T d", h=self.heads)
+        k = rearrange(k, "B T H W (h d) -> (B H W) h T d", h=self.heads)
+        v = rearrange(v, "B T H W (h d) -> (B H W) h T d", h=self.heads)
+
+        
+
+        # q = self.rotary_emb.rotate_queries_or_keys(q, self.rotary_emb.freqs)
+        # k = self.rotary_emb.rotate_queries_or_keys(k, self.rotary_emb.freqs)
+
+        q, k, v = map(lambda t: t.contiguous(), (q, k, v))
+
+        # if T == 21000:
+        #     # 手动计算缩放点积分数
+        #     _, _, _, d_k = q.shape
+        #     scores = torch.einsum("b h n d, b h m d -> b h n m", q, k) / (d_k ** 0.5)  # Shape: (B, T_q, T_k)
+
+        #     # 计算注意力图 (Attention Map)
+        #     attention_map = F.softmax(scores, dim=-1)  # Shape: (B, T_q, T_k)
+        #     b_, h_, n_, m_ = attention_map.shape
+        #     attention_map = attention_map.reshape(1, int(np.sqrt(b_/1)), int(np.sqrt(b_/1)), h_, n_, m_)
+        #     attention_map = attention_map.mean(3)
+
+        #     attn_bias = torch.zeros((T, T), dtype=q.dtype, device=q.device)
+        #     T_origin = T - self.reference_length
+        #     attn_bias[:T_origin, T_origin:] = 1
+        #     attn_bias[range(T), range(T)] = 1
+
+        #     attention_map = attention_map * attn_bias
+
+            # # print 注意力图
+            # import matplotlib.pyplot as plt
+            # fig, axes = plt.subplots(21000, 21000, figsize=(9, 9))  # 调整figsize以适配图像大小
+
+            # # 遍历3*3维度
+            # for i in range(21000):
+            #     for j in range(21000):
+            #         # 取出第(i, j)个子图像
+            #         img = attention_map[0, :, :, i, j].cpu().numpy()
+            #         axes[i, j].imshow(img, cmap='viridis')  # 可以自定义cmap
+            #         axes[i, j].axis('off')  # 隐藏坐标轴
+
+            # # 调整子图间距
+            # plt.tight_layout()
+            # plt.savefig('attention_map.png')
+            # import pdb; pdb.set_trace()
+            # plt.close()
+
+        attn_bias = torch.zeros((T, T), dtype=q.dtype, device=q.device)
+        attn_bias = attn_bias.masked_fill(attn_bias == 0, float('-inf'))
+        T_origin = T - self.reference_length
+        attn_bias[:T_origin, T_origin:] = 0
+        attn_bias[range(T), range(T)] = 0
+
+        # if T==121000:
+        #     import pdb;pdb.set_trace()
+
+        try:
+            x = F.scaled_dot_product_attention(query=q, key=k, value=v, attn_mask=attn_bias)
+        except:
+            import pdb;pdb.set_trace()
+
+        x = rearrange(x, "(B H W) h T d -> B T H W (h d)", B=B, H=H, W=W)
+        x = x.to(q.dtype)
+
+        # linear proj
+        x = self.to_out(x)
+        return x
+
+class MemFullAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        heads: int,
+        dim_head: int,
+        reference_length: int,
+        rotary_emb: RotaryEmbedding,
+        is_causal: bool = True
+    ):
+        super().__init__()
+        self.inner_dim = dim_head * heads
+        self.heads = heads
+        self.head_dim = dim_head
+        self.inner_dim = dim_head * heads
+        self.to_qkv = nn.Linear(dim, self.inner_dim * 3, bias=False)
+        self.to_out = nn.Linear(self.inner_dim, dim)
+
+        self.rotary_emb = rotary_emb
+        self.is_causal = is_causal
+
+        self.reference_length = reference_length
+
+        self.store = None
+
+    def forward(self, x: torch.Tensor, relative_embedding=False,
+                extra_condition=None,
+                state_embed_only_on_qk=False,
+                reference_length=None):
+
+        B, T, H, W, D = x.shape
+
+        if state_embed_only_on_qk:
+            q, k, _ = self.to_qkv(x+extra_condition).chunk(3, dim=-1)
+            _, _, v = self.to_qkv(x).chunk(3, dim=-1)
+        else:
+            q, k, v = self.to_qkv(x).chunk(3, dim=-1)
+
+        if relative_embedding:
+            length = reference_length+1
+            n_frames = T // length
+            x = x.reshape(B, n_frames, length, H, W, D)
+
+            x_list = []
+
+            for i in range(n_frames):
+                if i == n_frames-1:
+                    q_i = rearrange(q[:, i*length:], "B T H W (h d) -> B h (T H W) d", h=self.heads)
+                    k_i = rearrange(k[:, i*length+1:(i+1)*length], "B T H W (h d) -> B h (T H W) d", h=self.heads)
+                    v_i = rearrange(v[:, i*length+1:(i+1)*length], "B T H W (h d) -> B h (T H W) d", h=self.heads)                
+                else:
+                    q_i = rearrange(q[:, i*length:i*length+1], "B T H W (h d) -> B h (T H W) d", h=self.heads)
+                    k_i = rearrange(k[:, i*length+1:(i+1)*length], "B T H W (h d) -> B h (T H W) d", h=self.heads)
+                    v_i = rearrange(v[:, i*length+1:(i+1)*length], "B T H W (h d) -> B h (T H W) d", h=self.heads)
+
+                q_i, k_i, v_i = map(lambda t: t.contiguous(), (q_i, k_i, v_i))
+                x_i = F.scaled_dot_product_attention(query=q_i, key=k_i, value=v_i)
+                x_i = rearrange(x_i, "B h (T H W) d -> B T H W (h d)", B=B, H=H, W=W)
+                x_i = x_i.to(q.dtype)
+                x_list.append(x_i)
+        
+            x = torch.cat(x_list, dim=1)
+
+
+        else:
+            T_ = T - reference_length
+            q = rearrange(q, "B T H W (h d) -> B h (T H W) d", h=self.heads)
+            k = rearrange(k[:, T_:], "B T H W (h d) -> B h (T H W) d", h=self.heads)
+            v = rearrange(v[:, T_:], "B T H W (h d) -> B h (T H W) d", h=self.heads)
+
+            q, k, v = map(lambda t: t.contiguous(), (q, k, v))
+            x = F.scaled_dot_product_attention(query=q, key=k, value=v)
+            x = rearrange(x, "B h (T H W) d -> B T H W (h d)", B=B, H=H, W=W)
+            x = x.to(q.dtype)
+
+        # linear proj
+        x = self.to_out(x)
+
+        return x

algorithms/worldmem/models/diffusion.py

@@ -0,0 +1,594 @@
+from typing import Optional, Callable
+from collections import namedtuple
+from omegaconf import DictConfig
+import torch
+from torch import nn
+from torch.nn import functional as F
+from einops import rearrange
+from .utils import linear_beta_schedule, cosine_beta_schedule, sigmoid_beta_schedule, extract
+from .dit import DiT_models
+
+ModelPrediction = namedtuple("ModelPrediction", ["pred_noise", "pred_x_start", "model_out"])
+
+
+class Diffusion(nn.Module):
+    # Special thanks to lucidrains for the implementation of the base Diffusion model
+    # https://github.com/lucidrains/denoising-diffusion-pytorch
+
+    def __init__(
+        self,
+        x_shape: torch.Size,
+        reference_length: int,
+        action_cond_dim: int,
+        pose_cond_dim,
+        is_causal: bool,
+        cfg: DictConfig,
+        is_dit: bool=False,
+        use_plucker=False,
+        relative_embedding=False,
+        state_embed_only_on_qk=False,
+        use_memory_attention=False,
+        add_timestamp_embedding=False,
+        memory_token_cross_attention=False,
+        memory_cross_attn_layers=None,
+        ref_mode='sequential'
+    ):
+        super().__init__()
+        self.cfg = cfg
+
+        self.x_shape = x_shape
+        self.action_cond_dim = action_cond_dim
+        self.timesteps = cfg.timesteps
+        self.sampling_timesteps = cfg.sampling_timesteps
+        self.beta_schedule = cfg.beta_schedule
+        self.schedule_fn_kwargs = cfg.schedule_fn_kwargs
+        self.objective = cfg.objective
+        self.use_fused_snr = cfg.use_fused_snr
+        self.snr_clip = cfg.snr_clip
+        self.cum_snr_decay = cfg.cum_snr_decay
+        self.ddim_sampling_eta = cfg.ddim_sampling_eta
+        self.clip_noise = cfg.clip_noise
+        self.arch = cfg.architecture
+        self.stabilization_level = cfg.stabilization_level
+        self.is_causal = is_causal
+        self.is_dit = is_dit
+        self.reference_length = reference_length
+        self.pose_cond_dim = pose_cond_dim
+        self.use_plucker = use_plucker
+        self.relative_embedding = relative_embedding
+        self.state_embed_only_on_qk = state_embed_only_on_qk
+        self.use_memory_attention = use_memory_attention
+        self.add_timestamp_embedding = add_timestamp_embedding
+        self.memory_token_cross_attention = memory_token_cross_attention
+        self.memory_cross_attn_layers = memory_cross_attn_layers
+        self.ref_mode = ref_mode
+        if self.use_memory_attention:
+            raise ValueError(
+                "WorldMem reference-frame use_memory_attention has been removed from DiT. "
+                "Use memory_token_cross_attention=True for compact memory tokens."
+            )
+
+        self._build_model()
+        self._build_buffer()
+
+    def _build_model(self):
+        x_channel = self.x_shape[0]
+        if self.is_dit:
+            self.model = DiT_models["DiT-S/2"](action_cond_dim=self.action_cond_dim,
+                                            reference_length=self.reference_length,
+                                            memory_token_cross_attention=self.memory_token_cross_attention,
+                                            memory_cross_attn_layers=self.memory_cross_attn_layers,
+                                            ref_mode=self.ref_mode)
+        else:
+            raise NotImplementedError
+
+    def _build_buffer(self):
+        if self.beta_schedule == "linear":
+            beta_schedule_fn = linear_beta_schedule
+        elif self.beta_schedule == "cosine":
+            beta_schedule_fn = cosine_beta_schedule
+        elif self.beta_schedule == "sigmoid":
+            beta_schedule_fn = sigmoid_beta_schedule
+        else:
+            raise ValueError(f"unknown beta schedule {self.beta_schedule}")
+
+        betas = beta_schedule_fn(self.timesteps, **self.schedule_fn_kwargs)
+
+        alphas = 1.0 - betas
+        alphas_cumprod = torch.cumprod(alphas, dim=0)
+        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
+
+        # sampling related parameters
+        assert self.sampling_timesteps <= self.timesteps
+        self.is_ddim_sampling = self.sampling_timesteps < self.timesteps
+
+        # helper function to register buffer from float64 to float32
+        register_buffer = lambda name, val: self.register_buffer(name, val.to(torch.float32))
+
+        register_buffer("betas", betas)
+        register_buffer("alphas_cumprod", alphas_cumprod)
+        register_buffer("alphas_cumprod_prev", alphas_cumprod_prev)
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+
+        register_buffer("sqrt_alphas_cumprod", torch.sqrt(alphas_cumprod))
+        register_buffer("sqrt_one_minus_alphas_cumprod", torch.sqrt(1.0 - alphas_cumprod))
+        register_buffer("log_one_minus_alphas_cumprod", torch.log(1.0 - alphas_cumprod))
+        register_buffer("sqrt_recip_alphas_cumprod", torch.sqrt(1.0 / alphas_cumprod))
+        register_buffer("sqrt_recipm1_alphas_cumprod", torch.sqrt(1.0 / alphas_cumprod - 1))
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+
+        posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
+
+        # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
+
+        register_buffer("posterior_variance", posterior_variance)
+
+        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+
+        register_buffer(
+            "posterior_log_variance_clipped",
+            torch.log(posterior_variance.clamp(min=1e-20)),
+        )
+        register_buffer(
+            "posterior_mean_coef1",
+            betas * torch.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod),
+        )
+        register_buffer(
+            "posterior_mean_coef2",
+            (1.0 - alphas_cumprod_prev) * torch.sqrt(alphas) / (1.0 - alphas_cumprod),
+        )
+
+        # calculate p2 reweighting
+
+        # register_buffer(
+        #     "p2_loss_weight",
+        #     (self.p2_loss_weight_k + alphas_cumprod / (1 - alphas_cumprod))
+        #     ** -self.p2_loss_weight_gamma,
+        # )
+
+        # derive loss weight
+        # https://arxiv.org/abs/2303.09556
+        # snr: signal noise ratio
+        snr = alphas_cumprod / (1 - alphas_cumprod)
+        clipped_snr = snr.clone()
+        clipped_snr.clamp_(max=self.snr_clip)
+
+        register_buffer("clipped_snr", clipped_snr)
+        register_buffer("snr", snr)
+
+    def add_shape_channels(self, x):
+        return rearrange(x, f"... -> ...{' 1' * len(self.x_shape)}")
+
+    def model_predictions(self, x, t, action_cond=None, current_frame=None, 
+        pose_cond=None, mode="training", reference_length=None, frame_idx=None,
+        memory_tokens=None, memory_token_mask=None, memory_retrieval_tokens=None, memory_retrieval_mask=None,
+        **memory_kwargs):
+        x = x.permute(1,0,2,3,4)
+        action_cond = action_cond.permute(1,0,2)
+        if pose_cond is not None and pose_cond[0] is not None:
+            try:
+                pose_cond = pose_cond.permute(1,0,2)
+            except:
+                pass
+        t = t.permute(1,0)
+        model_output = self.model(x, t, action_cond, current_frame=current_frame, pose_cond=pose_cond, 
+            mode=mode, reference_length=reference_length, frame_idx=frame_idx,
+            memory_tokens=memory_tokens, memory_token_mask=memory_token_mask,
+            memory_retrieval_tokens=memory_retrieval_tokens,
+            memory_retrieval_mask=memory_retrieval_mask, **memory_kwargs)
+        model_output = model_output.permute(1,0,2,3,4)
+        x = x.permute(1,0,2,3,4)
+        t = t.permute(1,0)        
+
+        if self.objective == "pred_noise":
+            pred_noise = torch.clamp(model_output, -self.clip_noise, self.clip_noise)
+            x_start = self.predict_start_from_noise(x, t, pred_noise)
+
+        elif self.objective == "pred_x0":
+            x_start = model_output
+            pred_noise = self.predict_noise_from_start(x, t, x_start)
+
+        elif self.objective == "pred_v":
+            v = model_output
+            x_start = self.predict_start_from_v(x, t, v)
+            pred_noise = self.predict_noise_from_start(x, t, x_start)
+
+
+        return ModelPrediction(pred_noise, x_start, model_output)
+
+    def predict_start_from_noise(self, x_t, t, noise):
+        return (
+            extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
+        )
+
+    def predict_noise_from_start(self, x_t, t, x0):
+        return (extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - x0) / extract(
+            self.sqrt_recipm1_alphas_cumprod, t, x_t.shape
+        )
+
+    def predict_v(self, x_start, t, noise):
+        return (
+            extract(self.sqrt_alphas_cumprod, t, x_start.shape) * noise
+            - extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * x_start
+        )
+
+    def predict_start_from_v(self, x_t, t, v):
+        return (
+            extract(self.sqrt_alphas_cumprod, t, x_t.shape) * x_t
+            - extract(self.sqrt_one_minus_alphas_cumprod, t, x_t.shape) * v
+        )
+
+    def q_mean_variance(self, x_start, t):
+        mean = extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+        variance = extract(1.0 - self.alphas_cumprod, t, x_start.shape)
+        log_variance = extract(self.log_one_minus_alphas_cumprod, t, x_start.shape)
+        return mean, variance, log_variance
+
+    def q_posterior(self, x_start, x_t, t):
+        posterior_mean = (
+            extract(self.posterior_mean_coef1, t, x_t.shape) * x_start
+            + extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = extract(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def q_sample(self, x_start, t, noise=None):
+        if noise is None:
+            noise = torch.randn_like(x_start)
+            noise = torch.clamp(noise, -self.clip_noise, self.clip_noise)
+        return (
+            extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+            + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
+        )
+
+    def p_mean_variance(
+        self,
+        x,
+        t,
+        action_cond=None,
+        pose_cond=None,
+        reference_length=None,
+        frame_idx=None,
+        memory_tokens=None,
+        memory_token_mask=None,
+        memory_retrieval_tokens=None,
+        memory_retrieval_mask=None,
+        **memory_kwargs,
+    ):
+        model_pred = self.model_predictions(x=x, t=t, action_cond=action_cond, 
+            pose_cond=pose_cond, reference_length=reference_length,
+            frame_idx=frame_idx, memory_tokens=memory_tokens, memory_token_mask=memory_token_mask,
+            memory_retrieval_tokens=memory_retrieval_tokens,
+            memory_retrieval_mask=memory_retrieval_mask, **memory_kwargs)
+        x_start = model_pred.pred_x_start
+        return self.q_posterior(x_start=x_start, x_t=x, t=t)
+
+    def compute_loss_weights(self, noise_levels: torch.Tensor):
+
+        snr = self.snr[noise_levels]
+        clipped_snr = self.clipped_snr[noise_levels]
+        normalized_clipped_snr = clipped_snr / self.snr_clip
+        normalized_snr = snr / self.snr_clip
+
+        if not self.use_fused_snr:
+            # min SNR reweighting
+            match self.objective:
+                case "pred_noise":
+                    return clipped_snr / snr
+                case "pred_x0":
+                    return clipped_snr
+                case "pred_v":
+                    return clipped_snr / (snr + 1)
+
+        cum_snr = torch.zeros_like(normalized_snr)
+        for t in range(0, noise_levels.shape[0]):
+            if t == 0:
+                cum_snr[t] = normalized_clipped_snr[t]
+            else:
+                cum_snr[t] = self.cum_snr_decay * cum_snr[t - 1] + (1 - self.cum_snr_decay) * normalized_clipped_snr[t]
+
+        cum_snr = F.pad(cum_snr[:-1], (0, 0, 1, 0), value=0.0)
+        clipped_fused_snr = 1 - (1 - cum_snr * self.cum_snr_decay) * (1 - normalized_clipped_snr)
+        fused_snr = 1 - (1 - cum_snr * self.cum_snr_decay) * (1 - normalized_snr)
+
+        match self.objective:
+            case "pred_noise":
+                return clipped_fused_snr / fused_snr
+            case "pred_x0":
+                return clipped_fused_snr * self.snr_clip
+            case "pred_v":
+                return clipped_fused_snr * self.snr_clip / (fused_snr * self.snr_clip + 1)
+            case _:
+                raise ValueError(f"unknown objective {self.objective}")
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        action_cond: Optional[torch.Tensor],
+        pose_cond,
+        noise_levels: torch.Tensor,
+        reference_length,
+        frame_idx=None,
+        memory_tokens=None,
+        memory_token_mask=None,
+        memory_retrieval_tokens=None,
+        memory_retrieval_mask=None,
+        **memory_kwargs,
+    ):
+        noise = torch.randn_like(x)
+        noise = torch.clamp(noise, -self.clip_noise, self.clip_noise)
+        
+        noised_x = self.q_sample(x_start=x, t=noise_levels, noise=noise)
+
+        model_pred = self.model_predictions(x=noised_x, t=noise_levels, action_cond=action_cond, 
+                                    pose_cond=pose_cond,reference_length=reference_length, frame_idx=frame_idx,
+                                    memory_tokens=memory_tokens, memory_token_mask=memory_token_mask,
+                                    memory_retrieval_tokens=memory_retrieval_tokens,
+                                    memory_retrieval_mask=memory_retrieval_mask, **memory_kwargs)
+
+        pred = model_pred.model_out
+        x_pred = model_pred.pred_x_start
+
+        if self.objective == "pred_noise":
+            target = noise
+        elif self.objective == "pred_x0":
+            target = x
+        elif self.objective == "pred_v":
+            target = self.predict_v(x, noise_levels, noise)
+        else:
+            raise ValueError(f"unknown objective {self.objective}")
+
+        # 训练的时候每个frame随便给噪声
+        loss = F.mse_loss(pred, target.detach(), reduction="none")
+        loss_weight = self.compute_loss_weights(noise_levels)
+
+        loss_weight = loss_weight.view(*loss_weight.shape, *((1,) * (loss.ndim - 2)))
+
+        loss = loss * loss_weight
+
+        return x_pred, loss
+
+    def sample_step(
+        self,
+        x: torch.Tensor,
+        action_cond: Optional[torch.Tensor],
+        pose_cond,
+        curr_noise_level: torch.Tensor,
+        next_noise_level: torch.Tensor,
+        guidance_fn: Optional[Callable] = None,
+        current_frame=None,
+        mode="training",
+        reference_length=None,
+        frame_idx=None,
+        memory_tokens=None,
+        memory_token_mask=None,
+        memory_retrieval_tokens=None,
+        memory_retrieval_mask=None,
+        **memory_kwargs,
+    ):
+        real_steps = torch.linspace(-1, self.timesteps - 1, steps=self.sampling_timesteps + 1, device=x.device).long()
+
+        # convert noise levels (0 ~ sampling_timesteps) to real noise levels (-1 ~ timesteps - 1)
+        curr_noise_level = real_steps[curr_noise_level]
+        next_noise_level = real_steps[next_noise_level]
+
+        if self.is_ddim_sampling:
+            return self.ddim_sample_step(
+                x=x,
+                action_cond=action_cond,
+                pose_cond=pose_cond,
+                curr_noise_level=curr_noise_level,
+                next_noise_level=next_noise_level,
+                guidance_fn=guidance_fn,
+                current_frame=current_frame,
+                mode=mode,
+                reference_length=reference_length,
+                frame_idx=frame_idx,
+                memory_tokens=memory_tokens,
+                memory_token_mask=memory_token_mask,
+                memory_retrieval_tokens=memory_retrieval_tokens,
+                memory_retrieval_mask=memory_retrieval_mask,
+                **memory_kwargs,
+            )
+
+        # FIXME: temporary code for checking ddpm sampling
+        assert torch.all(
+            (curr_noise_level - 1 == next_noise_level) | ((curr_noise_level == -1) & (next_noise_level == -1))
+        ), "Wrong noise level given for ddpm sampling."
+
+        assert (
+            self.sampling_timesteps == self.timesteps
+        ), "sampling_timesteps should be equal to timesteps for ddpm sampling."
+
+        return self.ddpm_sample_step(
+            x=x,
+            action_cond=action_cond,
+            pose_cond=pose_cond,
+            curr_noise_level=curr_noise_level,
+            guidance_fn=guidance_fn,
+            reference_length=reference_length,
+            frame_idx=frame_idx,
+            memory_tokens=memory_tokens,
+            memory_token_mask=memory_token_mask,
+            memory_retrieval_tokens=memory_retrieval_tokens,
+            memory_retrieval_mask=memory_retrieval_mask,
+            **memory_kwargs,
+        )
+
+    def ddpm_sample_step(
+        self,
+        x: torch.Tensor,
+        action_cond: Optional[torch.Tensor],
+        pose_cond,
+        curr_noise_level: torch.Tensor,
+        guidance_fn: Optional[Callable] = None,
+        reference_length=None,
+        frame_idx=None,
+        memory_tokens=None,
+        memory_token_mask=None,
+        memory_retrieval_tokens=None,
+        memory_retrieval_mask=None,
+        **memory_kwargs,
+    ):
+        clipped_curr_noise_level = torch.where(
+            curr_noise_level < 0,
+            torch.full_like(curr_noise_level, self.stabilization_level - 1, dtype=torch.long),
+            curr_noise_level,
+        )
+
+        # treating as stabilization would require us to scale with sqrt of alpha_cum
+        orig_x = x.clone().detach()
+        scaled_context = self.q_sample(
+            x,
+            clipped_curr_noise_level,
+            noise=torch.zeros_like(x),
+        )
+        x = torch.where(self.add_shape_channels(curr_noise_level < 0), scaled_context, orig_x)
+
+        if guidance_fn is not None:
+            raise NotImplementedError("Guidance function is not implemented for ddpm sampling yet.")
+
+        else:
+            model_mean, _, model_log_variance = self.p_mean_variance(
+                x=x,
+                t=clipped_curr_noise_level,
+                action_cond=action_cond,
+                pose_cond=pose_cond,
+                reference_length=reference_length,
+                frame_idx=frame_idx,
+                memory_tokens=memory_tokens,
+                memory_token_mask=memory_token_mask,
+                memory_retrieval_tokens=memory_retrieval_tokens,
+                memory_retrieval_mask=memory_retrieval_mask,
+                **memory_kwargs,
+            )
+
+        noise = torch.where(
+            self.add_shape_channels(clipped_curr_noise_level > 0),
+            torch.randn_like(x),
+            0,
+        )
+        noise = torch.clamp(noise, -self.clip_noise, self.clip_noise)
+        x_pred = model_mean + torch.exp(0.5 * model_log_variance) * noise
+
+        # only update frames where the noise level decreases
+        return torch.where(self.add_shape_channels(curr_noise_level == -1), orig_x, x_pred)
+
+    def ddim_sample_step(
+        self,
+        x: torch.Tensor,
+        action_cond: Optional[torch.Tensor],
+        pose_cond,
+        curr_noise_level: torch.Tensor,
+        next_noise_level: torch.Tensor,
+        guidance_fn: Optional[Callable] = None,
+        current_frame=None,
+        mode="training",
+        reference_length=None,
+        frame_idx=None,
+        memory_tokens=None,
+        memory_token_mask=None,
+        memory_retrieval_tokens=None,
+        memory_retrieval_mask=None,
+        **memory_kwargs,
+    ):
+        # convert noise level -1 to self.stabilization_level - 1
+        clipped_curr_noise_level = torch.where(
+            curr_noise_level < 0,
+            torch.full_like(curr_noise_level, self.stabilization_level - 1, dtype=torch.long),
+            curr_noise_level,
+        )
+
+        # treating as stabilization would require us to scale with sqrt of alpha_cum
+        orig_x = x.clone().detach()
+        scaled_context = self.q_sample(
+            x,
+            clipped_curr_noise_level,
+            noise=torch.zeros_like(x),
+        )
+        x = torch.where(self.add_shape_channels(curr_noise_level < 0), scaled_context, orig_x)
+
+        alpha = self.alphas_cumprod[clipped_curr_noise_level]
+        alpha_next = torch.where(
+            next_noise_level < 0,
+            torch.ones_like(next_noise_level),
+            self.alphas_cumprod[next_noise_level],
+        )
+        sigma = torch.where(
+            next_noise_level < 0,
+            torch.zeros_like(next_noise_level),
+            self.ddim_sampling_eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt(),
+        )
+        c = (1 - alpha_next - sigma**2).sqrt()
+
+        alpha_next = self.add_shape_channels(alpha_next)
+        c = self.add_shape_channels(c)
+        sigma = self.add_shape_channels(sigma)
+
+        if guidance_fn is not None:
+            with torch.enable_grad():
+                x = x.detach().requires_grad_()
+
+                model_pred = self.model_predictions(
+                    x=x,
+                    t=clipped_curr_noise_level,
+                    action_cond=action_cond,
+                    pose_cond=pose_cond,
+                    current_frame=current_frame,
+                    mode=mode,
+                    reference_length=reference_length,
+                    frame_idx=frame_idx,
+                    memory_tokens=memory_tokens,
+                    memory_token_mask=memory_token_mask,
+                    memory_retrieval_tokens=memory_retrieval_tokens,
+                    memory_retrieval_mask=memory_retrieval_mask,
+                    **memory_kwargs,
+                )
+
+                guidance_loss = guidance_fn(model_pred.pred_x_start)
+                grad = -torch.autograd.grad(
+                    guidance_loss,
+                    x,
+                )[0]
+
+                pred_noise = model_pred.pred_noise + (1 - alpha_next).sqrt() * grad
+                x_start = self.predict_start_from_noise(x, clipped_curr_noise_level, pred_noise)
+
+        else:
+            # print(clipped_curr_noise_level)
+            model_pred = self.model_predictions(
+                x=x,
+                t=clipped_curr_noise_level,
+                action_cond=action_cond,
+                pose_cond=pose_cond,
+                current_frame=current_frame,
+                mode=mode,
+                reference_length=reference_length,
+                frame_idx=frame_idx,
+                memory_tokens=memory_tokens,
+                memory_token_mask=memory_token_mask,
+                memory_retrieval_tokens=memory_retrieval_tokens,
+                memory_retrieval_mask=memory_retrieval_mask,
+                **memory_kwargs,
+            )
+            x_start = model_pred.pred_x_start
+            pred_noise = model_pred.pred_noise
+
+        noise = torch.randn_like(x)
+        noise = torch.clamp(noise, -self.clip_noise, self.clip_noise)
+
+        x_pred = x_start * alpha_next.sqrt() + pred_noise * c + sigma * noise
+
+        # only update frames where the noise level decreases
+        mask = curr_noise_level == next_noise_level
+        x_pred = torch.where(
+            self.add_shape_channels(mask),
+            orig_x,
+            x_pred,
+        )
+
+        return x_pred

algorithms/worldmem/models/dit.py

@@ -0,0 +1,899 @@
+"""
+References:
+    - DiT: https://github.com/facebookresearch/DiT/blob/main/models.py
+    - Diffusion Forcing: https://github.com/buoyancy99/diffusion-forcing/blob/main/algorithms/diffusion_forcing/models/unet3d.py
+    - Latte: https://github.com/Vchitect/Latte/blob/main/models/latte.py
+"""
+
+from typing import Optional, Literal
+import torch
+from torch import nn
+from .rotary_embedding_torch import RotaryEmbedding
+from einops import rearrange
+from .attention import SpatialAxialAttention, TemporalAxialAttention
+from timm.models.vision_transformer import Mlp
+from timm.layers.helpers import to_2tuple
+import math
+from collections import namedtuple
+from typing import Optional, Callable
+
+def modulate(x, shift, scale):
+    fixed_dims = [1] * len(shift.shape[1:])
+    shift = shift.repeat(x.shape[0] // shift.shape[0], *fixed_dims)
+    scale = scale.repeat(x.shape[0] // scale.shape[0], *fixed_dims)
+    while shift.dim() < x.dim():
+        shift = shift.unsqueeze(-2)
+        scale = scale.unsqueeze(-2)
+    return x * (1 + scale) + shift
+
+def gate(x, g):
+    fixed_dims = [1] * len(g.shape[1:])
+    g = g.repeat(x.shape[0] // g.shape[0], *fixed_dims)
+    while g.dim() < x.dim():
+        g = g.unsqueeze(-2)
+    return g * x
+
+
+class PatchEmbed(nn.Module):
+    """2D Image to Patch Embedding"""
+
+    def __init__(
+        self,
+        img_height=256,
+        img_width=256,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        norm_layer=None,
+        flatten=True,
+    ):
+        super().__init__()
+        img_size = (img_height, img_width)
+        patch_size = to_2tuple(patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x, random_sample=False):
+        B, C, H, W = x.shape
+        assert random_sample or (H == self.img_size[0] and W == self.img_size[1]), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        
+        x = self.proj(x)
+        if self.flatten:
+            x = rearrange(x, "B C H W -> B (H W) C")
+        else:
+            x = rearrange(x, "B C H W -> B H W C")
+        x = self.norm(x)
+        return x
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+
+    def __init__(self, hidden_size, frequency_embedding_size=256, freq_type='time_step'):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),  # hidden_size is diffusion model hidden size
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+        self.freq_type = freq_type
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000, freq_type='time_step'):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+
+        if freq_type == 'time_step':
+            freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(device=t.device)
+        elif freq_type == 'spatial': # ~(-5 5)
+            freqs = torch.linspace(1.0, half, half).to(device=t.device) * torch.pi
+        elif freq_type == 'angle': # 0-360
+            freqs = torch.linspace(1.0, half, half).to(device=t.device) * torch.pi / 180
+
+
+        args = t[:, None].float() * freqs[None]
+        
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size, freq_type=self.freq_type)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer of DiT.
+    """
+
+    def __init__(self, hidden_size, patch_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+MEMORY_TYPE_NAMES = ("anchor", "dynamic", "revisit")
+MEMORY_TYPE_ANCHOR = 0
+MEMORY_TYPE_DYNAMIC = 1
+MEMORY_TYPE_REVISIT = 2
+
+
+class MemoryTokenCrossAttention(nn.Module):
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, num_memory_types=3):
+        super().__init__()
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.num_heads = num_heads
+        self.num_memory_types = num_memory_types
+        self.norm_q = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm_mem = nn.LayerNorm(hidden_size, eps=1e-6)
+        self.attn = nn.MultiheadAttention(hidden_size, num_heads, batch_first=True)
+        self.norm_mlp = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp = Mlp(
+            in_features=hidden_size,
+            hidden_features=mlp_hidden_dim,
+            act_layer=approx_gelu,
+            drop=0,
+        )
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
+        self.memory_type_embed = nn.Embedding(num_memory_types, hidden_size)
+        self.memory_type_scale = nn.Parameter(torch.ones(num_memory_types, hidden_size))
+        self.memory_type_gate = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, num_memory_types, bias=True))
+        self.last_gate_mean = None
+        self.last_delta_ratio = None
+        self.last_valid_fraction = None
+        self.last_type_gate_mean = None
+        for type_name in MEMORY_TYPE_NAMES[:num_memory_types]:
+            setattr(self, f"last_type_gate_{type_name}_mean", None)
+        nn.init.normal_(self.memory_type_embed.weight, std=0.02)
+        self.reset_identity_init()
+
+    def reset_identity_init(self):
+        nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.memory_type_gate[-1].weight, 0)
+        nn.init.constant_(self.memory_type_gate[-1].bias, 0)
+
+    def _attend(self, query, memory_tokens, memory_token_mask=None, memory_token_gate=None):
+        if memory_token_mask is None and memory_token_gate is None:
+            out, _ = self.attn(query, memory_tokens, memory_tokens, need_weights=False)
+            return out, None
+
+        if memory_token_mask is None:
+            memory_token_mask = torch.ones(
+                memory_tokens.shape[:2],
+                device=memory_tokens.device,
+                dtype=torch.bool,
+            )
+        else:
+            memory_token_mask = memory_token_mask.bool()
+        gate_tensor = None
+        if memory_token_gate is not None:
+            if tuple(memory_token_gate.shape) != tuple(memory_tokens.shape[:2]):
+                raise ValueError(
+                    f"memory_token_gate must have shape {tuple(memory_tokens.shape[:2])}, "
+                    f"got {tuple(memory_token_gate.shape)}"
+                )
+            gate_tensor = memory_token_gate.to(device=memory_tokens.device, dtype=query.dtype)
+            memory_token_mask = memory_token_mask & (gate_tensor > 0)
+        valid_rows = memory_token_mask.any(dim=1)
+        out = torch.zeros_like(query)
+        if valid_rows.any():
+            attn_mask = None
+            key_padding_mask = ~memory_token_mask[valid_rows]
+            if gate_tensor is not None:
+                gate_bias = torch.log(gate_tensor[valid_rows].clamp_min(1.0e-6))
+                gate_bias = gate_bias[:, None, :].expand(-1, query.shape[1], -1)
+                attn_mask = gate_bias.repeat_interleave(self.num_heads, dim=0)
+                float_padding_mask = torch.zeros_like(gate_tensor[valid_rows], dtype=query.dtype)
+                key_padding_mask = float_padding_mask.masked_fill(key_padding_mask, float("-inf"))
+            attended, _ = self.attn(
+                query[valid_rows],
+                memory_tokens[valid_rows],
+                memory_tokens[valid_rows],
+                key_padding_mask=key_padding_mask,
+                attn_mask=attn_mask,
+                need_weights=False,
+            )
+            out[valid_rows] = attended.to(out.dtype)
+        return out, valid_rows
+
+    def _apply_memory_type(self, memory_tokens, memory_type_ids):
+        if memory_type_ids is None:
+            return memory_tokens
+        memory_type_ids = memory_type_ids.to(device=memory_tokens.device, dtype=torch.long)
+        type_embed = self.memory_type_embed(memory_type_ids).to(memory_tokens.dtype)
+        type_scale = self.memory_type_scale[memory_type_ids].to(memory_tokens.dtype)
+        while type_embed.dim() < memory_tokens.dim():
+            type_embed = type_embed.unsqueeze(0)
+            type_scale = type_scale.unsqueeze(0)
+        return memory_tokens * type_scale + type_embed
+
+    def _store_type_gate_diagnostics(self, stage_gate):
+        with torch.no_grad():
+            detached = stage_gate.detach().float()
+            self.last_type_gate_mean = detached.mean()
+            for type_idx, type_name in enumerate(MEMORY_TYPE_NAMES[: self.num_memory_types]):
+                setattr(self, f"last_type_gate_{type_name}_mean", detached[..., type_idx].mean())
+
+    def _type_stage_gate(self, c, memory_tokens, memory_type_ids):
+        if memory_type_ids is None:
+            return None
+        memory_type_ids = memory_type_ids.to(device=memory_tokens.device, dtype=torch.long)
+        stage_gate = torch.sigmoid(self.memory_type_gate(c)).to(memory_tokens.dtype)
+        self._store_type_gate_diagnostics(stage_gate)
+        if memory_tokens.dim() == 4:
+            batch_size, num_frames, num_tokens = memory_tokens.shape[:3]
+            if memory_type_ids.dim() == 1:
+                gather_ids = memory_type_ids.view(1, 1, num_tokens).expand(batch_size, num_frames, num_tokens)
+            elif tuple(memory_type_ids.shape) == (batch_size, num_frames, num_tokens):
+                gather_ids = memory_type_ids
+            else:
+                raise ValueError(
+                    "rank-4 memory_type_ids must have shape (M,) or (B,T,M), "
+                    f"got {tuple(memory_type_ids.shape)}"
+                )
+            return torch.gather(stage_gate, dim=-1, index=gather_ids)
+        if memory_tokens.dim() == 3:
+            batch_size, num_tokens = memory_tokens.shape[:2]
+            if memory_type_ids.dim() != 1:
+                raise ValueError("rank-3 memory_type_ids must have shape (M,)")
+            gather_ids = memory_type_ids.view(1, 1, num_tokens).expand(batch_size, stage_gate.shape[1], num_tokens)
+            return torch.gather(stage_gate, dim=-1, index=gather_ids).mean(dim=1)
+        raise ValueError(f"memory_tokens must be rank 3 or 4, got rank {memory_tokens.dim()}")
+
+    def _combine_memory_gate(self, memory_tokens, memory_token_gate, type_stage_gate):
+        combined_gate = type_stage_gate
+        if memory_token_gate is not None:
+            if tuple(memory_token_gate.shape) != tuple(memory_tokens.shape[:-1]):
+                raise ValueError(
+                    f"memory_token_gate must have shape {tuple(memory_tokens.shape[:-1])}, "
+                    f"got {tuple(memory_token_gate.shape)}"
+                )
+            stream_gate = memory_token_gate.to(device=memory_tokens.device, dtype=memory_tokens.dtype)
+            combined_gate = stream_gate if combined_gate is None else combined_gate * stream_gate
+        return combined_gate
+
+    def _valid_mask(self, valid_rows, batch_size, num_frames, dtype, device):
+        if valid_rows is None:
+            return None
+        valid_rows = valid_rows.to(device=device, dtype=dtype)
+        if valid_rows.numel() == batch_size:
+            return valid_rows.view(batch_size, 1, 1, 1, 1)
+        if valid_rows.numel() == batch_size * num_frames:
+            return rearrange(valid_rows, "(b t) -> b t", b=batch_size, t=num_frames)[:, :, None, None, None]
+        raise ValueError(f"valid_rows has incompatible shape: {tuple(valid_rows.shape)}")
+
+    def _gate_valid_mask(self, valid_rows, batch_size, num_frames, dtype, device):
+        if valid_rows is None:
+            return None
+        valid_rows = valid_rows.to(device=device, dtype=dtype)
+        if valid_rows.numel() == batch_size:
+            return valid_rows.view(batch_size, 1, 1)
+        if valid_rows.numel() == batch_size * num_frames:
+            return rearrange(valid_rows, "(b t) -> b t", b=batch_size, t=num_frames)[:, :, None]
+        raise ValueError(f"valid_rows has incompatible shape: {tuple(valid_rows.shape)}")
+
+    def _residual_gate(self, residual_gate, batch_size, num_frames, dtype, device):
+        if residual_gate is None:
+            return None
+        if not torch.is_tensor(residual_gate):
+            return torch.tensor(float(residual_gate), dtype=dtype, device=device).view(1, 1, 1, 1, 1)
+        gate_tensor = residual_gate.to(device=device, dtype=dtype)
+        if gate_tensor.dim() == 0:
+            gate_tensor = gate_tensor.view(1, 1, 1, 1, 1)
+        elif gate_tensor.dim() == 1:
+            if gate_tensor.numel() == batch_size:
+                gate_tensor = gate_tensor.view(batch_size, 1, 1, 1, 1)
+            elif gate_tensor.numel() == batch_size * num_frames:
+                gate_tensor = rearrange(gate_tensor, "(b t) -> b t", b=batch_size, t=num_frames)[:, :, None, None, None]
+            else:
+                raise ValueError(f"residual_gate has incompatible shape: {tuple(gate_tensor.shape)}")
+        elif gate_tensor.dim() == 2:
+            if tuple(gate_tensor.shape) != (batch_size, num_frames):
+                raise ValueError(f"residual_gate must have shape (B,T), got {tuple(gate_tensor.shape)}")
+            gate_tensor = gate_tensor[:, :, None, None, None]
+        elif gate_tensor.dim() == 3:
+            if tuple(gate_tensor.shape[:2]) != (batch_size, num_frames):
+                raise ValueError(f"residual_gate must start with (B,T), got {tuple(gate_tensor.shape)}")
+            gate_tensor = gate_tensor[:, :, :, None, None]
+        else:
+            while gate_tensor.dim() < 5:
+                gate_tensor = gate_tensor.unsqueeze(-1)
+        return gate_tensor
+
+    def _store_diagnostics(self, output, base, gate_msa, gate_mlp, valid_rows):
+        with torch.no_grad():
+            batch_size, num_frames = base.shape[:2]
+            gate_values = torch.cat(
+                [gate_msa.detach().float().abs(), gate_mlp.detach().float().abs()],
+                dim=-1,
+            )
+            gate_mask = self._gate_valid_mask(
+                valid_rows,
+                batch_size,
+                num_frames,
+                dtype=gate_values.dtype,
+                device=gate_values.device,
+            )
+            if gate_mask is not None:
+                gate_values = gate_values * gate_mask
+                self.last_valid_fraction = valid_rows.detach().float().mean()
+                valid_count = (gate_mask.sum() * gate_values.shape[-1]).clamp_min(1.0)
+                self.last_gate_mean = gate_values.sum() / valid_count
+            else:
+                self.last_valid_fraction = base.detach().new_tensor(1.0, dtype=torch.float32)
+                self.last_gate_mean = gate_values.mean()
+
+            delta_norm = (output.detach().float() - base.detach().float()).norm()
+            base_norm = base.detach().float().norm()
+            self.last_delta_ratio = delta_norm / (base_norm + 1e-6)
+
+    def forward(
+        self,
+        x,
+        c,
+        memory_tokens,
+        memory_token_mask=None,
+        residual_base=None,
+        return_delta=False,
+        residual_gate=None,
+        memory_type_ids=None,
+        memory_token_gate=None,
+    ):
+        B, T, H, W, D = x.shape
+        if residual_base is None:
+            residual_base = x
+        m_shift_msa, m_scale_msa, m_gate_msa, m_shift_mlp, m_scale_mlp, m_gate_mlp = (
+            self.adaLN_modulation(c).chunk(6, dim=-1)
+        )
+        query_source = modulate(self.norm_q(x), m_shift_msa, m_scale_msa)
+        type_stage_gate = self._type_stage_gate(c, memory_tokens, memory_type_ids)
+        effective_token_gate = self._combine_memory_gate(memory_tokens, memory_token_gate, type_stage_gate)
+        if memory_tokens.dim() == 3:
+            query = rearrange(query_source, "b t h w d -> b (t h w) d")
+            memory_tokens = self._apply_memory_type(self.norm_mem(memory_tokens), memory_type_ids)
+            valid_rows = None
+            if memory_token_mask is not None:
+                if tuple(memory_token_mask.shape) != tuple(memory_tokens.shape[:2]):
+                    raise ValueError(
+                        f"legacy memory mask must have shape {tuple(memory_tokens.shape[:2])}, "
+                        f"got {tuple(memory_token_mask.shape)}"
+                    )
+            out, valid_rows = self._attend(
+                query,
+                memory_tokens,
+                memory_token_mask=memory_token_mask,
+                memory_token_gate=effective_token_gate,
+            )
+            out = rearrange(out, "b (t h w) d -> b t h w d", t=T, h=H, w=W)
+        elif memory_tokens.dim() == 4:
+            assert memory_tokens.shape[:2] == (B, T), (
+                f"per-frame memory tokens must have shape (B, T, M, D), got {tuple(memory_tokens.shape)}"
+            )
+            query = rearrange(query_source, "b t h w d -> (b t) (h w) d")
+            memory_tokens = self._apply_memory_type(self.norm_mem(memory_tokens), memory_type_ids)
+            memory_tokens = rearrange(memory_tokens, "b t m d -> (b t) m d")
+            if effective_token_gate is not None:
+                effective_token_gate = rearrange(effective_token_gate, "b t m -> (b t) m")
+            valid_rows = None
+            if memory_token_mask is not None:
+                expected_mask_shape = (B, T, memory_tokens.shape[1])
+                if tuple(memory_token_mask.shape) != expected_mask_shape:
+                    raise ValueError(
+                        f"per-frame memory mask must have shape {expected_mask_shape}, "
+                        f"got {tuple(memory_token_mask.shape)}"
+                    )
+                memory_token_mask = rearrange(memory_token_mask.bool(), "b t m -> (b t) m")
+            out, valid_rows = self._attend(
+                query,
+                memory_tokens,
+                memory_token_mask=memory_token_mask,
+                memory_token_gate=effective_token_gate,
+            )
+            out = rearrange(out, "(b t) (h w) d -> b t h w d", b=B, t=T, h=H, w=W)
+        else:
+            raise ValueError(f"memory_tokens must be rank 3 or 4, got rank {memory_tokens.dim()}")
+
+        valid_mask = self._valid_mask(valid_rows, B, T, dtype=out.dtype, device=out.device)
+        residual_gate_tensor = self._residual_gate(residual_gate, B, T, dtype=out.dtype, device=out.device)
+        attn_delta = gate(out, m_gate_msa)
+        if valid_mask is not None:
+            attn_delta = attn_delta * valid_mask
+        if residual_gate_tensor is not None:
+            attn_delta = attn_delta * residual_gate_tensor
+        output = residual_base + attn_delta
+
+        mlp_delta = gate(self.mlp(modulate(self.norm_mlp(output), m_shift_mlp, m_scale_mlp)), m_gate_mlp)
+        if valid_mask is not None:
+            mlp_delta = mlp_delta * valid_mask
+        if residual_gate_tensor is not None:
+            mlp_delta = mlp_delta * residual_gate_tensor
+        output = output + mlp_delta
+        self._store_diagnostics(output, residual_base, m_gate_msa, m_gate_mlp, valid_rows)
+        if return_delta:
+            return attn_delta + mlp_delta
+        return output
+
+class SpatioTemporalDiTBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        num_heads,
+        reference_length,
+        mlp_ratio=4.0,
+        is_causal=True,
+        spatial_rotary_emb: Optional[RotaryEmbedding] = None,
+        temporal_rotary_emb: Optional[RotaryEmbedding] = None,
+        use_memory_token_cross_attention=False,
+        ref_mode='sequential'
+    ):
+        super().__init__()
+        self.is_causal = is_causal
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+
+        self.s_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.s_attn = SpatialAxialAttention(
+            hidden_size,
+            heads=num_heads,
+            dim_head=hidden_size // num_heads,
+            rotary_emb=spatial_rotary_emb
+        )
+        self.s_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.s_mlp = Mlp(
+            in_features=hidden_size,
+            hidden_features=mlp_hidden_dim,
+            act_layer=approx_gelu,
+            drop=0,
+        )
+        self.s_adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
+
+        self.t_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.t_attn = TemporalAxialAttention(
+            hidden_size,
+            heads=num_heads,
+            dim_head=hidden_size // num_heads,
+            is_causal=is_causal,
+            rotary_emb=temporal_rotary_emb,
+            reference_length=reference_length
+        )
+        self.t_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.t_mlp = Mlp(
+            in_features=hidden_size,
+            hidden_features=mlp_hidden_dim,
+            act_layer=approx_gelu,
+            drop=0,
+        )
+        self.t_adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
+
+        self.reference_length = reference_length
+        self.use_memory_token_cross_attention = use_memory_token_cross_attention
+        if self.use_memory_token_cross_attention:
+            self.memory_token_cross_attn = MemoryTokenCrossAttention(hidden_size, num_heads, mlp_ratio=mlp_ratio)
+
+        self.ref_mode = ref_mode
+
+        if self.ref_mode == 'parallel':
+            self.parallel_map = nn.Linear(hidden_size, hidden_size)
+
+    def _expand_memory_stream(self, tokens, mask, stream_gate, type_idx, batch_size, num_frames):
+        if tokens is None or tokens.shape[-2] == 0:
+            return None
+        if tokens.dim() == 3:
+            if tokens.shape[0] != batch_size:
+                raise ValueError(f"rank-3 memory tokens must start with B={batch_size}, got {tuple(tokens.shape)}")
+            tokens = tokens[:, None].expand(-1, num_frames, -1, -1)
+            if mask is None:
+                mask = torch.ones(tokens.shape[:3], device=tokens.device, dtype=torch.bool)
+            elif mask.dim() == 2:
+                mask = mask[:, None].expand(-1, num_frames, -1)
+            elif mask.dim() != 3:
+                raise ValueError(f"rank-3 stream mask must have rank 2 or 3, got {tuple(mask.shape)}")
+        elif tokens.dim() == 4:
+            if tuple(tokens.shape[:2]) != (batch_size, num_frames):
+                raise ValueError(
+                    f"rank-4 memory tokens must start with (B,T)={(batch_size, num_frames)}, "
+                    f"got {tuple(tokens.shape)}"
+                )
+            if mask is None:
+                mask = torch.ones(tokens.shape[:3], device=tokens.device, dtype=torch.bool)
+            elif mask.dim() != 3:
+                raise ValueError(f"rank-4 stream mask must have rank 3, got {tuple(mask.shape)}")
+        else:
+            raise ValueError(f"memory stream tokens must be rank 3 or 4, got rank {tokens.dim()}")
+        if tuple(mask.shape) != tuple(tokens.shape[:3]):
+            raise ValueError(f"memory stream mask must have shape {tuple(tokens.shape[:3])}, got {tuple(mask.shape)}")
+        gate_tensor = self._expand_memory_stream_gate(stream_gate, tokens)
+        type_ids = torch.full((tokens.shape[2],), int(type_idx), device=tokens.device, dtype=torch.long)
+        return tokens, mask.to(device=tokens.device, dtype=torch.bool), gate_tensor, type_ids
+
+    def _expand_memory_stream_gate(self, stream_gate, tokens):
+        batch_size, num_frames, num_tokens = tokens.shape[:3]
+        if stream_gate is None:
+            return torch.ones(tokens.shape[:3], device=tokens.device, dtype=tokens.dtype)
+        if not torch.is_tensor(stream_gate):
+            return torch.full(tokens.shape[:3], float(stream_gate), device=tokens.device, dtype=tokens.dtype)
+        gate_tensor = stream_gate.to(device=tokens.device, dtype=tokens.dtype)
+        if gate_tensor.dim() == 0:
+            return gate_tensor.view(1, 1, 1).expand(batch_size, num_frames, num_tokens)
+        if gate_tensor.dim() == 1:
+            if gate_tensor.numel() != batch_size:
+                raise ValueError(f"rank-1 memory gate must have B={batch_size} values, got {tuple(gate_tensor.shape)}")
+            return gate_tensor.view(batch_size, 1, 1).expand(batch_size, num_frames, num_tokens)
+        if gate_tensor.dim() == 2:
+            if tuple(gate_tensor.shape) == (batch_size, num_frames):
+                return gate_tensor[:, :, None].expand(batch_size, num_frames, num_tokens)
+            if tuple(gate_tensor.shape) == (batch_size, num_tokens):
+                return gate_tensor[:, None, :].expand(batch_size, num_frames, num_tokens)
+            raise ValueError(
+                f"rank-2 memory gate must have shape (B,T) or (B,M), got {tuple(gate_tensor.shape)}"
+            )
+        if gate_tensor.dim() == 3:
+            if tuple(gate_tensor.shape) == (batch_size, num_frames, 1):
+                return gate_tensor.expand(batch_size, num_frames, num_tokens)
+            if tuple(gate_tensor.shape) == (batch_size, num_frames, num_tokens):
+                return gate_tensor
+            raise ValueError(
+                f"rank-3 memory gate must have shape (B,T,1) or (B,T,M), got {tuple(gate_tensor.shape)}"
+            )
+        raise ValueError(f"memory gate rank must be <=3, got rank {gate_tensor.dim()}")
+
+    def _pack_typed_memory_streams(
+        self,
+        batch_size,
+        num_frames,
+        memory_tokens=None,
+        memory_token_mask=None,
+        memory_dynamic_tokens=None,
+        memory_dynamic_mask=None,
+        memory_retrieval_tokens=None,
+        memory_retrieval_mask=None,
+        memory_anchor_gate=None,
+        memory_dynamic_gate=None,
+        memory_retrieval_gate=None,
+    ):
+        streams = []
+        for tokens, mask, stream_gate, type_idx in (
+            (memory_tokens, memory_token_mask, memory_anchor_gate, MEMORY_TYPE_ANCHOR),
+            (memory_dynamic_tokens, memory_dynamic_mask, memory_dynamic_gate, MEMORY_TYPE_DYNAMIC),
+            (memory_retrieval_tokens, memory_retrieval_mask, memory_retrieval_gate, MEMORY_TYPE_REVISIT),
+        ):
+            expanded = self._expand_memory_stream(tokens, mask, stream_gate, type_idx, batch_size, num_frames)
+            if expanded is not None:
+                streams.append(expanded)
+        if not streams:
+            return None
+        packed_tokens = torch.cat([item[0] for item in streams], dim=2)
+        packed_mask = torch.cat([item[1] for item in streams], dim=2)
+        packed_gate = torch.cat([item[2] for item in streams], dim=2)
+        packed_type_ids = torch.cat([item[3] for item in streams], dim=0)
+        valid_gate = packed_gate.masked_fill(~packed_mask, 0)
+        residual_gate = valid_gate.max(dim=2).values
+        return packed_tokens, packed_mask, packed_gate, packed_type_ids, residual_gate
+
+    def forward(self, x, c, current_frame=None, timestep=None, is_last_block=False, 
+        pose_cond=None, mode="training", c_action_cond=None, reference_length=None,
+        memory_tokens=None, memory_token_mask=None, memory_dynamic_tokens=None, memory_dynamic_mask=None,
+        memory_retrieval_tokens=None, memory_retrieval_mask=None, memory_anchor_gate=None,
+        memory_dynamic_gate=None, memory_retrieval_gate=None):
+        B, T, H, W, D = x.shape
+
+        # spatial block
+        
+        s_shift_msa, s_scale_msa, s_gate_msa, s_shift_mlp, s_scale_mlp, s_gate_mlp = self.s_adaLN_modulation(c).chunk(6, dim=-1)
+        x = x + gate(self.s_attn(modulate(self.s_norm1(x), s_shift_msa, s_scale_msa)), s_gate_msa)
+        x = x + gate(self.s_mlp(modulate(self.s_norm2(x), s_shift_mlp, s_scale_mlp)), s_gate_mlp)
+
+        # temporal block
+        if c_action_cond is not None:
+            t_shift_msa, t_scale_msa, t_gate_msa, t_shift_mlp, t_scale_mlp, t_gate_mlp = self.t_adaLN_modulation(c_action_cond).chunk(6, dim=-1)
+        else:
+            t_shift_msa, t_scale_msa, t_gate_msa, t_shift_mlp, t_scale_mlp, t_gate_mlp = self.t_adaLN_modulation(c).chunk(6, dim=-1)
+        
+        x_t = x + gate(self.t_attn(modulate(self.t_norm1(x), t_shift_msa, t_scale_msa)), t_gate_msa)
+        x_t = x_t + gate(self.t_mlp(modulate(self.t_norm2(x_t), t_shift_mlp, t_scale_mlp)), t_gate_mlp)
+
+        if self.ref_mode == 'sequential':
+            x = x_t
+
+        if self.use_memory_token_cross_attention:
+            memory_base = x
+            packed_memory = self._pack_typed_memory_streams(
+                B,
+                T,
+                memory_tokens=memory_tokens,
+                memory_token_mask=memory_token_mask,
+                memory_dynamic_tokens=memory_dynamic_tokens,
+                memory_dynamic_mask=memory_dynamic_mask,
+                memory_retrieval_tokens=memory_retrieval_tokens,
+                memory_retrieval_mask=memory_retrieval_mask,
+                memory_anchor_gate=memory_anchor_gate,
+                memory_dynamic_gate=memory_dynamic_gate,
+                memory_retrieval_gate=memory_retrieval_gate,
+            )
+            if packed_memory is not None:
+                packed_tokens, packed_mask, packed_gate, packed_type_ids, residual_gate = packed_memory
+                x = self.memory_token_cross_attn(
+                    memory_base,
+                    c,
+                    packed_tokens,
+                    packed_mask,
+                    residual_gate=residual_gate,
+                    memory_type_ids=packed_type_ids,
+                    memory_token_gate=packed_gate,
+                )
+
+        if self.ref_mode == 'parallel':
+            x = x_t + self.parallel_map(x)
+
+        return x
+
+
+class DiT(nn.Module):
+    """
+    Diffusion model with a Transformer backbone.
+    """
+
+    def __init__(
+        self,
+        input_h=18,
+        input_w=32,
+        patch_size=2,
+        in_channels=16,
+        hidden_size=1024,
+        depth=12,
+        num_heads=16,
+        mlp_ratio=4.0,
+        action_cond_dim=25,
+        max_frames=32,
+        reference_length=8,
+        memory_token_cross_attention=False,
+        memory_cross_attn_layers=None,
+        ref_mode='sequential'
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+        self.max_frames = max_frames
+
+        self.x_embedder = PatchEmbed(input_h, input_w, patch_size, in_channels, hidden_size, flatten=False)
+        self.t_embedder = TimestepEmbedder(hidden_size)
+
+        self.spatial_rotary_emb = RotaryEmbedding(dim=hidden_size // num_heads // 2, freqs_for="pixel", max_freq=256)
+        self.temporal_rotary_emb = RotaryEmbedding(dim=hidden_size // num_heads)
+
+        self.external_cond = nn.Linear(action_cond_dim, hidden_size) if action_cond_dim > 0 else nn.Identity()
+        if memory_cross_attn_layers is None:
+            memory_cross_attn_layer_set = None
+        else:
+            memory_cross_attn_layer_set = {int(layer_idx) for layer_idx in memory_cross_attn_layers}
+            invalid_layers = sorted(
+                layer_idx for layer_idx in memory_cross_attn_layer_set if layer_idx < 0 or layer_idx >= depth
+            )
+            if invalid_layers:
+                raise ValueError(
+                    f"memory_cross_attn_layers contains invalid indices {invalid_layers} for depth={depth}"
+                )
+
+        self.blocks = nn.ModuleList(
+            [
+                SpatioTemporalDiTBlock(
+                    hidden_size,
+                    num_heads,
+                    mlp_ratio=mlp_ratio,
+                    is_causal=True,
+                    reference_length=reference_length,
+                    spatial_rotary_emb=self.spatial_rotary_emb,
+                    temporal_rotary_emb=self.temporal_rotary_emb,
+                    use_memory_token_cross_attention=memory_token_cross_attention
+                    and (memory_cross_attn_layer_set is None or block_idx in memory_cross_attn_layer_set),
+                    ref_mode=ref_mode
+                )
+                for block_idx in range(depth)
+            ]
+        )
+        self.memory_token_cross_attention = memory_token_cross_attention
+        self.memory_cross_attn_layers = (
+            None if memory_cross_attn_layer_set is None else tuple(sorted(memory_cross_attn_layer_set))
+        )
+        self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        self.apply(_basic_init)
+
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
+        w = self.x_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.x_embedder.proj.bias, 0)
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers in DiT blocks:
+        for block in self.blocks:
+            nn.init.constant_(block.s_adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.s_adaLN_modulation[-1].bias, 0)
+            nn.init.constant_(block.t_adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.t_adaLN_modulation[-1].bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+        if self.memory_token_cross_attention:
+            for block in self.blocks:
+                memory_adapter = getattr(block, "memory_token_cross_attn", None)
+                if memory_adapter is not None:
+                    memory_adapter.reset_identity_init()
+
+    def memory_adapter_delta_diagnostics(self):
+        diagnostics = {}
+        ratios = []
+        type_gate_values = {type_name: [] for type_name in MEMORY_TYPE_NAMES}
+        shared_type_gate_values = []
+        for block in self.blocks:
+            adapter = getattr(block, "memory_token_cross_attn", None)
+            if adapter is None:
+                continue
+            ratio = getattr(adapter, "last_delta_ratio", None)
+            if ratio is not None:
+                ratios.append(torch.as_tensor(ratio).detach().float())
+            type_gate = getattr(adapter, "last_type_gate_mean", None)
+            if type_gate is not None:
+                shared_type_gate_values.append(torch.as_tensor(type_gate).detach().float())
+            for type_name in MEMORY_TYPE_NAMES:
+                value = getattr(adapter, f"last_type_gate_{type_name}_mean", None)
+                if value is not None:
+                    type_gate_values[type_name].append(torch.as_tensor(value).detach().float())
+        if ratios:
+            values = torch.stack(ratios)
+            diagnostics["memory_adapter_delta_ratio_max"] = float(values.max().item())
+            diagnostics["memory_adapter_delta_ratio_mean"] = float(values.mean().item())
+        if shared_type_gate_values:
+            values = torch.stack(shared_type_gate_values)
+            diagnostics["memory_adapter_type_gate_mean"] = float(values.mean().item())
+        for type_name, values_list in type_gate_values.items():
+            if values_list:
+                values = torch.stack(values_list)
+                diagnostics[f"memory_adapter_type_gate_{type_name}_mean"] = float(values.mean().item())
+        return diagnostics
+
+    def unpatchify(self, x):
+        """
+        x: (N, H, W, patch_size**2 * C)
+        imgs: (N, H, W, C)
+        """
+        c = self.out_channels
+        p = self.x_embedder.patch_size[0]
+        h = x.shape[1]
+        w = x.shape[2]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum("nhwpqc->nchpwq", x)
+        imgs = x.reshape(shape=(x.shape[0], c, h * p, w * p))
+        return imgs
+
+    def forward(
+        self,
+        x,
+        t,
+        action_cond=None,
+        pose_cond=None,
+        current_frame=None,
+        mode=None,
+        reference_length=None,
+        frame_idx=None,
+        memory_tokens=None,
+        memory_token_mask=None,
+        memory_dynamic_tokens=None,
+        memory_dynamic_mask=None,
+        memory_retrieval_tokens=None,
+        memory_retrieval_mask=None,
+        memory_anchor_gate=None,
+        memory_dynamic_gate=None,
+        memory_retrieval_gate=None,
+    ):
+        """
+        Forward pass of DiT.
+        x: (B, T, C, H, W) tensor of spatial inputs (images or latent representations of images)
+        t: (B, T,) tensor of diffusion timesteps
+        """
+
+        B, T, C, H, W = x.shape
+
+        # add spatial embeddings
+        x = rearrange(x, "b t c h w -> (b t) c h w")
+
+        x = self.x_embedder(x)  # (B*T, C, H, W) -> (B*T, H/2, W/2, D) , C = 16, D = d_model
+        # restore shape
+        x = rearrange(x, "(b t) h w d -> b t h w d", t=T)
+        # embed noise steps
+        t = rearrange(t, "b t -> (b t)")
+
+        c_t = self.t_embedder(t)  # (N, D)
+        c = c_t.clone()
+        c = rearrange(c, "(b t) d -> b t d", t=T)
+
+        if torch.is_tensor(action_cond):
+            c_action_cond = c + self.external_cond(action_cond)
+        else:
+            c_action_cond = None
+
+        for i, block in enumerate(self.blocks):
+            x = block(x, c, current_frame=current_frame, timestep=t, is_last_block= (i+1 == len(self.blocks)), 
+                mode=mode, c_action_cond=c_action_cond, reference_length=reference_length,
+                memory_tokens=memory_tokens, memory_token_mask=memory_token_mask,
+                memory_dynamic_tokens=memory_dynamic_tokens, memory_dynamic_mask=memory_dynamic_mask,
+                memory_retrieval_tokens=memory_retrieval_tokens, memory_retrieval_mask=memory_retrieval_mask,
+                memory_anchor_gate=memory_anchor_gate, memory_dynamic_gate=memory_dynamic_gate,
+                memory_retrieval_gate=memory_retrieval_gate)  # (N, T, H, W, D)
+        x = self.final_layer(x, c) # (N, T, H, W, patch_size ** 2 * out_channels)
+        # unpatchify
+        x = rearrange(x, "b t h w d -> (b t) h w d")
+        x = self.unpatchify(x)  # (N, out_channels, H, W)
+        x = rearrange(x, "(b t) c h w -> b t c h w", t=T)
+        return x
+
+
+def DiT_S_2(
+    action_cond_dim,
+    reference_length,
+    ref_mode,
+    memory_token_cross_attention=False,
+    memory_cross_attn_layers=None,
+):
+    return DiT(
+        patch_size=2,
+        hidden_size=1024,
+        depth=16,
+        num_heads=16,
+        action_cond_dim=action_cond_dim,
+        reference_length=reference_length,
+        memory_token_cross_attention=memory_token_cross_attention,
+        memory_cross_attn_layers=memory_cross_attn_layers,
+        ref_mode=ref_mode
+    )
+
+
+DiT_models = {"DiT-S/2": DiT_S_2}

algorithms/worldmem/models/pose_prediction.py

@@ -0,0 +1,42 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class PosePredictionNet(nn.Module):
+    def __init__(self, img_channels=16, img_feat_dim=256, pose_dim=5, action_dim=25, hidden_dim=128):
+        super(PosePredictionNet, self).__init__()
+        
+        self.cnn = nn.Sequential(
+            nn.Conv2d(img_channels, 32, kernel_size=3, stride=2, padding=1), 
+            nn.ReLU(),
+            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
+            nn.ReLU(),
+            nn.AdaptiveAvgPool2d((1, 1)) 
+        )
+        
+        self.fc_img = nn.Linear(128, img_feat_dim)  
+        
+        self.mlp_motion = nn.Sequential(
+            nn.Linear(pose_dim + action_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU()
+        )
+        
+        self.fc_out = nn.Sequential(
+            nn.Linear(img_feat_dim + hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, pose_dim)
+        )
+
+    def forward(self, img, action, pose):
+        img_feat = self.cnn(img).view(img.size(0), -1)
+        img_feat = self.fc_img(img_feat)
+        
+        motion_feat = self.mlp_motion(torch.cat([pose, action], dim=1))
+        fused_feat = torch.cat([img_feat, motion_feat], dim=1)
+        pose_next_pred = self.fc_out(fused_feat)
+        
+        return pose_next_pred

algorithms/worldmem/models/rotary_embedding_torch.py

@@ -0,0 +1,302 @@
+"""
+Adapted from https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
+"""
+
+from __future__ import annotations
+from math import pi, log
+
+import torch
+from torch.nn import Module, ModuleList
+from torch.amp import autocast
+from torch import nn, einsum, broadcast_tensors, Tensor
+
+from einops import rearrange, repeat
+
+from typing import Literal
+
+# helper functions
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+# broadcat, as tortoise-tts was using it
+
+
+def broadcat(tensors, dim=-1):
+    broadcasted_tensors = broadcast_tensors(*tensors)
+    return torch.cat(broadcasted_tensors, dim=dim)
+
+
+# rotary embedding helper functions
+
+
+def rotate_half(x):
+    x = rearrange(x, "... (d r) -> ... d r", r=2)
+    x1, x2 = x.unbind(dim=-1)
+    x = torch.stack((-x2, x1), dim=-1)
+    return rearrange(x, "... d r -> ... (d r)")
+
+
+@autocast("cuda", enabled=False)
+def apply_rotary_emb(freqs, t, start_index=0, scale=1.0, seq_dim=-2):
+    dtype = t.dtype
+
+    if t.ndim == 3:
+        seq_len = t.shape[seq_dim]
+        freqs = freqs[-seq_len:]
+
+    rot_dim = freqs.shape[-1]
+    end_index = start_index + rot_dim
+
+    assert rot_dim <= t.shape[-1], f"feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}"
+
+    # Split t into three parts: left, middle (to be transformed), and right
+    t_left = t[..., :start_index]
+    t_middle = t[..., start_index:end_index]
+    t_right = t[..., end_index:]
+
+    # Apply rotary embeddings without modifying t in place
+    t_transformed = (t_middle * freqs.cos() * scale) + (rotate_half(t_middle) * freqs.sin() * scale)
+
+    out = torch.cat((t_left, t_transformed, t_right), dim=-1)
+
+    return out.type(dtype)
+
+
+# learned rotation helpers
+
+
+def apply_learned_rotations(rotations, t, start_index=0, freq_ranges=None):
+    if exists(freq_ranges):
+        rotations = einsum("..., f -> ... f", rotations, freq_ranges)
+        rotations = rearrange(rotations, "... r f -> ... (r f)")
+
+    rotations = repeat(rotations, "... n -> ... (n r)", r=2)
+    return apply_rotary_emb(rotations, t, start_index=start_index)
+
+
+# classes
+
+
+class RotaryEmbedding(Module):
+    def __init__(
+        self,
+        dim,
+        custom_freqs: Tensor | None = None,
+        freqs_for: Literal["lang", "pixel", "constant"] = "lang",
+        theta=10000,
+        max_freq=10,
+        num_freqs=1,
+        learned_freq=False,
+        use_xpos=False,
+        xpos_scale_base=512,
+        interpolate_factor=1.0,
+        theta_rescale_factor=1.0,
+        seq_before_head_dim=False,
+        cache_if_possible=True,
+        cache_max_seq_len=8192,
+    ):
+        super().__init__()
+        # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
+        # has some connection to NTK literature
+        # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
+
+        theta *= theta_rescale_factor ** (dim / (dim - 2))
+
+        self.freqs_for = freqs_for
+
+        if exists(custom_freqs):
+            freqs = custom_freqs
+        elif freqs_for == "lang":
+            freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+        elif freqs_for == "pixel":
+            freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi
+        elif freqs_for == "spacetime":
+            time_freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+            freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi
+        elif freqs_for == "constant":
+            freqs = torch.ones(num_freqs).float()
+
+        if freqs_for == "spacetime":
+            self.time_freqs = nn.Parameter(time_freqs, requires_grad=learned_freq)
+        self.freqs = nn.Parameter(freqs, requires_grad=learned_freq)
+
+        self.cache_if_possible = cache_if_possible
+        self.cache_max_seq_len = cache_max_seq_len
+
+        self.register_buffer("cached_freqs", torch.zeros(cache_max_seq_len, dim), persistent=False)
+        self.register_buffer("cached_freqs_seq_len", torch.tensor(0), persistent=False)
+
+        self.learned_freq = learned_freq
+
+        # dummy for device
+
+        self.register_buffer("dummy", torch.tensor(0), persistent=False)
+
+        # default sequence dimension
+
+        self.seq_before_head_dim = seq_before_head_dim
+        self.default_seq_dim = -3 if seq_before_head_dim else -2
+
+        # interpolation factors
+
+        assert interpolate_factor >= 1.0
+        self.interpolate_factor = interpolate_factor
+
+        # xpos
+
+        self.use_xpos = use_xpos
+
+        if not use_xpos:
+            return
+
+        scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
+        self.scale_base = xpos_scale_base
+
+        self.register_buffer("scale", scale, persistent=False)
+        self.register_buffer("cached_scales", torch.zeros(cache_max_seq_len, dim), persistent=False)
+        self.register_buffer("cached_scales_seq_len", torch.tensor(0), persistent=False)
+
+        # add apply_rotary_emb as static method
+
+        self.apply_rotary_emb = staticmethod(apply_rotary_emb)
+
+    @property
+    def device(self):
+        return self.dummy.device
+
+    def get_seq_pos(self, seq_len, device, dtype, offset=0):
+        return (torch.arange(seq_len, device=device, dtype=dtype) + offset) / self.interpolate_factor
+
+    def rotate_queries_or_keys(self, t, freqs, seq_dim=None, offset=0, scale=None):
+        seq_dim = default(seq_dim, self.default_seq_dim)
+
+        assert not self.use_xpos or exists(scale), "you must use `.rotate_queries_and_keys` method instead and pass in both queries and keys, for length extrapolatable rotary embeddings"
+
+        device, dtype, seq_len = t.device, t.dtype, t.shape[seq_dim]
+
+        seq = self.get_seq_pos(seq_len, device=device, dtype=dtype, offset=offset)
+
+        seq_freqs = self.forward(seq, freqs, seq_len=seq_len, offset=offset)
+
+        if seq_dim == -3:
+            seq_freqs = rearrange(seq_freqs, "n d -> n 1 d")
+
+        return apply_rotary_emb(seq_freqs, t, scale=default(scale, 1.0), seq_dim=seq_dim)
+
+    def rotate_queries_with_cached_keys(self, q, k, seq_dim=None, offset=0):
+        dtype, device, seq_dim = (
+            q.dtype,
+            q.device,
+            default(seq_dim, self.default_seq_dim),
+        )
+
+        q_len, k_len = q.shape[seq_dim], k.shape[seq_dim]
+        assert q_len <= k_len
+
+        q_scale = k_scale = 1.0
+
+        if self.use_xpos:
+            seq = self.get_seq_pos(k_len, dtype=dtype, device=device)
+
+            q_scale = self.get_scale(seq[-q_len:]).type(dtype)
+            k_scale = self.get_scale(seq).type(dtype)
+
+        rotated_q = self.rotate_queries_or_keys(q, seq_dim=seq_dim, scale=q_scale, offset=k_len - q_len + offset)
+        rotated_k = self.rotate_queries_or_keys(k, seq_dim=seq_dim, scale=k_scale**-1)
+
+        rotated_q = rotated_q.type(q.dtype)
+        rotated_k = rotated_k.type(k.dtype)
+
+        return rotated_q, rotated_k
+
+    def rotate_queries_and_keys(self, q, k, freqs, seq_dim=None):
+        seq_dim = default(seq_dim, self.default_seq_dim)
+
+        assert self.use_xpos
+        device, dtype, seq_len = q.device, q.dtype, q.shape[seq_dim]
+
+        seq = self.get_seq_pos(seq_len, dtype=dtype, device=device)
+
+        seq_freqs = self.forward(seq, freqs, seq_len=seq_len)
+        scale = self.get_scale(seq, seq_len=seq_len).to(dtype)
+
+        if seq_dim == -3:
+            seq_freqs = rearrange(seq_freqs, "n d -> n 1 d")
+            scale = rearrange(scale, "n d -> n 1 d")
+
+        rotated_q = apply_rotary_emb(seq_freqs, q, scale=scale, seq_dim=seq_dim)
+        rotated_k = apply_rotary_emb(seq_freqs, k, scale=scale**-1, seq_dim=seq_dim)
+
+        rotated_q = rotated_q.type(q.dtype)
+        rotated_k = rotated_k.type(k.dtype)
+
+        return rotated_q, rotated_k
+
+    def get_scale(self, t: Tensor, seq_len: int | None = None, offset=0):
+        assert self.use_xpos
+
+        should_cache = self.cache_if_possible and exists(seq_len) and (offset + seq_len) <= self.cache_max_seq_len
+
+        if should_cache and exists(self.cached_scales) and (seq_len + offset) <= self.cached_scales_seq_len.item():
+            return self.cached_scales[offset : (offset + seq_len)]
+
+        scale = 1.0
+        if self.use_xpos:
+            power = (t - len(t) // 2) / self.scale_base
+            scale = self.scale ** rearrange(power, "n -> n 1")
+            scale = repeat(scale, "n d -> n (d r)", r=2)
+
+        if should_cache and offset == 0:
+            self.cached_scales[:seq_len] = scale.detach()
+            self.cached_scales_seq_len.copy_(seq_len)
+
+        return scale
+
+    def get_axial_freqs(self, *dims):
+        Colon = slice(None)
+        all_freqs = []
+
+        for ind, dim in enumerate(dims):
+            # only allow pixel freqs for last two dimensions
+            use_pixel = (self.freqs_for == "pixel" or self.freqs_for == "spacetime") and ind >= len(dims) - 2
+            if use_pixel:
+                pos = torch.linspace(-1, 1, steps=dim, device=self.device)
+            else:
+                pos = torch.arange(dim, device=self.device)
+
+            if self.freqs_for == "spacetime" and not use_pixel:
+                seq_freqs = self.forward(pos, self.time_freqs, seq_len=dim)
+            else:
+                seq_freqs = self.forward(pos, self.freqs, seq_len=dim)
+
+            all_axis = [None] * len(dims)
+            all_axis[ind] = Colon
+
+            new_axis_slice = (Ellipsis, *all_axis, Colon)
+            all_freqs.append(seq_freqs[new_axis_slice])
+
+        all_freqs = broadcast_tensors(*all_freqs)
+        return torch.cat(all_freqs, dim=-1)
+
+    @autocast("cuda", enabled=False)
+    def forward(self, t: Tensor, freqs: Tensor, seq_len=None, offset=0):
+        should_cache = self.cache_if_possible and not self.learned_freq and exists(seq_len) and self.freqs_for != "pixel" and (offset + seq_len) <= self.cache_max_seq_len
+
+        if should_cache and exists(self.cached_freqs) and (offset + seq_len) <= self.cached_freqs_seq_len.item():
+            return self.cached_freqs[offset : (offset + seq_len)].detach()
+
+        freqs = einsum("..., f -> ... f", t.type(freqs.dtype), freqs)
+        freqs = repeat(freqs, "... n -> ... (n r)", r=2)
+
+        if should_cache and offset == 0:
+            self.cached_freqs[:seq_len] = freqs.detach()
+            self.cached_freqs_seq_len.copy_(seq_len)
+
+        return freqs

algorithms/worldmem/models/utils.py

@@ -0,0 +1,163 @@
+"""
+Adapted from https://github.com/buoyancy99/diffusion-forcing/blob/main/algorithms/diffusion_forcing/models/utils.py
+Action format derived from VPT https://github.com/openai/Video-Pre-Training
+Adapted from https://github.com/etched-ai/open-oasis/blob/master/utils.py
+"""
+
+import math
+import torch
+from torch import nn
+from torchvision.io import read_image, read_video
+from torchvision.transforms.functional import resize
+from einops import rearrange
+from typing import Mapping, Sequence
+from einops import rearrange, parse_shape
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if callable(d) else d
+
+
+def extract(a, t, x_shape):
+    f, b = t.shape
+    out = a[t]
+    return out.reshape(f, b, *((1,) * (len(x_shape) - 2)))
+
+
+def linear_beta_schedule(timesteps):
+    """
+    linear schedule, proposed in original ddpm paper
+    """
+    scale = 1000 / timesteps
+    beta_start = scale * 0.0001
+    beta_end = scale * 0.02
+    return torch.linspace(beta_start, beta_end, timesteps, dtype=torch.float64)
+
+
+def cosine_beta_schedule(timesteps, s=0.008):
+    """
+    cosine schedule
+    as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
+    """
+    steps = timesteps + 1
+    t = torch.linspace(0, timesteps, steps, dtype=torch.float64) / timesteps
+    alphas_cumprod = torch.cos((t + s) / (1 + s) * math.pi * 0.5) ** 2
+    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+    return torch.clip(betas, 0, 0.999)
+
+
+
+def sigmoid_beta_schedule(timesteps, start=-3, end=3, tau=1, clamp_min=1e-5):
+    """
+    sigmoid schedule
+    proposed in https://arxiv.org/abs/2212.11972 - Figure 8
+    better for images > 64x64, when used during training
+    """
+    steps = timesteps + 1
+    t = torch.linspace(0, timesteps, steps, dtype=torch.float64) / timesteps
+    v_start = torch.tensor(start / tau).sigmoid()
+    v_end = torch.tensor(end / tau).sigmoid()
+    alphas_cumprod = (-((t * (end - start) + start) / tau).sigmoid() + v_end) / (v_end - v_start)
+    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+    return torch.clip(betas, 0, 0.999)
+
+
+ACTION_KEYS = [
+    "inventory",
+    "ESC",
+    "hotbar.1",
+    "hotbar.2",
+    "hotbar.3",
+    "hotbar.4",
+    "hotbar.5",
+    "hotbar.6",
+    "hotbar.7",
+    "hotbar.8",
+    "hotbar.9",
+    "forward",
+    "back",
+    "left",
+    "right",
+    "cameraX",
+    "cameraY",
+    "jump",
+    "sneak",
+    "sprint",
+    "swapHands",
+    "attack",
+    "use",
+    "pickItem",
+    "drop",
+]
+
+
+def one_hot_actions(actions: Sequence[Mapping[str, int]]) -> torch.Tensor:
+    actions_one_hot = torch.zeros(len(actions), len(ACTION_KEYS))
+    for i, current_actions in enumerate(actions):
+        for j, action_key in enumerate(ACTION_KEYS):
+            if action_key.startswith("camera"):
+                if action_key == "cameraX":
+                    value = current_actions["camera"][0]
+                elif action_key == "cameraY":
+                    value = current_actions["camera"][1]
+                else:
+                    raise ValueError(f"Unknown camera action key: {action_key}")
+                max_val = 20
+                bin_size = 0.5
+                num_buckets = int(max_val / bin_size)
+                value = (value - num_buckets) / num_buckets
+                assert -1 - 1e-3 <= value <= 1 + 1e-3, f"Camera action value must be in [-1, 1], got {value}"
+            else:
+                value = current_actions[action_key]
+                assert 0 <= value <= 1, f"Action value must be in [0, 1] got {value}"
+            actions_one_hot[i, j] = value
+
+    return actions_one_hot
+
+
+IMAGE_EXTENSIONS = {"png", "jpg", "jpeg"}
+VIDEO_EXTENSIONS = {"mp4"}
+
+
+def load_prompt(path, video_offset=None, n_prompt_frames=1):
+    if path.lower().split(".")[-1] in IMAGE_EXTENSIONS:
+        print("prompt is image; ignoring video_offset and n_prompt_frames")
+        prompt = read_image(path)
+        # add frame dimension
+        prompt = rearrange(prompt, "c h w -> 1 c h w")
+    elif path.lower().split(".")[-1] in VIDEO_EXTENSIONS:
+        prompt = read_video(path, pts_unit="sec")[0]
+        if video_offset is not None:
+            prompt = prompt[video_offset:]
+        prompt = prompt[:n_prompt_frames]
+    else:
+        raise ValueError(f"unrecognized prompt file extension; expected one in {IMAGE_EXTENSIONS} or {VIDEO_EXTENSIONS}")
+    assert prompt.shape[0] == n_prompt_frames, f"input prompt {path} had less than n_prompt_frames={n_prompt_frames} frames"
+    prompt = resize(prompt, (360, 640))
+    # add batch dimension
+    prompt = rearrange(prompt, "t c h w -> 1 t c h w")
+    prompt = prompt.float() / 255.0
+    return prompt
+
+
+def load_actions(path, action_offset=None):
+    if path.endswith(".actions.pt"):
+        actions = one_hot_actions(torch.load(path))
+    elif path.endswith(".one_hot_actions.pt"):
+        actions = torch.load(path, weights_only=True)
+    else:
+        raise ValueError("unrecognized action file extension; expected '*.actions.pt' or '*.one_hot_actions.pt'")
+    if action_offset is not None:
+        actions = actions[action_offset:]
+    actions = torch.cat([torch.zeros_like(actions[:1]), actions], dim=0)
+    # add batch dimension
+    actions = rearrange(actions, "t d -> 1 t d")
+    return actions

algorithms/worldmem/models/vae.py

@@ -0,0 +1,359 @@
+"""
+References:
+    - VQGAN: https://github.com/CompVis/taming-transformers
+    - MAE: https://github.com/facebookresearch/mae
+"""
+
+import numpy as np
+import math
+import functools
+from collections import namedtuple
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from timm.models.vision_transformer import Mlp
+from timm.layers.helpers import to_2tuple
+from rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb
+from .dit import PatchEmbed
+
+
+class DiagonalGaussianDistribution(object):
+    def __init__(self, parameters, deterministic=False, dim=1):
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=dim)
+        if dim == 1:
+            self.dims = [1, 2, 3]
+        elif dim == 2:
+            self.dims = [1, 2]
+        else:
+            raise NotImplementedError
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
+
+    def sample(self):
+        x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device)
+        return x
+
+    def mode(self):
+        return self.mean
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        frame_height,
+        frame_width,
+        qkv_bias=False,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.frame_height = frame_height
+        self.frame_width = frame_width
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.proj = nn.Linear(dim, dim)
+
+        rotary_freqs = RotaryEmbedding(
+            dim=head_dim // 4,
+            freqs_for="pixel",
+            max_freq=frame_height * frame_width,
+        ).get_axial_freqs(frame_height, frame_width)
+        self.register_buffer("rotary_freqs", rotary_freqs, persistent=False)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        assert N == self.frame_height * self.frame_width
+
+        q, k, v = self.qkv(x).chunk(3, dim=-1)
+
+        q = rearrange(
+            q,
+            "b (H W) (h d) -> b h H W d",
+            H=self.frame_height,
+            W=self.frame_width,
+            h=self.num_heads,
+        )
+        k = rearrange(
+            k,
+            "b (H W) (h d) -> b h H W d",
+            H=self.frame_height,
+            W=self.frame_width,
+            h=self.num_heads,
+        )
+        v = rearrange(
+            v,
+            "b (H W) (h d) -> b h H W d",
+            H=self.frame_height,
+            W=self.frame_width,
+            h=self.num_heads,
+        )
+
+        q = apply_rotary_emb(self.rotary_freqs, q)
+        k = apply_rotary_emb(self.rotary_freqs, k)
+
+        q = rearrange(q, "b h H W d -> b h (H W) d")
+        k = rearrange(k, "b h H W d -> b h (H W) d")
+        v = rearrange(v, "b h H W d -> b h (H W) d")
+
+        x = F.scaled_dot_product_attention(q, k, v)
+        x = rearrange(x, "b h N d -> b N (h d)")
+
+        x = self.proj(x)
+        return x
+
+
+class AttentionBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        frame_height,
+        frame_width,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        attn_causal=False,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+    ):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads,
+            frame_height,
+            frame_width,
+            qkv_bias=qkv_bias,
+        )
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+        )
+
+    def forward(self, x):
+        x = x + self.attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class AutoencoderKL(nn.Module):
+    def __init__(
+        self,
+        latent_dim,
+        input_height=256,
+        input_width=256,
+        patch_size=16,
+        enc_dim=768,
+        enc_depth=6,
+        enc_heads=12,
+        dec_dim=768,
+        dec_depth=6,
+        dec_heads=12,
+        mlp_ratio=4.0,
+        norm_layer=functools.partial(nn.LayerNorm, eps=1e-6),
+        use_variational=True,
+        **kwargs,
+    ):
+        super().__init__()
+        self.input_height = input_height
+        self.input_width = input_width
+        self.patch_size = patch_size
+        self.seq_h = input_height // patch_size
+        self.seq_w = input_width // patch_size
+        self.seq_len = self.seq_h * self.seq_w
+        self.patch_dim = 3 * patch_size**2
+
+        self.latent_dim = latent_dim
+        self.enc_dim = enc_dim
+        self.dec_dim = dec_dim
+
+        # patch
+        self.patch_embed = PatchEmbed(input_height, input_width, patch_size, 3, enc_dim)
+
+        # encoder
+        self.encoder = nn.ModuleList(
+            [
+                AttentionBlock(
+                    enc_dim,
+                    enc_heads,
+                    self.seq_h,
+                    self.seq_w,
+                    mlp_ratio,
+                    qkv_bias=True,
+                    norm_layer=norm_layer,
+                )
+                for i in range(enc_depth)
+            ]
+        )
+        self.enc_norm = norm_layer(enc_dim)
+
+        # bottleneck
+        self.use_variational = use_variational
+        mult = 2 if self.use_variational else 1
+        self.quant_conv = nn.Linear(enc_dim, mult * latent_dim)
+        self.post_quant_conv = nn.Linear(latent_dim, dec_dim)
+
+        # decoder
+        self.decoder = nn.ModuleList(
+            [
+                AttentionBlock(
+                    dec_dim,
+                    dec_heads,
+                    self.seq_h,
+                    self.seq_w,
+                    mlp_ratio,
+                    qkv_bias=True,
+                    norm_layer=norm_layer,
+                )
+                for i in range(dec_depth)
+            ]
+        )
+        self.dec_norm = norm_layer(dec_dim)
+        self.predictor = nn.Linear(dec_dim, self.patch_dim)  # decoder to patch
+
+        # initialize this weight first
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # initialization
+        # initialize nn.Linear and nn.LayerNorm
+        self.apply(self._init_weights)
+
+        # initialize patch_embed like nn.Linear (instead of nn.Conv2d)
+        w = self.patch_embed.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            # we use xavier_uniform following official JAX ViT:
+            nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0.0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0.0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def patchify(self, x):
+        # patchify
+        bsz, _, h, w = x.shape
+        x = x.reshape(
+            bsz,
+            3,
+            self.seq_h,
+            self.patch_size,
+            self.seq_w,
+            self.patch_size,
+        ).permute([0, 1, 3, 5, 2, 4])  # [b, c, h, p, w, p] --> [b, c, p, p, h, w]
+        x = x.reshape(bsz, self.patch_dim, self.seq_h, self.seq_w)  # --> [b, cxpxp, h, w]
+        x = x.permute([0, 2, 3, 1]).reshape(bsz, self.seq_len, self.patch_dim)  # --> [b, hxw, cxpxp]
+        return x
+
+    def unpatchify(self, x):
+        bsz = x.shape[0]
+        # unpatchify
+        x = x.reshape(bsz, self.seq_h, self.seq_w, self.patch_dim).permute([0, 3, 1, 2])  # [b, h, w, cxpxp] --> [b, cxpxp, h, w]
+        x = x.reshape(
+            bsz,
+            3,
+            self.patch_size,
+            self.patch_size,
+            self.seq_h,
+            self.seq_w,
+        ).permute([0, 1, 4, 2, 5, 3])  # [b, c, p, p, h, w] --> [b, c, h, p, w, p]
+        x = x.reshape(
+            bsz,
+            3,
+            self.input_height,
+            self.input_width,
+        )  # [b, c, hxp, wxp]
+        return x
+
+    def encode(self, x):
+        # patchify
+        x = self.patch_embed(x)
+
+        # encoder
+        for blk in self.encoder:
+            x = blk(x)
+        x = self.enc_norm(x)
+
+        # bottleneck
+        moments = self.quant_conv(x)
+        if not self.use_variational:
+            moments = torch.cat((moments, torch.zeros_like(moments)), 2)
+        posterior = DiagonalGaussianDistribution(moments, deterministic=(not self.use_variational), dim=2)
+        return posterior
+
+    def decode(self, z):
+        # bottleneck
+        z = self.post_quant_conv(z)
+
+        # decoder
+        for blk in self.decoder:
+            z = blk(z)
+        z = self.dec_norm(z)
+
+        # predictor
+        z = self.predictor(z)
+
+        # unpatchify
+        dec = self.unpatchify(z)
+        return dec
+
+    def autoencode(self, input, sample_posterior=True):
+        posterior = self.encode(input)
+        if self.use_variational and sample_posterior:
+            z = posterior.sample()
+        else:
+            z = posterior.mode()
+        dec = self.decode(z)
+        return dec, posterior, z
+
+    def get_input(self, batch, k):
+        x = batch[k]
+        if len(x.shape) == 3:
+            x = x[..., None]
+        x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
+        return x
+
+    def forward(self, inputs, labels, split="train"):
+        rec, post, latent = self.autoencode(inputs)
+        return rec, post, latent
+
+    def get_last_layer(self):
+        return self.predictor.weight
+
+
+def ViT_L_20_Shallow_Encoder(**kwargs):
+    if "latent_dim" in kwargs:
+        latent_dim = kwargs.pop("latent_dim")
+    else:
+        latent_dim = 16
+    return AutoencoderKL(
+        latent_dim=latent_dim,
+        patch_size=20,
+        enc_dim=1024,
+        enc_depth=6,
+        enc_heads=16,
+        dec_dim=1024,
+        dec_depth=12,
+        dec_heads=16,
+        input_height=360,
+        input_width=640,
+        **kwargs,
+    )
+
+
+VAE_models = {
+    "vit-l-20-shallow-encoder": ViT_L_20_Shallow_Encoder,
+}

configurations/algorithm/base_algo.yaml

@@ -0,0 +1,3 @@
+# This will be passed as the cfg to Algo.__init__(cfg) of your algorithm class
+
+debug: ${debug} # inherited from configurations/config.yaml

configurations/algorithm/base_pytorch_algo.yaml

@@ -0,0 +1,4 @@
+defaults:
+  - base_algo # inherits from configurations/algorithm/base_algo.yaml
+
+lr: ${experiment.training.lr}

configurations/algorithm/base_video_dit.yaml

@@ -0,0 +1,36 @@
+defaults:
+  - df_base
+
+n_frames: ${dataset.n_frames}
+frame_skip: ${dataset.frame_skip}
+metadata: ${dataset.metadata}
+
+# training hyperparameters
+weight_decay: 2e-3
+warmup_steps: 1000
+optimizer_beta: [0.9, 0.99]
+action_cond_dim: 25
+use_plucker: true
+
+diffusion:
+  # training
+  beta_schedule: sigmoid
+  objective: pred_v
+  use_fused_snr: True
+  cum_snr_decay: 0.96
+  clip_noise: 20.
+  # sampling
+  sampling_timesteps: 20
+  ddim_sampling_eta: 0.0
+  stabilization_level: 15
+  # architecture
+  architecture:
+    network_size: 64
+    attn_heads: 4
+    attn_dim_head: 64
+    dim_mults: [1, 2, 4, 8]
+    resolution: ${dataset.resolution}
+    attn_resolutions: [16, 32, 64, 128]
+    use_init_temporal_attn: True
+    use_linear_attn: True
+    time_emb_type: rotary

configurations/algorithm/dememwm_memory_dit.yaml

@@ -0,0 +1,103 @@
+
+defaults:
+  - base_video_dit
+  - _self_
+
+_name: dememwm_memory_dit
+
+# Standalone Memory-DiT path. Do not route through old SSM-memory config.
+memory_token_cross_attention: true
+memory_cross_attn_layers: null
+memory_condition_length: 0
+pose_cond_dim: 5
+log_video: false
+
+dememwm:
+  enabled: true
+  training_stage: stage_1  # fallback only when curriculum.enabled=false
+  debug_force_all_streams: false
+  curriculum:
+    enabled: true
+    full_stage_start_step: 60000
+    freeze_vae: true
+    dit_freeze:
+      enabled: true
+    lr:
+      dememwm_modules: 1.0e-4
+      memory_adapters: 1.0e-4
+      full_dit: 1.0e-5
+  # Current Conv2D memory projectors preserve latent H,W=(18,32).
+  # Pool sizes are resolved from projected spatial grid size and downsample ratios.
+  token_patch_size: 2
+  anchor:
+    enabled: true
+    anchor_indices: [0, 1, 2, 3]
+    allow_generated_as_anchor: false
+    diverse_selection: true
+    compress:
+      downsample_ratio: 4
+  dynamic:
+    enabled: true
+    exclude_latest_local_frames: 4
+    recent_frames: 8
+    conv_kernel_t: 3
+    conv_stride_t: 2
+  revisit:
+    enabled: true
+    deterministic_pose_retrieval: true
+    fov_overlap_threshold: 0.30
+    high_quality_fov_threshold: 0.70
+    plucker_weight: 0.10
+    max_frames: 2
+    # FoV geometry for coverage-based retrieval scoring.
+    # fov_half_h/v: half-angles (degrees) of the horizontal/vertical field of view.
+    # fov_radius: world-space radius of the sample sphere.
+    # fov_{yaw,pitch,depth}_samples: grid resolution for FoV point sampling.
+    fov_half_h: 52.5   # 105 deg total horizontal FoV
+    fov_half_v: 37.5   # 75 deg total vertical FoV
+    fov_radius: 30.0
+    fov_yaw_samples: 25
+    fov_pitch_samples: 20
+    fov_depth_samples: 20
+    pose_preselect_topk: 64
+    # Plucker descriptor grid for secondary pose-similarity scoring.
+    plucker_grid_h: 4
+    plucker_grid_w: 4
+    plucker_focal_length: 0.35
+    compress:
+      downsample_ratio: 4
+  stage_policy:
+    noise_bucket_logging: true
+  eval_ablation:
+    enabled: false
+    branch: A_plus_D_plus_R_normal
+  generated_history_proxy:
+    enabled: false
+    start_step: 0
+    ramp_steps: 1
+    max_prob: 0.0
+    noise_std: 0.25
+    dropout_prob: 0.0
+  injection:
+    dit_hidden_size: 1024
+    anchor_gate: 1.0
+    dynamic_gate: 1.0
+    revisit_gate: 1.0
+  cache:
+    enabled: true
+    device: cpu
+    keep_raw_latents: all
+    keep_compressed_records: true
+    keep_prefix_anchors: true
+    eviction_policy: none
+    no_evict: true
+    clear_between_videos: true
+    max_records: null
+    max_slots: null
+    on_capacity_exceeded: warn
+  checkpoint:
+    strict_dememwm_eval_load: true
+
+diffusion:
+  architecture:
+    network_size: 64

configurations/algorithm/df_base.yaml

@@ -0,0 +1,42 @@
+defaults:
+  - base_pytorch_algo
+
+# dataset-dependent configurations
+x_shape: ${dataset.observation_shape}
+frame_stack: 1
+frame_skip: 1
+data_mean: ${dataset.data_mean}
+data_std: ${dataset.data_std}
+external_cond_dim: 0 #${dataset.action_dim}
+context_frames: ${dataset.context_length}
+# training hyperparameters
+weight_decay: 1e-4
+warmup_steps: 10000
+optimizer_beta: [0.9, 0.999]
+# diffusion-related
+uncertainty_scale: 1
+guidance_scale: 0.0
+chunk_size: 1 # -1 for full trajectory diffusion, number to specify diffusion chunk size
+scheduling_matrix: autoregressive
+noise_level: random_all
+causal: True
+
+diffusion:
+  # training
+  objective: pred_x0
+  beta_schedule: cosine
+  schedule_fn_kwargs: {}
+  clip_noise: 20.0
+  use_snr: False
+  use_cum_snr: False
+  use_fused_snr: False
+  snr_clip: 5.0
+  cum_snr_decay: 0.98
+  timesteps: 1000
+  # sampling
+  sampling_timesteps: 50 # fixme, numer of diffusion steps, should be increased
+  ddim_sampling_eta: 1.0
+  stabilization_level: 10
+  # architecture
+  architecture:
+    network_size: 64

configurations/dataset/base_dataset.yaml

@@ -0,0 +1,3 @@
+# This will be passed as the cfg to Dataset.__init__(cfg) of your dataset class
+
+debug: ${debug} # inherited from configurations/config.yaml

configurations/dataset/base_video.yaml

@@ -0,0 +1,14 @@
+defaults:
+  - base_dataset
+
+metadata: "data/${dataset.name}/metadata.json"
+data_mean: "data/${dataset.name}/data_mean.npy"
+data_std: "data/${dataset.name}/data_std.npy"
+save_dir: ???
+n_frames: 32
+context_length: 4
+resolution: 128
+observation_shape: [3, "${dataset.resolution}", "${dataset.resolution}"]
+external_cond_dim: 0
+validation_multiplier: 1
+frame_skip: 1

configurations/dataset/video_minecraft.yaml

@@ -0,0 +1,14 @@
+defaults:
+  - base_video
+
+save_dir: data/minecraft_simple_backforward
+n_frames: 16 # TODO: increase later
+resolution: 128
+data_mean: 0.5
+data_std: 0.5
+action_cond_dim: 25
+context_length: 1
+frame_skip: 1
+validation_multiplier: 1
+
+_name: video_minecraft_oasis

configurations/dataset/video_minecraft_latent.yaml

@@ -0,0 +1,6 @@
+defaults:
+  - video_minecraft
+
+precomputed_feature_dir: /share_1/users/bonan_ding/worldmem_data/minecraft/vae_features
+
+_name: video_minecraft_latent

configurations/experiment/base_experiment.yaml

@@ -0,0 +1,2 @@
+debug: ${debug} # inherited from configurations/config.yaml
+tasks: [main] # tasks to run sequantially, such as [training, test], useful when your project has multiple stages and you want to run only a subset of them.

configurations/experiment/base_pytorch.yaml

@@ -0,0 +1,51 @@
+# inherites from base_experiment.yaml
+# most of the options have docs at https://lightning.ai/docs/pytorch/stable/common/trainer.html
+
+defaults:
+  - base_experiment
+
+tasks: [training] # tasks to run sequantially, change when your project has multiple stages and you want to run only a subset of them.
+num_nodes: 1 # number of gpu servers used in large scale distributed training
+
+training:
+  precision: 16-mixed # set float precision, 16-mixed is faster while 32 is more stable
+  compile: False # whether to compile the model with torch.compile
+  lr: 0.001 # learning rate
+  batch_size: 16 # training batch size; effective batch size is this number * gpu * nodes iff using distributed training
+  max_epochs: 1000 # set to -1 to train forever
+  max_steps: -1 # set to -1 to train forever, will override max_epochs
+  max_time: null # set to something like "00:12:00:00" to enable
+  data:
+    num_workers: 4 # number of CPU threads for data preprocessing.
+    shuffle: True # whether training data will be shuffled
+  optim:
+    accumulate_grad_batches: 1 # accumulate gradients for n batches before backprop
+    gradient_clip_val: 1.0 # clip gradients with norm above this value, set to 0 to disable
+  checkpointing:
+    # these are arguments to pytorch lightning's callback, `ModelCheckpoint` class
+    every_n_train_steps: 5000 # save a checkpoint every n train steps
+    every_n_epochs: null # mutually exclusive with ``every_n_train_steps`` and ``train_time_interval``
+    train_time_interval: null # in format of "00:12:00:00", mutually exclusive with ``every_n_train_steps`` and ``every_n_epochs``.
+    save_last: True # keep last.ckpt for automatic interrupted-run resume
+    enable_version_counter: False # If this is ``False``, later checkpoint will be overwrite previous ones.
+
+validation:
+  precision: 16-mixed
+  compile: False # whether to compile the model with torch.compile
+  batch_size: 16 # validation batch size per GPU; effective batch size is this number * gpu * nodes iff using distributed training
+  val_every_n_step: 2000 # controls how frequent do we run validation, can be float (fraction of epoches) or int (steps) or null (if val_every_n_epoch is set)
+  val_every_n_epoch: null # if you want to do validation every n epoches, requires val_every_n_step to be null.
+  limit_batch: null # if null, run through validation set. Otherwise limit the number of batches to use for validation.
+  inference_mode: True # whether to run validation in inference mode (enable_grad won't work!)
+  data:
+    num_workers: 4 # number of CPU threads for data preprocessing, for validation.
+    shuffle: False # whether validation data will be shuffled
+
+test:
+  precision: 16-mixed
+  compile: False # whether to compile the model with torch.compile
+  batch_size: 4 # test batch size per GPU; effective batch size is this number * gpu * nodes iff using distributed training
+  limit_batch: null # if null, run through test set. Otherwise limit the number of batches to use for test.
+  data:
+    num_workers: 4 # number of CPU threads for data preprocessing, for test.
+    shuffle: False # whether test data will be shuffled

configurations/experiment/exp_video.yaml

@@ -0,0 +1,31 @@
+defaults:
+  - base_pytorch
+
+tasks: [training]
+
+training:
+  lr: 2e-5
+  precision: 16-mixed
+  batch_size: 4
+  max_epochs: -1
+  max_steps: 2000005
+  checkpointing:
+    every_n_train_steps: 2500
+  optim:
+    gradient_clip_val: 1.0
+
+validation:
+  val_every_n_step: 2500
+  val_every_n_epoch: null
+  batch_size: 4
+  limit_batch: 1
+
+test:
+  limit_batch: 1
+  batch_size: 1
+
+logging:
+  metrics:
+    # - fvd
+    # - fid
+    # - lpips

configurations/training.yaml

@@ -0,0 +1,18 @@
+# configuration parsing starts here
+defaults:
+  - experiment: exp_video # experiment yaml file name in configurations/experiments folder [fixme]
+  - dataset: video_minecraft # dataset yaml file name in configurations/dataset folder [fixme]
+  - algorithm: dememwm_memory_dit # algorithm yaml file name in configurations/algorithm folder [fixme]
+  - cluster: null # optional, cluster yaml file name in configurations/cluster folder. Leave null for local compute
+
+debug: false # global debug flag will be passed into configuration of experiment, dataset and algorithm
+
+wandb:
+  entity: turlin # wandb account name / organization name [fixme]
+  project: DeMemWM # wandb project name; if not provided, defaults to root folder name [fixme]
+  mode: offline # set wandb logging to online, offline or dryrun
+
+resume: null # wandb run id to resume logging and loading checkpoint from
+load: null # wandb run id containing checkpoint or a path to a checkpoint file
+auto_resume: true # automatically resume training from output_dir/checkpoints when available
+resume_ckpt_path: null # explicit full Lightning checkpoint path for deterministic training resume

datasets/__init__.py

@@ -0,0 +1 @@
+from .video import MinecraftVideoDataset, MinecraftVideoLatentDataset

datasets/video/__init__.py

@@ -0,0 +1,2 @@
+from .minecraft_video_dataset import MinecraftVideoDataset
+from .minecraft_video_latent_dataset import MinecraftVideoLatentDataset