Add DeMemWM step 1-4 dataset scaffold

.exp_artifact/dememwm_rebuild_plan.md

@@ -0,0 +1,290 @@
+# DeMemWM Rebuild Plan
+
+## Goal
+
+Rebuild the DeMemWM idea on the `dememwm` branch under `WorldMem` while keeping the existing WorldMem implementation clean and usable as the baseline.
+
+The rebuilt design should move memory selection into the dataset, like the current WorldMem dataset path, and keep the DeMemWM algorithm focused on packing, noise handling, model calls, and target-only loss.
+
+## Current Status
+
+- Step 1: copied DeMemWM namespace exists, but `algorithms/worldmem` must remain read-only from this point on. The current working tree has no `algorithms/worldmem` diff.
+- Step 2: `dememwm_base` config and experiment registration exist.
+- Step 3: DeMemWM latent dataset exists and returns `[target][anchor][dynamic][revisit]` with `latents`, `actions`, `poses`, `frame_indices`, `memory_segments`, `memory_masks`, and `image_hw`.
+- Step 4: typed DeMemWM dataset selector exists. `worldmem` backend selection is not implemented yet; do not silently rely on `memory_selection.backend: worldmem` until that backend is added or explicitly rejected.
+- Step 5: not started. The current copied algorithm is not yet compatible with the new latent dataset contract; the known data-to-method integration issues are listed under Step 5 below.
+- Steps 6-10: not started.
+
+## 1. Keep `worldmem` Clean
+
+Do not modify `algorithms/worldmem` for DeMemWM-specific behavior. It remains the baseline/reference implementation.
+
+Create an independent algorithm namespace:
+
+```text
+algorithms/
+  worldmem/
+    ...
+
+  dememwm/
+    __init__.py
+    df_base.py
+    df_video.py
+    pose_prediction.py
+    models/
+      attention.py
+      diffusion.py
+      dit.py
+      vae.py
+      ...
+```
+
+Initial implementation step: copy `algorithms/worldmem` to `algorithms/dememwm`, then rebuild DeMemWM inside that copy.
+
+## 2. New Base Config
+
+Create:
+
+```text
+configurations/algorithm/dememwm_base.yaml
+```
+
+Do not start with `dememwm_frame_memory_ablation.yaml`. Ablation configs can be added later as separate overrides or dedicated YAML files.
+
+Register the new algorithm in `experiments/exp_video.py`:
+
+```python
+dememwm_base = DeMemWMMinecraft
+```
+
+The first explicit target run should use:
+
+```bash
+algorithm=dememwm_base
+dataset=video_minecraft_dememwm_latent
+```
+
+## 3. Dataset-Side Memory Selection
+
+Memory selection happens inside the dataset, not inside the algorithm.
+
+Create a DeMemWM latent dataset, for example:
+
+```text
+datasets/video/minecraft_video_dememwm_latent_dataset.py
+```
+
+It should load precomputed VAE latents and append memory frames directly:
+
+```text
+[target][anchor][dynamic][revisit]
+```
+
+Returned batch should include:
+
+```python
+latents
+actions
+poses
+frame_indices
+memory_segments
+memory_masks
+image_hw
+```
+
+Example segment metadata:
+
+```python
+memory_segments = {
+    "target": T,
+    "anchor": A,
+    "dynamic": D,
+    "revisit": R,
+}
+```
+
+The dataset should also return masks for padded or unavailable memory slots.
+
+## 4. Optional Dataset Selection Backends
+
+Add configurable memory selection:
+
+```yaml
+memory_selection:
+  enabled: true
+  backend: dememwm   # worldmem | dememwm
+  max_anchor_frames: 2
+  max_dynamic_frames: 4
+  max_revisit_frames: 2
+```
+
+### `worldmem` Backend
+
+Use FOV-style memory selection similar to current WorldMem. This backend is useful as a baseline and should append selected memory frames after the target sequence.
+
+### `dememwm` Backend
+
+Use typed memory streams:
+
+- `anchor`: prefix/context frames, optionally pose-diverse.
+- `dynamic`: recent frames before the target window.
+- `revisit`: older causal frames selected by FOV overlap plus Plucker score plus thresholds.
+
+Representative config:
+
+```yaml
+memory_selection:
+  fov_overlap_threshold: 0.6
+  min_total_selected_coverage: 0.1
+  local_context_exclusion_frames: 8
+  plucker_weight: 0.1
+  anchor_diverse_selection: true
+```
+
+Core invariants:
+
+- no future memory;
+- dynamic frames come before the target window;
+- revisit excludes local context;
+- anchor prefers prefix/context frames;
+- padded memory slots are masked.
+
+## 5. Thin DeMemWM Algorithm
+
+`algorithms/dememwm/df_video.py` should not select memory.
+
+Current status: not started. The copied WorldMem algorithm still expects the old tuple/raw-frame path, so the `algorithm=dememwm_base` and `dataset=video_minecraft_dememwm_latent` pairing is not yet runnable as intended.
+
+Its responsibilities:
+
+1. Preprocess the batch.
+2. Read `memory_segments` and `memory_masks`.
+3. Split target and memory by segment metadata.
+4. Zero action conditions for memory frames.
+5. Assign target noise and cleaner or clean memory noise.
+6. Call DeMemWM diffusion/DiT with memory metadata.
+7. Compute loss only on target frames.
+
+This keeps the division of responsibility simple:
+
+- dataset: chooses and appends memory;
+- algorithm: prepares noise/conditions and computes loss;
+- model: performs memory-aware denoising.
+
+Known data-to-method integration issues to fix in this step:
+
+- Update `_preprocess_batch` to consume the latent dataset dict instead of unpacking the old 4-tuple.
+- Treat `latents` as already VAE-encoded; do not call the VAE encoder on latent tensors.
+- Split returned sequences by `memory_segments` so appended memory frames are not predicted or scored as target frames.
+- Use `memory_masks` so padded memory slots do not become valid conditioning references.
+- Zero action conditions only for valid memory frames.
+- Use returned `image_hw` for Plucker/ray geometry instead of latent tensor height/width.
+- Remove validation/test algorithm-side memory reselection; validation/test should use dataset-provided memory just like training.
+- Keep target-only loss and target-only sampling updates.
+
+## 6. Isolated DeMemWM Model Changes
+
+Modify only the copied model files under `algorithms/dememwm/models`.
+
+Required changes:
+
+```text
+algorithms/dememwm/models/diffusion.py
+algorithms/dememwm/models/attention.py
+algorithms/dememwm/models/dit.py
+```
+
+### `diffusion.py`
+
+- Accept `memory_segments`, `memory_masks`, and memory pose metadata.
+- Return predictions only for the target segment.
+- Compute training loss only on target frames.
+- During sampling, update only target frames.
+
+### `attention.py`
+
+- Add segment-aware temporal masking.
+- Prevent normal temporal self-attention from leaking target/memory information across streams.
+- Preserve ordinary behavior when no memory metadata is provided.
+
+### `dit.py`
+
+Add separate reference attention branches:
+
+```text
+r_attn_anchor
+r_attn_dynamic
+r_attn_revisit
+```
+
+The target segment provides queries. Memory streams provide keys and values. Geometry-aware FOV/Plucker information can be used in this path.
+
+No DeMemWM-specific changes should be made to `algorithms/worldmem/models/*`.
+
+## 7. Avoid Heavy Algorithm-Side Memory Package
+
+Do not port the noisy `algorithms/worldmem/dememwm/*` structure from `~/DeMemWM`.
+
+Memory selection should live in the dataset path. If helper logic becomes too large for the dataset file, place it in dataset infrastructure:
+
+```text
+datasets/video/memory_selection.py
+```
+
+Start inside `minecraft_video_dememwm_latent_dataset.py` if possible, then split only if readability requires it.
+
+## 8. Runtime And Checkpoint Infrastructure
+
+Merge the useful infrastructure from `~/DeMemWM` into shared runtime files:
+
+- safer `load_custom_checkpoint`;
+- compatible prefix matching;
+- shape mismatch reports;
+- strict DeMemWM checkpoint checks;
+- zero new frame-memory gates when loading base WorldMem/Oasis weights;
+- auto-resume;
+- incomplete checkpoint filtering.
+
+This belongs in shared `main.py` and `experiments/exp_base.py`, because it benefits both WorldMem and DeMemWM.
+
+## 9. Tests
+
+Add tests early.
+
+Dataset tests:
+
+- dataset appends memory in `[target][anchor][dynamic][revisit]` order;
+- `memory_segments` sums to returned sequence length;
+- masks match padded/unavailable memory slots;
+- `worldmem` selector returns FOV-selected memory;
+- `dememwm` selector returns anchor/dynamic/revisit correctly;
+- causal invariants are enforced.
+
+Model/algorithm tests:
+
+- registered `algorithm=dememwm_base` + `dataset=video_minecraft_dememwm_latent` smoke path consumes one collated batch;
+- DeMemWM does not VAE-encode precomputed latents;
+- validation/test does not predict or score appended memory frames;
+- Plucker/ray geometry uses `image_hw`, not latent tensor size;
+- padded memory slots are ignored through `memory_masks`;
+- temporal mask blocks cross-stream leakage;
+- diffusion loss only uses target frames;
+- memory action conditions are zeroed;
+- cleaner memory noise is bounded by target noise;
+- DeMemWM checkpoint mismatch reporting catches missing memory adapter keys.
+
+## 10. Launch Scripts Later
+
+Do not work on Berzelius launch scripts until the local smoke path works.
+
+First milestone:
+
+```bash
+python -m main \
+  +name=dememwm_smoke \
+  algorithm=dememwm_base \
+  dataset=video_minecraft_dememwm_latent
+```
+
+Use a tiny dataset slice and verify one dataloader pass plus one forward/training step before moving to full training scripts.
+

algorithms/dememwm/__init__.py

@@ -0,0 +1,37 @@
+from collections.abc import Mapping
+
+from omegaconf import open_dict
+
+from .df_video import WorldMemMinecraft
+from .pose_prediction import PosePrediction
+
+
+_SEGMENT_KEYS = ("anchor", "dynamic", "revisit")
+
+
+def _cfg_get(cfg, key: str, default=None):
+    if cfg is None:
+        return default
+    if isinstance(cfg, Mapping):
+        return cfg.get(key, default)
+    return getattr(cfg, key, default)
+
+
+def _derive_memory_condition_length(cfg) -> int:
+    memory_cfg = _cfg_get(cfg, "memory_selection")
+    if memory_cfg is None:
+        value = _cfg_get(cfg, "memory_condition_length")
+        if value is None:
+            raise ValueError("DeMemWM requires memory_selection or memory_condition_length")
+        return int(value)
+    return sum(int(_cfg_get(memory_cfg, f"max_{key}_frames", 0)) for key in _SEGMENT_KEYS)
+
+
+class DeMemWMMinecraft(WorldMemMinecraft):
+    """DeMemWM namespace wrapper for the copied WorldMem video algorithm."""
+
+    def __init__(self, cfg):
+        if _cfg_get(cfg, "memory_condition_length") is None:
+            with open_dict(cfg):
+                cfg.memory_condition_length = _derive_memory_condition_length(cfg)
+        super().__init__(cfg)

algorithms/dememwm/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/dememwm/df_video.py

@@ -0,0 +1,954 @@
+import os
+import random
+import math
+import numpy as np
+import torch
+import torch.distributed as dist
+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 WorldMemMinecraft(DiffusionForcingBase):
+    """
+    Video generation for MineCraft with memory.
+    """
+
+    def __init__(self, cfg: DictConfig):
+        """
+        Initialize the WorldMemMinecraft 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(sync_on_compute=False)
+        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_start(self) -> None:
+        self._reset_metric_accumulators()
+
+    def on_test_epoch_start(self) -> None:
+        self._reset_metric_accumulators()
+
+    def on_validation_epoch_end(self) -> None:
+        self._on_eval_epoch_end()
+
+    def on_test_epoch_end(self) -> None:
+        self._on_eval_epoch_end()
+
+    def _reset_metric_accumulators(self) -> None:
+        self._metric_device = next(self.validation_lpips_model.parameters()).device
+        self._mse_sum = torch.tensor(0.0, device=self._metric_device)
+        self._mse_count = torch.tensor(0.0, device=self._metric_device)
+        self._psnr_sum = torch.tensor(0.0, device=self._metric_device)
+        self._psnr_count = torch.tensor(0.0, device=self._metric_device)
+        self._lpips_sum = torch.tensor(0.0, device=self._metric_device)
+        self._lpips_count = torch.tensor(0.0, device=self._metric_device)
+
+    def _update_metric_accumulators(self, xs_pred: torch.Tensor, xs_gt: torch.Tensor) -> None:
+        if not hasattr(self, "_metric_device"):
+            self._reset_metric_accumulators()
+
+        xs_pred_device = xs_pred.to(self._metric_device)
+        xs_gt_device = xs_gt.to(self._metric_device)
+        metric_dict = get_validation_metrics_for_videos(
+            xs_pred_device,
+            xs_gt_device,
+            lpips_model=self.validation_lpips_model,
+            lpips_batch_size=self.lpips_batch_size,
+        )
+
+        mse_count = torch.tensor(float(xs_pred_device.numel()), device=self._metric_device)
+        psnr_count = torch.tensor(float(xs_pred_device.shape[1]), device=self._metric_device)
+        lpips_count = torch.tensor(float(xs_pred_device.shape[0] * xs_pred_device.shape[1]), device=self._metric_device)
+
+        self._mse_sum += metric_dict["mse"].detach() * mse_count
+        self._mse_count += mse_count
+        self._psnr_sum += metric_dict["psnr"].detach() * psnr_count
+        self._psnr_count += psnr_count
+        self._lpips_sum += torch.tensor(float(metric_dict["lpips"]), device=self._metric_device) * lpips_count
+        self._lpips_count += lpips_count
+
+        del xs_pred_device, xs_gt_device
+
+    def _on_eval_epoch_end(self) -> None:
+        if not hasattr(self, "_metric_device"):
+            return
+
+        if dist.is_available() and dist.is_initialized():
+            for tensor in (
+                self._mse_sum,
+                self._mse_count,
+                self._psnr_sum,
+                self._psnr_count,
+                self._lpips_sum,
+                self._lpips_count,
+            ):
+                dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+
+        if self._mse_count.item() > 0:
+            self.log_dict(
+                {
+                    "mse": self._mse_sum / self._mse_count.clamp_min(1.0),
+                    "psnr": self._psnr_sum / self._psnr_count.clamp_min(1.0),
+                    "lpips": self._lpips_sum / self._lpips_count.clamp_min(1.0),
+                },
+                sync_dist=False,
+            )
+
+        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))
+
+        if self.logger and self.log_video:
+            log_video(
+                xs_pred,
+                xs_decode,
+                step=batch_idx,
+                namespace=namespace + "_vis",
+                context_frames=self.context_frames,
+                logger=self.logger.experiment,
+                save_local=self.save_local,
+                local_save_dir=self.local_save_dir,
+            )
+
+        self._update_metric_accumulators(xs_pred, xs_decode)
+        return
+
+    def test_step(self, batch, batch_idx) -> STEP_OUTPUT:
+        return self.validation_step(batch, batch_idx, namespace="test")
+
+    @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/dememwm/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/dememwm/models/cameractrl_module.py

@@ -0,0 +1,12 @@
+import torch.nn as nn
+class SimpleCameraPoseEncoder(nn.Module):
+    def __init__(self, c_in, c_out, hidden_dim=128):
+        super(SimpleCameraPoseEncoder, self).__init__()
+        self.model = nn.Sequential(
+            nn.Linear(c_in, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, c_out)
+        )
+    def forward(self, x):
+        return self.model(x)
+

algorithms/dememwm/models/diffusion.py

@@ -0,0 +1,520 @@
+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,
+        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.ref_mode = ref_mode
+
+        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,
+                                            pose_cond_dim=self.pose_cond_dim, reference_length=self.reference_length,
+                                            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)
+        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):
+        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)
+        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):
+        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)
+        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
+    ):
+        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)
+
+        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
+    ):
+        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
+            )
+
+        # 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
+        )
+
+    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,
+    ):
+        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
+            )
+
+        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
+    ):
+        # 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
+                )
+
+                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
+            )
+            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/dememwm/models/dit.py

@@ -0,0 +1,572 @@
+"""
+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, MemTemporalAxialAttention, MemFullAttention
+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
+from .cameractrl_module import SimpleCameraPoseEncoder
+
+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
+
+
+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,
+        reference_rotary_emb=None,
+        use_plucker=False,
+        relative_embedding=False,
+        state_embed_only_on_qk=False,
+        use_memory_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.use_memory_attention = use_memory_attention
+        if self.use_memory_attention:
+            self.r_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+            self.ref_type = "full_ref"
+            if self.ref_type == "temporal_ref":
+                self.r_attn = MemTemporalAxialAttention(
+                    hidden_size,
+                    heads=num_heads,
+                    dim_head=hidden_size // num_heads,
+                    is_causal=is_causal,
+                    rotary_emb=None
+                )
+            elif self.ref_type == "full_ref":
+                self.r_attn = MemFullAttention(
+                    hidden_size,
+                    heads=num_heads,
+                    dim_head=hidden_size // num_heads,
+                    is_causal=is_causal,
+                    rotary_emb=reference_rotary_emb,
+                    reference_length=reference_length
+                )
+            self.r_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+            self.r_mlp = Mlp(
+                in_features=hidden_size,
+                hidden_features=mlp_hidden_dim,
+                act_layer=approx_gelu,
+                drop=0,
+            )
+            self.r_adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
+
+            self.use_plucker = use_plucker
+            if use_plucker:
+                self.pose_cond_mlp = nn.Linear(hidden_size, hidden_size)
+                self.temporal_pose_cond_mlp = nn.Linear(hidden_size, hidden_size)
+
+        self.reference_length = reference_length
+        self.relative_embedding = relative_embedding
+        self.state_embed_only_on_qk = state_embed_only_on_qk
+
+        self.ref_mode = ref_mode
+
+        if self.ref_mode == 'parallel':
+            self.parallel_map = nn.Linear(hidden_size, hidden_size)
+
+    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):
+        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
+
+        # memory block
+        relative_embedding = self.relative_embedding # and mode == "training"
+
+        if self.use_memory_attention:
+            r_shift_msa, r_scale_msa, r_gate_msa, r_shift_mlp, r_scale_mlp, r_gate_mlp = self.r_adaLN_modulation(c).chunk(6, dim=-1)
+
+            if pose_cond is not None:
+                if self.use_plucker:
+                    input_cond = self.pose_cond_mlp(pose_cond)
+                    
+                    if relative_embedding:
+                        n_frames = x.shape[1] - reference_length
+                        x1_relative_embedding = []
+                        r_shift_msa_relative_embedding = []
+                        r_scale_msa_relative_embedding = []
+                        for i in range(n_frames):
+                            x1_relative_embedding.append(torch.cat([x[:,i:i+1], x[:, -reference_length:]], dim=1).clone())
+                            r_shift_msa_relative_embedding.append(torch.cat([r_shift_msa[:,i:i+1], r_shift_msa[:, -reference_length:]], dim=1).clone())
+                            r_scale_msa_relative_embedding.append(torch.cat([r_scale_msa[:,i:i+1], r_scale_msa[:, -reference_length:]], dim=1).clone())
+                        x1_zero_frame = torch.cat(x1_relative_embedding, dim=1)
+                        r_shift_msa = torch.cat(r_shift_msa_relative_embedding, dim=1)
+                        r_scale_msa = torch.cat(r_scale_msa_relative_embedding, dim=1)
+
+                        # if current_frame == 18:
+                        #     import pdb;pdb.set_trace()
+
+                        if self.state_embed_only_on_qk:
+                            attn_input = x1_zero_frame
+                            extra_condition = input_cond                        
+                        else:
+                            attn_input = input_cond + x1_zero_frame
+                            extra_condition = None
+                    else:
+                        attn_input = input_cond + x
+                        extra_condition = None
+                    # print("input_cond2:", input_cond.abs().mean())
+                    # print("c:", c.abs().mean())
+                    # input_cond = x1
+
+                    x = x + gate(self.r_attn(modulate(self.r_norm1(attn_input), r_shift_msa, r_scale_msa), 
+                                    relative_embedding=relative_embedding, 
+                                    extra_condition=extra_condition, 
+                                    state_embed_only_on_qk=self.state_embed_only_on_qk,
+                                    reference_length=reference_length), r_gate_msa)
+                else:
+                    # pose_cond *= 0
+                    x = x + gate(self.r_attn(modulate(self.r_norm1(x+pose_cond[:,:,None, None]), r_shift_msa, r_scale_msa), 
+                                    current_frame=current_frame, timestep=timestep,
+                                    is_last_block=is_last_block,
+                                    reference_length=reference_length), r_gate_msa)
+            else:
+                x = x + gate(self.r_attn(modulate(self.r_norm1(x), r_shift_msa, r_scale_msa), current_frame=current_frame, timestep=timestep,
+                                is_last_block=is_last_block), r_gate_msa)
+
+            x = x + gate(self.r_mlp(modulate(self.r_norm2(x), r_shift_mlp, r_scale_mlp)), r_gate_mlp)
+
+        if self.ref_mode == 'parallel':
+            x = x_t + self.parallel_map(x)
+
+        return x
+
+        # print((x1-x2).abs().sum())
+        # r_shift_msa, r_scale_msa, r_gate_msa, r_shift_mlp, r_scale_mlp, r_gate_mlp = self.r_adaLN_modulation(c).chunk(6, dim=-1)
+        # x2 = x1 + gate(self.r_attn(modulate(self.r_norm1(x_), r_shift_msa, r_scale_msa)), r_gate_msa)
+        # x2 = gate(self.r_mlp(modulate(self.r_norm2(x2), r_shift_mlp, r_scale_mlp)), r_gate_mlp)
+        # x = x1 + x2
+
+        # print(x.mean())
+        # 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,
+        pose_cond_dim=4,
+        max_frames=32,
+        reference_length=8,
+        use_plucker=False,
+        relative_embedding=False,
+        state_embed_only_on_qk=False,
+        use_memory_attention=False,
+        add_timestamp_embedding=False,
+        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.add_timestamp_embedding = add_timestamp_embedding
+        if self.add_timestamp_embedding:
+            self.timestamp_embedding = TimestepEmbedder(hidden_size)
+
+        frame_h, frame_w = self.x_embedder.grid_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.reference_rotary_emb = RotaryEmbedding(dim=hidden_size // num_heads // 2, freqs_for="pixel", max_freq=256)
+        self.reference_rotary_emb = None
+
+        self.external_cond = nn.Linear(action_cond_dim, hidden_size) if action_cond_dim > 0 else nn.Identity()
+
+        # self.pose_cond = nn.Linear(pose_cond_dim, hidden_size) if pose_cond_dim > 0 else nn.Identity()
+        
+        self.use_plucker = use_plucker
+        if not self.use_plucker:
+            self.position_embedder = TimestepEmbedder(hidden_size, freq_type='spatial')
+            self.angle_embedder = TimestepEmbedder(hidden_size, freq_type='angle')
+        else:
+            self.pose_embedder = SimpleCameraPoseEncoder(c_in=6, c_out=hidden_size)
+
+        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,
+                    reference_rotary_emb=self.reference_rotary_emb,
+                    use_plucker=self.use_plucker,
+                    relative_embedding=relative_embedding,
+                    state_embed_only_on_qk=state_embed_only_on_qk,
+                    use_memory_attention=use_memory_attention,
+                    ref_mode=ref_mode
+                )
+                for _ in range(depth)
+            ]
+        )
+        self.use_memory_attention = use_memory_attention
+        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)
+
+        if self.use_memory_attention:
+            if not self.use_plucker:
+                nn.init.normal_(self.position_embedder.mlp[0].weight, std=0.02)
+                nn.init.normal_(self.position_embedder.mlp[2].weight, std=0.02)
+
+                nn.init.normal_(self.angle_embedder.mlp[0].weight, std=0.02)
+                nn.init.normal_(self.angle_embedder.mlp[2].weight, std=0.02)
+            
+            if self.add_timestamp_embedding:
+                nn.init.normal_(self.timestamp_embedding.mlp[0].weight, std=0.02)
+                nn.init.normal_(self.timestamp_embedding.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)
+
+            if self.use_plucker and self.use_memory_attention:
+                nn.init.constant_(block.pose_cond_mlp.weight, 0)
+                nn.init.constant_(block.pose_cond_mlp.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)
+
+    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):
+        """
+        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):
+            try:
+                c_action_cond = c + self.external_cond(action_cond)
+            except:
+                import pdb;pdb.set_trace()
+        else:
+            c_action_cond = None
+        
+        if torch.is_tensor(pose_cond):
+            if not self.use_plucker:
+                pose_cond = pose_cond.to(action_cond.dtype)
+                b_, t_, d_ = pose_cond.shape
+                pos_emb = self.position_embedder(rearrange(pose_cond[...,:3], "b t d -> (b t d)"))
+                angle_emb = self.angle_embedder(rearrange(pose_cond[...,3:], "b t d -> (b t d)"))
+                pos_emb = rearrange(pos_emb, "(b t d) c -> b t d c", b=b_, t=t_, d=3).sum(-2)
+                angle_emb = rearrange(angle_emb, "(b t d) c -> b t d c", b=b_, t=t_, d=2).sum(-2)
+                pc = pos_emb + angle_emb
+            else:
+                pose_cond = pose_cond[:, :, ::40, ::40]
+                # pc = self.pose_embedder(pose_cond)[0]
+                # pc = pc.permute(0,2,3,4,1)
+                pc = self.pose_embedder(pose_cond)
+                pc = pc.permute(1,0,2,3,4)
+
+                if torch.is_tensor(frame_idx) and self.add_timestamp_embedding:
+                    bb = frame_idx.shape[1]
+                    frame_idx = rearrange(frame_idx, "t b -> (b t)")
+                    frame_idx = self.timestamp_embedding(frame_idx)
+                    frame_idx = rearrange(frame_idx, "(b t) d -> b t d", b=bb)
+                    pc = pc + frame_idx[:, :, None, None]
+
+                # pc = pc + rearrange(c_t.clone(), "(b t) d -> b t d", t=T)[:,:,None,None] # add time condition for different timestep scaling 
+        else:
+            pc = 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)), 
+                pose_cond=pc, mode=mode, c_action_cond=c_action_cond, reference_length=reference_length)  # (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, pose_cond_dim, reference_length, 
+use_plucker, relative_embedding, 
+state_embed_only_on_qk, use_memory_attention, add_timestamp_embedding,
+ref_mode):
+    return DiT(
+        patch_size=2,
+        hidden_size=1024,
+        depth=16,
+        num_heads=16,
+        action_cond_dim=action_cond_dim,
+        pose_cond_dim=pose_cond_dim,
+        reference_length=reference_length,
+        use_plucker=use_plucker,
+        relative_embedding=relative_embedding,
+        state_embed_only_on_qk=state_embed_only_on_qk,
+        use_memory_attention=use_memory_attention,
+        add_timestamp_embedding=add_timestamp_embedding,
+        ref_mode=ref_mode
+    )
+
+
+DiT_models = {"DiT-S/2": DiT_S_2}

algorithms/dememwm/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/dememwm/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/dememwm/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/dememwm/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,
+}

algorithms/dememwm/pose_prediction.py

@@ -0,0 +1,374 @@
+from omegaconf import DictConfig
+import torch
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from algorithms.common.metrics import (
+    FrechetInceptionDistance,
+    LearnedPerceptualImagePatchSimilarity,
+    FrechetVideoDistance,
+)
+from .df_base import DiffusionForcingBase
+from utils.logging_utils import log_video, get_validation_metrics_for_videos
+from .models.vae import VAE_models
+from .models.dit import DiT_models
+from einops import rearrange
+from torch import autocast
+import numpy as np
+from tqdm import tqdm
+import torch.nn.functional as F
+from .models.pose_prediction import PosePredictionNet
+import torchvision.transforms.functional as TF
+import random
+from torchvision.transforms import InterpolationMode
+from PIL import Image
+import math
+from packaging import version as pver
+import torch.distributed as dist
+import matplotlib.pyplot as plt
+
+import torch
+import math
+import wandb
+
+import torch.nn as nn
+from algorithms.common.base_pytorch_algo import BasePytorchAlgo
+
+class PosePrediction(BasePytorchAlgo):
+
+    def __init__(self, cfg: DictConfig):
+
+        super().__init__(cfg)
+
+    def _build_model(self):
+        self.pose_prediction_model = PosePredictionNet()
+        vae = VAE_models["vit-l-20-shallow-encoder"]()
+        self.vae = vae.eval()
+
+    def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
+        xs, conditions, pose_conditions= batch
+        pose_conditions[:,:,3:] = pose_conditions[:,:,3:] // 15
+        xs = self.encode(xs)
+
+        b,f,c,h,w = xs.shape
+        xs = xs[:,:-1].reshape(-1, c, h, w)
+        conditions = conditions[:,1:].reshape(-1, 25)
+        offset_gt = pose_conditions[:,1:] - pose_conditions[:,:-1]
+        pose_conditions = pose_conditions[:,:-1].reshape(-1, 5)
+        offset_gt = offset_gt.reshape(-1, 5)
+        offset_gt[:, 3][offset_gt[:, 3]==23] = -1
+        offset_gt[:, 3][offset_gt[:, 3]==-23] = 1
+        offset_gt[:, 4][offset_gt[:, 4]==23] = -1
+        offset_gt[:, 4][offset_gt[:, 4]==-23] = 1
+
+        offset_pred = self.pose_prediction_model(xs, conditions, pose_conditions)
+        criterion = nn.MSELoss()
+        loss = criterion(offset_pred, offset_gt)
+        if batch_idx % 200 == 0:
+            self.log("training/loss", loss.cpu())
+        output_dict = {
+            "loss": loss}
+        return output_dict
+
+    def encode(self, x):
+        # vae encoding
+        B = x.shape[1]
+        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():
+            with autocast("cuda", dtype=torch.half):
+                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)
+        # x = x[:, :n_prompt_frames]
+        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():
+            with autocast("cuda", dtype=torch.half):
+                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 validation_step(self, batch, batch_idx, namespace="validation") -> STEP_OUTPUT:
+        xs, conditions, pose_conditions= batch
+        pose_conditions[:,:,3:] = pose_conditions[:,:,3:] // 15
+        xs = self.encode(xs)
+
+        b,f,c,h,w = xs.shape
+        xs = xs[:,:-1].reshape(-1, c, h, w)
+        conditions = conditions[:,1:].reshape(-1, 25)
+        offset_gt = pose_conditions[:,1:] - pose_conditions[:,:-1]
+        pose_conditions = pose_conditions[:,:-1].reshape(-1, 5)
+        offset_gt = offset_gt.reshape(-1, 5)
+        offset_gt[:, 3][offset_gt[:, 3]==23] = -1
+        offset_gt[:, 3][offset_gt[:, 3]==-23] = 1
+        offset_gt[:, 4][offset_gt[:, 4]==23] = -1
+        offset_gt[:, 4][offset_gt[:, 4]==-23] = 1
+
+        offset_pred = self.pose_prediction_model(xs, conditions, pose_conditions)
+
+        criterion = nn.MSELoss()
+        loss = criterion(offset_pred, offset_gt)
+
+        if batch_idx % 200 == 0:
+            self.log("validation/loss", loss.cpu())
+        output_dict = {
+            "loss": loss}
+        return
+
+    @torch.no_grad()
+    def interactive(self, batch, context_frames, device):
+        with torch.cuda.amp.autocast():
+            condition_similar_length = self.condition_similar_length
+            # xs_raw, conditions, pose_conditions, c2w_mat, masks, frame_idx = self._preprocess_batch(batch)
+
+            first_frame, new_conditions, new_pose_conditions, new_c2w_mat, new_frame_idx = batch
+
+            if self.frames is None:
+                first_frame_encode = self.encode(first_frame[None, None].to(device))
+                self.frames = first_frame_encode.to(device)
+                self.actions = new_conditions[None, None].to(device)
+                self.poses = new_pose_conditions[None, None].to(device)
+                self.memory_c2w = new_c2w_mat[None, None].to(device)
+                self.frame_idx = torch.tensor([[new_frame_idx]]).to(device)
+                return first_frame
+            else:
+                self.actions = torch.cat([self.actions, new_conditions[None, None].to(device)])
+                self.poses = torch.cat([self.poses, new_pose_conditions[None, None].to(device)])
+                self.memory_c2w = torch.cat([self.memory_c2w, new_c2w_mat[None, None].to(device)])
+                self.frame_idx = torch.cat([self.frame_idx, torch.tensor([[new_frame_idx]]).to(device)])
+
+            conditions = self.actions.clone()
+            pose_conditions = self.poses.clone()
+            c2w_mat = self.memory_c2w .clone()
+            frame_idx = self.frame_idx.clone()
+
+
+            curr_frame = 0
+            horizon = 1
+            batch_size = 1
+            n_frames = curr_frame + horizon
+            # context
+            n_context_frames = context_frames // self.frame_stack
+            xs_pred = self.frames[:n_context_frames].clone()
+            curr_frame += n_context_frames
+
+            pbar = tqdm(total=n_frames, initial=curr_frame, desc="Sampling")
+
+            # generation on frame
+            scheduling_matrix = self._generate_scheduling_matrix(horizon)
+            chunk = torch.randn((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 + horizon - self.n_tokens)
+
+            pbar.set_postfix(
+                {
+                    "start": start_frame,
+                    "end": curr_frame + horizon,
+                }
+            )
+
+            if condition_similar_length:
+
+                if curr_frame < condition_similar_length:
+                    random_idx = [i for i in range(curr_frame)] + [0] * (condition_similar_length-curr_frame)
+                    random_idx = np.repeat(np.array(random_idx)[:,None], xs_pred.shape[1], -1)
+                else:
+                    num_samples = 10000
+                    radius = 30
+                    samples = torch.rand((num_samples, 1), device=pose_conditions.device)
+                    angles = 2 * np.pi * torch.rand((num_samples,), device=pose_conditions.device)
+                    # points = radius * torch.sqrt(samples) * torch.stack((torch.cos(angles), torch.sin(angles)), dim=1)
+                    
+                    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(points, pose_conditions[curr_frame, :, [0, 2]], pose_conditions[curr_frame, :, -1], fov_half)
+
+                    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_fov_list = []
+                    for pc in pose_conditions[:curr_frame]:
+                        in_fov_list.append(is_inside_fov_3d_hv(points, pc[:, :3], pc[:, -2], pc[:, -1],
+                                                        fov_half_h, fov_half_v))
+                    
+                    in_fov_list = torch.stack(in_fov_list)
+                    # v3
+                    random_idx = []
+
+                    for csl in range(self.condition_similar_length // 2):
+                        overlap_ratio = ((in_fov1[None].bool() & in_fov_list).sum(1))/in_fov1.sum()
+                        # mask = distance > (in_fov1.bool().sum(0) / 4)
+                        #_, r_idx = torch.topk(overlap_ratio / tensor_max_with_number((frame_idx[curr_frame] - frame_idx[:curr_frame]), 10), k=1, dim=0)
+                        
+                        # if csl > self.condition_similar_length:
+                        # 	_, r_idx = torch.topk(overlap_ratio, k=1, dim=0)
+                        # else:
+                        # 	_, r_idx = torch.topk(overlap_ratio / tensor_max_with_number((frame_idx[curr_frame] - frame_idx[:curr_frame]), 10), k=1, dim=0)
+
+                        _, r_idx = torch.topk(overlap_ratio, k=1, dim=0)
+                        # _, r_idx = torch.topk(overlap_ratio / tensor_max_with_number((frame_idx[curr_frame] - frame_idx[:curr_frame]), 10), k=1, dim=0)
+
+                        # if curr_frame >=93:
+                        #     import pdb;pdb.set_trace()
+
+                        # start_time = time.time()
+                        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.append(r_idx)
+
+                    for bi in range(conditions.shape[1]):
+                        if len(torch.nonzero(conditions[:,bi,24] == 1))==0:
+                            pass
+                        else:
+                            last_idx = torch.nonzero(conditions[:,bi,24] == 1)[-1]
+                            in_fov_list[:last_idx,:,bi] = False
+
+                    for csl in range(self.condition_similar_length // 2):
+                        overlap_ratio = ((in_fov1[None].bool() & in_fov_list).sum(1))/in_fov1.sum()
+                        # mask = distance > (in_fov1.bool().sum(0) / 4)
+                        #_, r_idx = torch.topk(overlap_ratio / tensor_max_with_number((frame_idx[curr_frame] - frame_idx[:curr_frame]), 10), k=1, dim=0)
+                        
+                        # if csl > self.condition_similar_length:
+                        # 	_, r_idx = torch.topk(overlap_ratio, k=1, dim=0)
+                        # else:
+                        # 	_, r_idx = torch.topk(overlap_ratio / tensor_max_with_number((frame_idx[curr_frame] - frame_idx[:curr_frame]), 10), k=1, dim=0)
+
+                        _, r_idx = torch.topk(overlap_ratio, k=1, dim=0)
+                        # _, r_idx = torch.topk(overlap_ratio / tensor_max_with_number((frame_idx[curr_frame] - frame_idx[:curr_frame]), 10), k=1, dim=0)
+
+                        # if curr_frame >=93:
+                        #     import pdb;pdb.set_trace()
+
+                        # start_time = time.time()
+                        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.append(r_idx)
+                    
+                    random_idx = torch.cat(random_idx).cpu()
+                    condition_similar_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)
+
+            if condition_similar_length:
+                # import pdb;pdb.set_trace()
+                padding = torch.zeros((condition_similar_length,) + conditions.shape[1:], device=conditions.device, dtype=conditions.dtype)
+                input_condition = torch.cat([conditions[start_frame : curr_frame + horizon], padding], dim=0)
+                if self.pose_cond_dim:
+                    # if not self.use_plucker:
+                    input_pose_condition = torch.cat([pose_conditions[start_frame : curr_frame + horizon], pose_conditions[random_idx[:,range(xs_pred.shape[1])], range(xs_pred.shape[1])]], dim=0).clone()
+
+                if self.use_plucker:
+                    if self.all_zero_frame:
+                        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(xs_pred.shape[1])], range(xs_pred.shape[1])]]).clone(), 0, focal_length=self.focal_length, is_old_setting=self.old_setting).to(xs_pred.dtype))
+                            frame_idx_list.append(torch.cat([frame_idx[i:i+1]-frame_idx[i:i+1], frame_idx[random_idx[:,range(xs_pred.shape[1])], range(xs_pred.shape[1])]-frame_idx[i:i+1]]))
+                        input_pose_condition = torch.cat(input_pose_condition)
+                        frame_idx_list = torch.cat(frame_idx_list)
+
+                        # print(frame_idx_list[:,0])
+                    else:
+                        # print(curr_frame-start_frame)
+                        # input_pose_condition = torch.cat([c2w_mat[start_frame : curr_frame + horizon], c2w_mat[random_idx[:,range(xs_pred.shape[1])], range(xs_pred.shape[1])]], dim=0).clone()
+                        # import pdb;pdb.set_trace()
+                        if self.last_frame_refer:
+                            input_pose_condition = torch.cat([c2w_mat[start_frame : curr_frame + horizon], c2w_mat[-1:]], dim=0).clone()
+                        else:
+                            input_pose_condition = torch.cat([c2w_mat[start_frame : curr_frame + horizon], c2w_mat[random_idx[:,range(xs_pred.shape[1])], range(xs_pred.shape[1])]], dim=0).clone()
+                        
+                        if self.zero_curr:
+                            # print("="*50)
+                            input_pose_condition = convert_to_plucker(input_pose_condition, curr_frame-start_frame, focal_length=self.focal_length, is_old_setting=self.old_setting)
+                            # input_pose_condition[:curr_frame-start_frame] = input_pose_condition[curr_frame-start_frame:curr_frame-start_frame+1]
+                        # input_pose_condition = convert_to_plucker(input_pose_condition, -self.condition_similar_length-1, focal_length=self.focal_length)
+                        else:
+                            input_pose_condition = convert_to_plucker(input_pose_condition, -condition_similar_length, focal_length=self.focal_length, is_old_setting=self.old_setting)
+                        frame_idx_list = None
+                else:
+                    input_pose_condition = torch.cat([pose_conditions[start_frame : curr_frame + horizon], pose_conditions[random_idx[:,range(xs_pred.shape[1])], range(xs_pred.shape[1])]], dim=0).clone()
+                    frame_idx_list = None
+            else:
+                input_condition = conditions[start_frame : curr_frame + horizon]
+                input_pose_condition = None
+                frame_idx_list = None
+                
+            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 condition_similar_length:
+                    from_noise_levels = np.concatenate([from_noise_levels, np.zeros((condition_similar_length,from_noise_levels.shape[-1]), dtype=np.int32)], axis=0)
+                    to_noise_levels = np.concatenate([to_noise_levels, np.zeros((condition_similar_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)
+
+
+                if input_pose_condition is not None:
+                    input_pose_condition = input_pose_condition.to(xs_pred.dtype)
+                
+                xs_pred[start_frame:] = self.diffusion_model.sample_step(
+                    xs_pred[start_frame:],
+                    input_condition,
+                    input_pose_condition,
+                    from_noise_levels[start_frame:],
+                    to_noise_levels[start_frame:],
+                    current_frame=curr_frame,
+                    mode="validation",
+                    reference_length=condition_similar_length,
+                    frame_idx=frame_idx_list
+                )
+
+                # if curr_frame > 14:
+                #     import pdb;pdb.set_trace()
+
+                # if xs_pred_back is not None:
+                #     xs_pred = torch.cat([xs_pred[:6], xs_pred_back[6:12], xs_pred[6:]], dim=0)
+        
+            # import pdb;pdb.set_trace()
+            if condition_similar_length: # and curr_frame+1!=n_frames:
+                xs_pred = xs_pred[:-condition_similar_length]
+
+            curr_frame += horizon
+            pbar.update(horizon)
+
+            self.frames = torch.cat([self.frames, xs_pred[n_context_frames:]])
+
+            xs_pred = self.decode(xs_pred[n_context_frames:])
+
+            return xs_pred[-1,0].cpu()
+

configurations/algorithm/dememwm_base.yaml

@@ -0,0 +1,48 @@
+defaults:
+  - df_base
+
+n_frames: ${dataset.n_frames}
+frame_skip: ${dataset.frame_skip}
+metadata: ${dataset.metadata}
+memory_selection: ${dataset.memory_selection}
+
+n_tokens: ${dataset.n_frames}
+action_cond_dim: ${dataset.action_cond_dim}
+pose_cond_dim: 5
+use_plucker: true
+relative_embedding: true
+state_embed_only_on_qk: true
+use_memory_attention: true
+add_timestamp_embedding: true
+ref_mode: sequential
+focal_length: 0.35
+log_video: false
+save_local: true
+require_pose_prediction: false
+
+# training hyperparameters
+weight_decay: 2e-3
+warmup_steps: 1000
+optimizer_beta: [0.9, 0.99]
+
+diffusion:
+  beta_schedule: sigmoid
+  objective: pred_v
+  use_fused_snr: True
+  cum_snr_decay: 0.96
+  clip_noise: 20.
+  sampling_timesteps: 20
+  ddim_sampling_eta: 0.0
+  stabilization_level: 15
+  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
+
+_name: dememwm_base

configurations/dataset/video_minecraft_dememwm_latent.yaml

@@ -0,0 +1,33 @@
+defaults:
+  - base_dataset
+
+save_dir: data/minecraft_simple_backforward_latents
+metadata: null
+n_frames: 8
+n_frames_valid: 500
+frame_skip: 1
+resolution: 128
+image_hw: [360, 640]
+observation_shape: [16, 18, 32]
+data_mean: 0.0
+data_std: 1.0
+action_cond_dim: 25
+context_length: 100
+single_eval_clip: true
+shuffle_clips: true
+
+memory_selection:
+  enabled: true
+  max_anchor_frames: 2
+  max_dynamic_frames: 4
+  max_revisit_frames: 2
+  pose_similarity_threshold: 0.6
+  training_use_plucker: false
+  training_plucker_weight: 0.1
+  fov_overlap_threshold: 0.6
+  min_total_selected_coverage: 0.1
+  local_context_exclusion_frames: 8
+  plucker_moment_radius: 30.0
+  anchor_diverse_selection: true
+
+_name: video_minecraft_dememwm_latent

datasets/video/__init__.py

@@ -1 +1,2 @@
-from .minecraft_video_dataset import MinecraftVideoDataset
+from .minecraft_video_dataset import MinecraftVideoDataset
+from .minecraft_video_dememwm_latent_dataset import MinecraftVideoDeMemWMLatentDataset

datasets/video/memory_selection.py

@@ -0,0 +1,427 @@
+"""Dataset-side memory selection for DeMemWM latent samples."""
+
+from __future__ import annotations
+
+from typing import Dict, Mapping, MutableMapping, Tuple
+
+import numpy as np
+import torch
+
+
+SEGMENT_KEYS = ("anchor", "dynamic", "revisit")
+
+
+def cfg_get(cfg, key: str, default=None):
+    if cfg is None:
+        return default
+    if isinstance(cfg, Mapping):
+        return cfg.get(key, default)
+    return getattr(cfg, key, default)
+
+
+def _as_pose_array(poses) -> np.ndarray:
+    poses = np.asarray(poses, dtype=np.float32)
+    if poses.ndim != 2 or poses.shape[1] < 5:
+        raise ValueError("poses must have shape [num_frames, >=5]")
+    return poses[:, :5]
+
+
+def _as_pose_tensor(poses) -> torch.Tensor:
+    pose_tensor = poses if torch.is_tensor(poses) else torch.as_tensor(poses, dtype=torch.float32)
+    pose_tensor = pose_tensor.to(dtype=torch.float32)
+    if pose_tensor.ndim != 2 or pose_tensor.shape[1] < 5:
+        raise ValueError("poses must have shape [num_frames, >=5]")
+    return pose_tensor[:, :5]
+
+
+def _wrap_degrees(diff: torch.Tensor) -> torch.Tensor:
+    return torch.abs(torch.remainder(diff + 180.0, 360.0) - 180.0)
+
+
+def _pose_directions(poses: torch.Tensor) -> torch.Tensor:
+    poses = _as_pose_tensor(poses)
+    pitch = torch.deg2rad(poses[:, 3])
+    yaw = torch.deg2rad(poses[:, 4])
+    cos_pitch = torch.cos(pitch)
+    directions = torch.stack(
+        [
+            torch.sin(yaw) * cos_pitch,
+            torch.sin(pitch),
+            torch.cos(yaw) * cos_pitch,
+        ],
+        dim=-1,
+    )
+    return directions / torch.linalg.vector_norm(directions, dim=-1, keepdim=True).clamp_min(1e-6)
+
+
+def _sample_points_in_sphere(center: torch.Tensor, cfg=None) -> torch.Tensor:
+    num_points = max(1, int(cfg_get(cfg, "fov_num_points", 10000)))
+    radius = float(cfg_get(cfg, "fov_radius", cfg_get(cfg, "fov_spatial_radius", 30.0)))
+    seed = int(cfg_get(cfg, "fov_sample_seed", 0))
+    device = center.device
+    dtype = center.dtype
+    generator = torch.Generator(device=device)
+    generator.manual_seed(seed)
+
+    samples_r = torch.rand(num_points, device=device, dtype=dtype, generator=generator)
+    samples_phi = torch.rand(num_points, device=device, dtype=dtype, generator=generator)
+    samples_u = torch.rand(num_points, device=device, dtype=dtype, generator=generator)
+    r = radius * torch.pow(samples_r, 1.0 / 3.0)
+    phi = 2.0 * torch.pi * samples_phi
+    theta = torch.acos(1.0 - 2.0 * samples_u)
+
+    offsets = torch.stack(
+        [
+            r * torch.sin(theta) * torch.cos(phi),
+            r * torch.sin(theta) * torch.sin(phi),
+            r * torch.cos(theta),
+        ],
+        dim=-1,
+    )
+    return center[None, :3] + offsets
+
+
+def _inside_fov_points(points: torch.Tensor, poses: torch.Tensor, cfg=None) -> torch.Tensor:
+    points = points.to(dtype=torch.float32)
+    poses = _as_pose_tensor(poses).to(device=points.device, dtype=torch.float32)
+    if points.numel() == 0 or poses.numel() == 0:
+        return torch.zeros((poses.shape[0], points.shape[0]), device=points.device, dtype=torch.bool)
+
+    fov_half_h = float(cfg_get(cfg, "fov_half_h", 105.0 / 2.0))
+    fov_half_v = float(cfg_get(cfg, "fov_half_v", 75.0 / 2.0))
+    vectors = points[None, :, :] - poses[:, None, :3]
+    x = vectors[..., 0]
+    y = vectors[..., 1]
+    z = vectors[..., 2]
+    azimuth = torch.rad2deg(torch.atan2(x, z))
+    elevation = torch.rad2deg(torch.atan2(y, torch.sqrt(x * x + z * z)))
+    yaw_diff = _wrap_degrees(azimuth - poses[:, None, 4])
+    pitch_diff = _wrap_degrees(elevation - poses[:, None, 3])
+    return (yaw_diff < fov_half_h) & (pitch_diff < fov_half_v)
+
+
+def _target_fov_points(target_poses: torch.Tensor, cfg=None) -> torch.Tensor:
+    """Sample a WorldMem-style sphere and keep points visible in the target FOV."""
+
+    target_poses = _as_pose_tensor(target_poses)
+    if target_poses.shape[0] == 0:
+        return torch.zeros((0, 3), device=target_poses.device, dtype=target_poses.dtype)
+
+    points = _sample_points_in_sphere(target_poses[0], cfg)
+    target_inside = _inside_fov_points(points, target_poses, cfg).any(dim=0)
+    return points[target_inside]
+
+
+def _pose_preselect(candidates: np.ndarray, poses_t: torch.Tensor, target_positions: np.ndarray, cfg=None) -> np.ndarray:
+    topk = cfg_get(cfg, "pose_preselect_topk", 64)
+    if topk is None:
+        return candidates
+    topk = int(topk)
+    if topk <= 0 or len(candidates) <= topk:
+        return candidates
+
+    candidates_t = torch.as_tensor(candidates, device=poses_t.device, dtype=torch.long)
+    targets_t = torch.as_tensor(target_positions, device=poses_t.device, dtype=torch.long)
+    radius = float(cfg_get(cfg, "fov_radius", cfg_get(cfg, "fov_spatial_radius", 30.0)))
+    candidate_poses = poses_t.index_select(0, candidates_t)
+    target_poses = poses_t.index_select(0, targets_t)
+    candidate_forward = _pose_directions(candidate_poses)
+    target_forward = _pose_directions(target_poses)
+    translation_norm = torch.linalg.vector_norm(candidate_poses[:, None, :3] - target_poses[None, :, :3], dim=-1) / max(radius, 1e-6)
+    dot = (candidate_forward[:, None, :] * target_forward[None, :, :]).sum(dim=-1).clamp(-1.0, 1.0)
+    angular = torch.acos(dot) / torch.pi
+    pose_distance = (translation_norm + angular).min(dim=1).values
+    rank = pose_distance.to(dtype=torch.float64) - candidates_t.to(dtype=torch.float64) * 1e-12
+    selected = torch.topk(rank, k=min(topk, int(candidates_t.numel())), largest=False, sorted=True).indices
+    return candidates_t.index_select(0, selected).cpu().numpy()
+
+
+def _candidate_fov_masks(poses_t: torch.Tensor, candidates: np.ndarray, target_positions: np.ndarray, cfg=None) -> Tuple[torch.Tensor, torch.Tensor]:
+    targets_t = torch.as_tensor(target_positions, device=poses_t.device, dtype=torch.long)
+    points = _target_fov_points(poses_t.index_select(0, targets_t), cfg)
+    if len(candidates) == 0 or points.shape[0] == 0:
+        return (
+            torch.zeros((len(candidates), 0), device=poses_t.device, dtype=torch.bool),
+            torch.zeros((len(candidates),), device=poses_t.device, dtype=torch.float32),
+        )
+
+    candidates_t = torch.as_tensor(candidates, device=poses_t.device, dtype=torch.long)
+    chunk_size = int(cfg_get(cfg, "candidate_chunk_size", 64))
+    chunk_size = len(candidates) if chunk_size <= 0 else chunk_size
+    inside_parts = []
+    score_parts = []
+    for start in range(0, len(candidates), chunk_size):
+        chunk = candidates_t[start : start + chunk_size]
+        inside = _inside_fov_points(points, poses_t.index_select(0, chunk), cfg)
+        inside_parts.append(inside)
+        score_parts.append(inside.float().mean(dim=1))
+    return torch.cat(inside_parts, dim=0), torch.cat(score_parts, dim=0)
+
+
+def _plucker_scores_tensor(candidate_poses: torch.Tensor, target_poses: torch.Tensor, cfg=None) -> torch.Tensor:
+    candidate_poses = _as_pose_tensor(candidate_poses)
+    target_poses = _as_pose_tensor(target_poses).to(device=candidate_poses.device)
+    if candidate_poses.shape[0] == 0 or target_poses.shape[0] == 0:
+        return torch.zeros((candidate_poses.shape[0],), device=candidate_poses.device, dtype=torch.float32)
+
+    moment_radius = float(cfg_get(cfg, "plucker_moment_radius", 30.0))
+    cand_dir = _pose_directions(candidate_poses)
+    target_dir = _pose_directions(target_poses)
+    cand_moment = torch.linalg.cross(candidate_poses[:, :3], cand_dir, dim=-1)
+    target_moment = torch.linalg.cross(target_poses[:, :3], target_dir, dim=-1)
+
+    direction_sim = (cand_dir @ target_dir.T).clamp(0.0, 1.0)
+    moment_dist = torch.linalg.vector_norm(cand_moment[:, None, :] - target_moment[None, :, :], dim=-1)
+    moment_sim = torch.exp(-moment_dist / max(moment_radius, 1e-6))
+    return (direction_sim * moment_sim).max(dim=1).values.to(dtype=torch.float32)
+
+
+def _empty_selection(counts: Mapping[str, int]) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray]]:
+    indices = {}
+    masks = {}
+    for key in SEGMENT_KEYS:
+        count = int(counts.get(key, 0))
+        indices[key] = np.full((count,), -1, dtype=np.int64)
+        masks[key] = np.zeros((count,), dtype=bool)
+    return indices, masks
+
+
+def _write_segment(indices: MutableMapping[str, np.ndarray], masks: MutableMapping[str, np.ndarray], key: str, selected) -> None:
+    selected = np.asarray(selected, dtype=np.int64)
+    count = len(indices[key])
+    if count == 0 or len(selected) == 0:
+        return
+    selected = selected[:count]
+    indices[key][: len(selected)] = selected
+    masks[key][: len(selected)] = True
+
+
+def _best_row(rows: torch.Tensor, gains: torch.Tensor, fov_values: torch.Tensor, plucker: torch.Tensor, gaps: torch.Tensor, candidates_t: torch.Tensor) -> int:
+    active = rows
+    for values in (gains, fov_values, plucker, -gaps.to(dtype=torch.float32), -candidates_t.to(dtype=torch.float32)):
+        vals = values.index_select(0, active)
+        active = active.index_select(0, torch.nonzero(vals == vals.max(), as_tuple=False).flatten())
+        if active.numel() == 1:
+            break
+    return int(active[0].item())
+
+
+def _select_by_pose_similarity(
+    poses: np.ndarray,
+    candidates: np.ndarray,
+    target_positions: np.ndarray,
+    cfg,
+    count: int,
+    rng=None,
+) -> np.ndarray:
+    if count <= 0 or len(candidates) == 0 or len(target_positions) == 0:
+        return np.empty((0,), dtype=np.int64)
+
+    poses_t = _as_pose_tensor(poses)
+    candidates_t = torch.as_tensor(candidates.astype(np.int64), device=poses_t.device, dtype=torch.long)
+    targets_t = torch.as_tensor(target_positions.astype(np.int64), device=poses_t.device, dtype=torch.long)
+    candidate_poses = poses_t.index_select(0, candidates_t)
+    target_poses = poses_t.index_select(0, targets_t)
+
+    # Training follows WorldMem's cheap pose-similarity sampling, not the 10k-point FOV selector.
+    spatial_scale = max(float(cfg_get(cfg, "pose_similarity_radius", cfg_get(cfg, "fov_radius", 30.0))), 1e-6)
+    angle_scale = max(float(cfg_get(cfg, "pose_similarity_angle_scale", 180.0)), 1e-6)
+    spatial_delta = torch.linalg.vector_norm(candidate_poses[:, None, :3] - target_poses[None, :, :3], dim=-1)
+    angle_delta = _wrap_degrees(candidate_poses[:, None, 3:] - target_poses[None, :, 3:]).mean(dim=-1)
+    spatial_sim = (1.0 - spatial_delta / spatial_scale).clamp(0.0, 1.0)
+    angle_sim = (1.0 - angle_delta / angle_scale).clamp(0.0, 1.0)
+    pose_scores = ((spatial_sim + angle_sim) * 0.5).max(dim=1).values
+
+    threshold = float(cfg_get(cfg, "pose_similarity_threshold", 0.6))
+    valid = pose_scores >= threshold
+    if not bool(valid.any()):
+        return np.empty((0,), dtype=np.int64)
+
+    valid_idx = torch.nonzero(valid, as_tuple=False).flatten()
+    valid_candidates = candidates_t.index_select(0, valid_idx)
+    valid_scores = pose_scores.index_select(0, valid_idx)
+
+    if bool(cfg_get(cfg, "training_use_plucker", False)):
+        plucker = _plucker_scores_tensor(
+            poses_t.index_select(0, valid_candidates),
+            target_poses,
+            cfg,
+        )
+        rank_scores = valid_scores + float(cfg_get(cfg, "training_plucker_weight", 0.1)) * plucker
+        order = torch.argsort(rank_scores, descending=True)[:count]
+        selected = valid_candidates.index_select(0, order).cpu().numpy()
+    else:
+        valid_np = valid_candidates.cpu().numpy()
+        size = min(count, len(valid_np))
+        choice = np.random.choice if rng is None else rng.choice
+        selected = choice(valid_np, size=size, replace=False)
+
+    return np.sort(np.asarray(selected, dtype=np.int64))
+
+
+def _select_by_point_union(
+    poses: np.ndarray,
+    candidates: np.ndarray,
+    target_positions: np.ndarray,
+    cfg,
+    count: int,
+    *,
+    fov_threshold: float | None = None,
+    min_total_coverage: float | None = None,
+    use_plucker: bool = False,
+) -> np.ndarray:
+    if count <= 0 or len(candidates) == 0 or len(target_positions) == 0:
+        return np.empty((0,), dtype=np.int64)
+
+    poses_t = _as_pose_tensor(poses)
+    candidates = _pose_preselect(candidates.astype(np.int64), poses_t, target_positions, cfg)
+    inside, fov_values = _candidate_fov_masks(poses_t, candidates, target_positions, cfg)
+    if inside.shape[0] == 0:
+        return np.empty((0,), dtype=np.int64)
+
+    valid = torch.ones((len(candidates),), device=poses_t.device, dtype=torch.bool)
+    if fov_threshold is not None:
+        valid &= fov_values >= float(fov_threshold)
+    if not bool(valid.any()):
+        return np.empty((0,), dtype=np.int64)
+
+    valid_idx = torch.nonzero(valid, as_tuple=False).flatten()
+    candidates_t = torch.as_tensor(candidates, device=poses_t.device, dtype=torch.long).index_select(0, valid_idx)
+    inside = inside.index_select(0, valid_idx)
+    fov_values = fov_values.index_select(0, valid_idx)
+    target_t = torch.as_tensor(target_positions, device=poses_t.device, dtype=torch.long)
+    if use_plucker:
+        plucker = _plucker_scores_tensor(poses_t.index_select(0, candidates_t), poses_t.index_select(0, target_t), cfg)
+    else:
+        plucker = torch.zeros((candidates_t.shape[0],), device=poses_t.device, dtype=torch.float32)
+
+    remaining = torch.ones((candidates_t.shape[0],), device=poses_t.device, dtype=torch.bool)
+    covered = torch.zeros((inside.shape[1],), device=poses_t.device, dtype=torch.bool)
+    selected_rows = []
+    first_target = int(target_positions[0])
+    gaps = first_target - candidates_t
+
+    for _ in range(count):
+        gains = (inside & ~covered[None, :]).float().mean(dim=1)
+        rows = torch.nonzero(remaining, as_tuple=False).flatten()
+        if rows.numel() == 0:
+            break
+        row = _best_row(rows, gains, fov_values, plucker, gaps, candidates_t)
+        if float(gains[row].item()) <= 0.0 and float(fov_values[row].item()) <= 0.0:
+            break
+        selected_rows.append(row)
+        covered |= inside[row]
+        remaining[row] = False
+
+    if not selected_rows:
+        return np.empty((0,), dtype=np.int64)
+
+    coverage = float(covered.float().mean().item()) if covered.numel() else 0.0
+    if min_total_coverage is not None and coverage < float(min_total_coverage):
+        return np.empty((0,), dtype=np.int64)
+    selected = candidates_t.index_select(0, torch.as_tensor(selected_rows, device=poses_t.device, dtype=torch.long))
+    return np.sort(selected.cpu().numpy().astype(np.int64))
+
+
+def _select_anchor(candidates: np.ndarray, count: int, cfg) -> np.ndarray:
+    if count <= 0 or len(candidates) == 0:
+        return np.empty((0,), dtype=np.int64)
+    if not bool(cfg_get(cfg, "anchor_diverse_selection", True)) or len(candidates) <= count:
+        return candidates[:count].astype(np.int64)
+
+    positions = np.linspace(0, len(candidates) - 1, num=count, dtype=np.int64)
+    selected = list(dict.fromkeys(candidates[positions].tolist()))
+    if len(selected) < count:
+        for idx in candidates:
+            if int(idx) not in selected:
+                selected.append(int(idx))
+            if len(selected) == count:
+                break
+    return np.asarray(selected[:count], dtype=np.int64)
+
+
+def _select_dynamic(target_start: int, count: int) -> np.ndarray:
+    if count <= 0 or target_start <= 0:
+        return np.empty((0,), dtype=np.int64)
+    start = max(0, target_start - count)
+    return np.arange(start, target_start, dtype=np.int64)[-count:]
+
+
+def _select_revisit(
+    poses: np.ndarray,
+    target_positions: np.ndarray,
+    cfg,
+    count: int,
+    excluded: np.ndarray,
+    split: str,
+    rng=None,
+) -> np.ndarray:
+    if count <= 0:
+        return np.empty((0,), dtype=np.int64)
+
+    target_start = int(target_positions[0])
+    exclusion = max(0, int(cfg_get(cfg, "local_context_exclusion_frames", 8)))
+    causal_stop = max(0, target_start - exclusion)
+    candidates = np.arange(0, causal_stop, dtype=np.int64)
+    if len(candidates) == 0:
+        return np.empty((0,), dtype=np.int64)
+
+    if len(excluded) > 0:
+        candidates = candidates[~np.isin(candidates, excluded)]
+    if len(candidates) == 0:
+        return np.empty((0,), dtype=np.int64)
+
+    if split == "training":
+        return _select_by_pose_similarity(poses, candidates, target_positions, cfg, count, rng=rng)
+
+    return _select_by_point_union(
+        poses,
+        candidates,
+        target_positions,
+        cfg,
+        count,
+        fov_threshold=float(cfg_get(cfg, "fov_overlap_threshold", 0.6)),
+        min_total_coverage=float(cfg_get(cfg, "min_total_selected_coverage", 0.1)),
+        use_plucker=True,
+    )
+
+
+def select_memory_indices(poses, target_positions, cfg=None, split: str = "training", rng=None):
+    """Select causal memory indices and masks for [anchor][dynamic][revisit]."""
+
+    poses = _as_pose_array(poses)
+    target_positions = np.asarray(target_positions, dtype=np.int64)
+    target_positions = target_positions[(target_positions >= 0) & (target_positions < len(poses))]
+
+    counts = {
+        "anchor": int(cfg_get(cfg, "max_anchor_frames", 0)),
+        "dynamic": int(cfg_get(cfg, "max_dynamic_frames", 0)),
+        "revisit": int(cfg_get(cfg, "max_revisit_frames", 0)),
+    }
+    indices, masks = _empty_selection(counts)
+
+    if not bool(cfg_get(cfg, "enabled", True)) or len(target_positions) == 0:
+        return indices, masks
+
+    target_start = int(target_positions[0])
+
+    dynamic = _select_dynamic(target_start, counts["dynamic"])
+    anchor_stop = max(0, target_start - counts["dynamic"])
+    anchor_candidates = np.arange(0, anchor_stop, dtype=np.int64)
+    anchor = _select_anchor(anchor_candidates, counts["anchor"], cfg)
+    excluded = np.concatenate([anchor, dynamic]) if len(anchor) or len(dynamic) else np.empty((0,), dtype=np.int64)
+    revisit = _select_revisit(poses, target_positions, cfg, counts["revisit"], excluded, split, rng=rng)
+
+    _write_segment(indices, masks, "anchor", anchor)
+    _write_segment(indices, masks, "dynamic", dynamic)
+    _write_segment(indices, masks, "revisit", revisit)
+    return indices, masks
+
+
+def memory_segment_lengths(target_length: int, cfg=None) -> Dict[str, int]:
+    return {
+        "target": int(target_length),
+        "anchor": int(cfg_get(cfg, "max_anchor_frames", 0)),
+        "dynamic": int(cfg_get(cfg, "max_dynamic_frames", 0)),
+        "revisit": int(cfg_get(cfg, "max_revisit_frames", 0)),
+    }

datasets/video/minecraft_video_dememwm_latent_dataset.py

@@ -0,0 +1,176 @@
+"""Minecraft latent dataset with dataset-side DeMemWM memory selection."""
+
+from __future__ import annotations
+
+from pathlib import Path
+import numpy as np
+import torch
+from omegaconf import DictConfig
+
+from .memory_selection import SEGMENT_KEYS, cfg_get, memory_segment_lengths, select_memory_indices
+
+
+class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
+    """Load precomputed VAE latents and append causal memory frames.
+
+    Expected files are ``.npz`` archives under ``save_dir/<split>`` with
+    ``latents``, ``actions``, and ``poses`` arrays. Optional ``frame_indices``
+    and ``image_hw`` arrays are propagated when present.
+    """
+
+    def __init__(self, cfg: DictConfig, split: str = "training"):
+        super().__init__()
+        self.cfg = cfg
+        self.split = split
+        self.save_dir = Path(cfg.save_dir)
+        self.n_frames = int(cfg_get(cfg, "n_frames_valid", cfg.n_frames)) if split in {"validation", "test"} else int(cfg.n_frames)
+        self.frame_skip = int(cfg_get(cfg, "frame_skip", 1))
+        self.target_span = max(1, (self.n_frames - 1) * self.frame_skip + 1)
+        self.action_cond_dim = int(cfg_get(cfg, "action_cond_dim", 0))
+        self.memory_selection = cfg_get(cfg, "memory_selection", {})
+        self.single_eval_clip = split in {"validation", "test"} and bool(cfg_get(cfg, "single_eval_clip", True))
+        self._target_offsets = np.arange(self.n_frames, dtype=np.int64) * self.frame_skip
+        self._memory_segments = memory_segment_lengths(self.n_frames, self.memory_selection)
+        self.memory_condition_length = sum(self._memory_segments[key] for key in SEGMENT_KEYS)
+        self.data_paths = self.get_data_paths(split)
+        if not self.data_paths:
+            raise FileNotFoundError(f"No latent .npz files found for split '{split}' under {self.save_dir}")
+
+        self.lengths = np.asarray([self._file_length(path) for path in self.data_paths], dtype=np.int64)
+        clips = np.clip(self.lengths - self.target_span + 1, a_min=1, a_max=None).astype(np.int64)
+        if self.single_eval_clip:
+            clips = np.ones_like(clips)
+        self.clips_per_file = clips
+        self.cum_clips_per_file = np.cumsum(self.clips_per_file)
+        self.idx_remap = list(range(self.__len__()))
+        if split == "training" and bool(cfg_get(cfg, "shuffle_clips", True)):
+            rng = np.random.default_rng(int(cfg_get(cfg, "seed", 0)))
+            rng.shuffle(self.idx_remap)
+
+    def get_data_paths(self, split: str):
+        split_dir = self.save_dir / split
+        data_dir = split_dir if split_dir.exists() else self.save_dir
+        return sorted(path for path in data_dir.glob("**/*.npz") if path.is_file())
+
+    def __len__(self):
+        return int(self.cum_clips_per_file[-1])
+
+    def split_idx(self, idx: int):
+        file_idx = int(np.argmax(self.cum_clips_per_file > idx))
+        prev = int(np.pad(self.cum_clips_per_file, (1, 0))[file_idx])
+        return file_idx, int(idx - prev)
+
+    def __getitem__(self, idx: int):
+        return self.load_data(self.idx_remap[idx])
+
+    def load_data(self, idx: int):
+        # === 1. Resolve clip and load arrays ===
+        file_idx, frame_idx = self.split_idx(int(idx))
+        arrays = self._load_arrays(self.data_paths[file_idx])
+        latents = arrays["latents"]
+        actions = arrays["actions"]
+        poses_pool = arrays["poses"]
+        frame_ids = arrays["frame_indices"]
+
+        # === 2. Select target and memory indices ===
+        target_positions = frame_idx + self._target_offsets
+        memory_indices, memory_masks = select_memory_indices(
+            poses_pool,
+            target_positions,
+            self.memory_selection,
+            split=self.split,
+        )
+        ordered_memory = [memory_indices[key] for key in SEGMENT_KEYS]
+        source_positions = np.concatenate([target_positions, *ordered_memory]).astype(np.int64)
+
+        # === 3. Gather arrays and normalize poses ===
+        latent_seq, valid_mask = self._gather_array(latents, source_positions)
+        action_seq, _ = self._gather_array(actions, source_positions)
+        pose_seq, _ = self._gather_array(poses_pool, source_positions)
+        frame_index_seq = self._gather_frame_indices(frame_ids, source_positions)
+
+        target_valid = valid_mask[: self.n_frames]
+        if target_valid.any():
+            ref_idx = int(target_positions[target_valid][0])
+        else:
+            ref_idx = min(max(frame_idx, 0), len(poses_pool) - 1)
+        pose_seq = self._normalize_poses(pose_seq, poses_pool[ref_idx])
+
+        masks = {"target": torch.as_tensor(target_valid, dtype=torch.bool)}
+        for key in SEGMENT_KEYS:
+            masks[key] = torch.as_tensor(memory_masks[key], dtype=torch.bool)
+
+        # === 4. Return all items ===
+        return {
+            "latents": torch.as_tensor(latent_seq).float(),
+            "actions": torch.as_tensor(action_seq).float(),
+            "poses": torch.as_tensor(pose_seq).float(),
+            "frame_indices": torch.as_tensor(frame_index_seq, dtype=torch.long),
+            "memory_segments": dict(self._memory_segments),
+            "memory_masks": masks,
+            "image_hw": torch.as_tensor(arrays["image_hw"], dtype=torch.long),
+        }
+
+    def _file_length(self, path: Path) -> int:
+        with np.load(path, allow_pickle=False) as data:
+            return int(data["latents"].shape[0])
+
+    def _load_arrays(self, path: Path):
+        with np.load(path, allow_pickle=False) as data:
+            latents = np.asarray(data["latents"])
+            actions = np.asarray(data["actions"])
+            poses = np.asarray(data["poses"], dtype=np.float32)
+            if "frame_indices" in data:
+                frame_indices = np.asarray(data["frame_indices"], dtype=np.int64)
+            else:
+                frame_indices = np.arange(len(latents), dtype=np.int64)
+
+            if "image_hw" in data:
+                image_hw = np.asarray(data["image_hw"], dtype=np.int64).reshape(-1)[:2]
+            else:
+                image_hw = np.asarray(cfg_get(self.cfg, "image_hw", [cfg_get(self.cfg, "resolution", 0), cfg_get(self.cfg, "resolution", 0)]), dtype=np.int64)
+
+        self._validate_arrays(path, latents, actions, poses, frame_indices)
+        return {
+            "latents": latents,
+            "actions": actions,
+            "poses": poses,
+            "frame_indices": frame_indices,
+            "image_hw": image_hw,
+        }
+
+    def _validate_arrays(self, path: Path, latents, actions, poses, frame_indices) -> None:
+        length = len(latents)
+        if len(actions) != length or len(poses) != length or len(frame_indices) != length:
+            raise ValueError(f"Latents/actions/poses/frame_indices length mismatch in {path}")
+        if poses.ndim != 2 or poses.shape[1] < 5:
+            raise ValueError(f"poses in {path} must have shape [num_frames, >=5]")
+        if self.action_cond_dim and actions.ndim == 2 and actions.shape[1] != self.action_cond_dim:
+            raise ValueError(
+                f"actions in {path} have dim {actions.shape[1]}, expected {self.action_cond_dim}; "
+                "precompute or configure the latent action dimension explicitly"
+            )
+
+    @staticmethod
+    def _gather_array(array: np.ndarray, indices: np.ndarray):
+        output = np.zeros((len(indices), *array.shape[1:]), dtype=array.dtype)
+        valid = (indices >= 0) & (indices < len(array))
+        if valid.any():
+            output[valid] = array[indices[valid]]
+        return output, valid
+
+    @staticmethod
+    def _gather_frame_indices(frame_ids: np.ndarray, indices: np.ndarray):
+        output = np.full((len(indices),), -1, dtype=np.int64)
+        valid = (indices >= 0) & (indices < len(frame_ids))
+        if valid.any():
+            output[valid] = frame_ids[indices[valid]]
+        return output
+
+    @staticmethod
+    def _normalize_poses(poses: np.ndarray, ref_pose: np.ndarray):
+        poses = poses.astype(np.float32, copy=True)
+        poses[:, :3] -= ref_pose[:3]
+        poses[:, -1] = -poses[:, -1]
+        poses[:, 3:] %= 360.0
+        return poses

experiments/exp_video.py

@@ -1,8 +1,10 @@
 from datasets.video import (
-    MinecraftVideoDataset
+    MinecraftVideoDataset,
+    MinecraftVideoDeMemWMLatentDataset,
 )
 
 from algorithms.worldmem import WorldMemMinecraft
+from algorithms.dememwm import DeMemWMMinecraft
 from .exp_base import BaseLightningExperiment
 
 
@@ -13,9 +15,12 @@ class VideoPredictionExperiment(BaseLightningExperiment):
 
     compatible_algorithms = dict(
         df_video_worldmemminecraft=WorldMemMinecraft,
+        dememwm_base=DeMemWMMinecraft,
     )
 
     compatible_datasets = dict(
         # video datasets
         video_minecraft=MinecraftVideoDataset,
+        video_minecraft_dememwm_latent=MinecraftVideoDeMemWMLatentDataset,
     )
+

tests/test_dememwm_latent_dataset.py

@@ -0,0 +1,187 @@
+import tempfile
+import unittest
+from pathlib import Path
+
+import numpy as np
+from omegaconf import OmegaConf
+
+from datasets.video.memory_selection import select_memory_indices
+from datasets.video.minecraft_video_dememwm_latent_dataset import MinecraftVideoDeMemWMLatentDataset
+
+
+def _selection_cfg(**overrides):
+    cfg = {
+        "enabled": True,
+        "max_anchor_frames": 2,
+        "max_dynamic_frames": 2,
+        "max_revisit_frames": 2,
+        "pose_similarity_threshold": 0.6,
+        "training_use_plucker": False,
+        "training_plucker_weight": 0.1,
+        "fov_overlap_threshold": 0.0,
+        "min_total_selected_coverage": 0.0,
+        "local_context_exclusion_frames": 2,
+        "plucker_moment_radius": 30.0,
+        "anchor_diverse_selection": False,
+        "fov_num_points": 512,
+        "fov_sample_seed": 0,
+        "pose_preselect_topk": 64,
+        "candidate_chunk_size": 64,
+    }
+    cfg.update(overrides)
+    return OmegaConf.create(cfg)
+
+
+def _poses(num_frames):
+    poses = np.zeros((num_frames, 5), dtype=np.float32)
+    poses[:, 0] = np.arange(num_frames, dtype=np.float32)
+    poses[:, 4] = np.arange(num_frames, dtype=np.float32) * 3.0
+    return poses
+
+
+class MemorySelectionTests(unittest.TestCase):
+    def test_revisit_uses_sampled_fov_selection(self):
+        poses = np.array(
+            [
+                [0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 180],
+                [0, 0, 0, 0, 180],
+                [0, 0, 0, 0, 180],
+                [0, 0, 0, 0, 180],
+                [0, 0, 0, 0, 0],
+            ],
+            dtype=np.float32,
+        )
+        cfg = _selection_cfg(
+            max_anchor_frames=0,
+            max_dynamic_frames=0,
+            max_revisit_frames=1,
+            fov_overlap_threshold=0.5,
+            local_context_exclusion_frames=1,
+        )
+        indices, masks = select_memory_indices(poses, np.array([5]), cfg, split="validation")
+
+        self.assertEqual(indices["revisit"].tolist(), [0])
+        self.assertEqual(masks["revisit"].tolist(), [True])
+
+    def test_training_revisit_uses_pose_similarity_threshold(self):
+        poses = np.array(
+            [
+                [0.2, 0, 0, 0, 2],
+                [40, 0, 0, 0, 180],
+                [45, 0, 0, 0, 180],
+                [50, 0, 0, 0, 180],
+                [55, 0, 0, 0, 180],
+                [0, 0, 0, 0, 0],
+            ],
+            dtype=np.float32,
+        )
+        cfg = _selection_cfg(
+            max_anchor_frames=0,
+            max_dynamic_frames=0,
+            max_revisit_frames=1,
+            pose_similarity_threshold=0.6,
+            local_context_exclusion_frames=1,
+        )
+        indices, masks = select_memory_indices(
+            poses,
+            np.array([5]),
+            cfg,
+            split="training",
+            rng=np.random.default_rng(0),
+        )
+
+        self.assertEqual(indices["revisit"].tolist(), [0])
+        self.assertEqual(masks["revisit"].tolist(), [True])
+
+    def test_training_revisit_pads_when_pose_matches_are_short(self):
+        poses = np.array(
+            [
+                [0.2, 0, 0, 0, 2],
+                [40, 0, 0, 0, 180],
+                [45, 0, 0, 0, 180],
+                [50, 0, 0, 0, 180],
+                [55, 0, 0, 0, 180],
+                [0, 0, 0, 0, 0],
+            ],
+            dtype=np.float32,
+        )
+        cfg = _selection_cfg(
+            max_anchor_frames=0,
+            max_dynamic_frames=0,
+            max_revisit_frames=2,
+            pose_similarity_threshold=0.95,
+            local_context_exclusion_frames=1,
+        )
+        indices, masks = select_memory_indices(
+            poses,
+            np.array([5]),
+            cfg,
+            split="training",
+            rng=np.random.default_rng(0),
+        )
+
+        self.assertEqual(indices["revisit"].tolist(), [0, -1])
+        self.assertEqual(masks["revisit"].tolist(), [True, False])
+
+    def test_dememwm_selection_is_typed_and_causal(self):
+        poses = _poses(12)
+        indices, masks = select_memory_indices(poses, np.array([6, 7, 8]), _selection_cfg())
+
+        self.assertEqual(indices["anchor"].tolist(), [0, 1])
+        self.assertEqual(indices["dynamic"].tolist(), [4, 5])
+        self.assertTrue(np.all(indices["revisit"][masks["revisit"]] < 4))
+        for key in ("anchor", "dynamic", "revisit"):
+            self.assertTrue(np.all(indices[key][masks[key]] < 6))
+        selected = np.concatenate([indices[key][masks[key]] for key in ("anchor", "dynamic", "revisit")])
+        self.assertEqual(len(selected), len(np.unique(selected)))
+
+
+class DeMemWMLatentDatasetTests(unittest.TestCase):
+    def test_dataset_returns_target_anchor_dynamic_revisit_contract(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            split_dir = root / "training"
+            split_dir.mkdir()
+            num_frames = 12
+            np.savez(
+                split_dir / "sample.npz",
+                latents=np.arange(num_frames * 2, dtype=np.float32).reshape(num_frames, 1, 1, 2),
+                actions=np.ones((num_frames, 25), dtype=np.float32),
+                poses=_poses(num_frames),
+                frame_indices=np.arange(100, 100 + num_frames, dtype=np.int64),
+                image_hw=np.array([360, 640], dtype=np.int64),
+            )
+            cfg = OmegaConf.create(
+                {
+                    "save_dir": str(root),
+                    "n_frames": 3,
+                    "n_frames_valid": 3,
+                    "frame_skip": 1,
+                    "action_cond_dim": 25,
+                    "resolution": 128,
+                    "image_hw": [360, 640],
+                    "shuffle_clips": False,
+                    "single_eval_clip": True,
+                    "memory_selection": _selection_cfg(),
+                }
+            )
+            dataset = MinecraftVideoDeMemWMLatentDataset(cfg, split="training")
+            sample = dataset[6]
+
+        self.assertEqual(sample["memory_segments"], {"target": 3, "anchor": 2, "dynamic": 2, "revisit": 2})
+        self.assertEqual(tuple(sample["latents"].shape), (9, 1, 1, 2))
+        self.assertEqual(sample["frame_indices"][:3].tolist(), [106, 107, 108])
+        self.assertEqual(sample["frame_indices"][3:5].tolist(), [100, 101])
+        self.assertEqual(sample["frame_indices"][5:7].tolist(), [104, 105])
+        self.assertTrue(all(idx < 104 for idx in sample["frame_indices"][7:9].tolist() if idx >= 0))
+        self.assertEqual(len(sample["frame_indices"][7:9]), 2)
+        self.assertTrue(sample["memory_masks"]["target"].all().item())
+        self.assertTrue(sample["memory_masks"]["anchor"].all().item())
+        self.assertTrue(sample["memory_masks"]["dynamic"].all().item())
+        self.assertTrue(sample["memory_masks"]["revisit"].all().item())
+        self.assertEqual(sample["image_hw"].tolist(), [360, 640])
+
+
+if __name__ == "__main__":
+    unittest.main()