Clean up DeMemWM naming and memory path

algorithms/dememwm/__init__.py

@@ -1,37 +1,2 @@
-from collections.abc import Mapping
-
-from omegaconf import open_dict
-
-from .df_video import WorldMemMinecraft
+from .df_video import DeMemWMMinecraft
 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_video.py

@@ -12,7 +12,7 @@ from PIL import Image
 from packaging import version as pver
 from einops import rearrange
 from tqdm import tqdm
-from omegaconf import DictConfig
+from omegaconf import DictConfig, open_dict
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 from algorithms.common.metrics import (
     LearnedPerceptualImagePatchSimilarity,
@@ -29,6 +29,24 @@ _DEMEMWM_SEGMENT_KEYS = ("target", "anchor", "dynamic", "revisit")
 _DEMEMWM_STREAM_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 _DEMEMWM_STREAM_KEYS)
+
+
 def _segment_value_to_int(value, key):
     if torch.is_tensor(value):
         flat = value.reshape(-1)
@@ -432,18 +450,22 @@ def save_tensor_as_png(tensor, file_path):
     # Save image
     image.save(file_path)
 
-class WorldMemMinecraft(DiffusionForcingBase):
+class DeMemWMMinecraft(DiffusionForcingBase):
     """
-    Video generation for MineCraft with memory.
+    DeMemWM video generation for MineCraft with frame memory.
     """
 
     def __init__(self, cfg: DictConfig):
         """
-        Initialize the WorldMemMinecraft class with the given configuration.
+        Initialize the DeMemWMMinecraft class with the given configuration.
 
         Args:
             cfg (DictConfig): Configuration object.
         """
+        if _cfg_get(cfg, "memory_condition_length") is None:
+            with open_dict(cfg):
+                cfg.memory_condition_length = _derive_memory_condition_length(cfg)
+
         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"):
@@ -471,7 +493,9 @@ class WorldMemMinecraft(DiffusionForcingBase):
     def _build_model(self):
 
         self.diffusion_model = Diffusion(
-            reference_length=self.memory_condition_length,
+            # DeMemWM injects memory through explicit frame_memory_segments;
+            # keep the legacy reference_length path disabled by default.
+            reference_length=0,
             x_shape=self.x_stacked_shape,
             action_cond_dim=self.action_cond_dim,
             pose_cond_dim=self.pose_cond_dim,

algorithms/dememwm/models/attention.py

@@ -3,13 +3,11 @@ Based on https://github.com/buoyancy99/diffusion-forcing/blob/main/algorithms/di
 """
 
 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__(
@@ -224,179 +222,3 @@ class SpatialAxialAttention(nn.Module):
         # 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/diffusion.py

@@ -360,7 +360,7 @@ class Diffusion(nn.Module):
         else:
             raise ValueError(f"unknown objective {self.objective}")
 
-        # 训练的时候每个frame随便给噪声
+        # Give random noise to each frame during training
         loss = F.mse_loss(pred, target.detach(), reduction="none")
         loss_weight = self.compute_loss_weights(target_noise_levels)
 

algorithms/dememwm/models/dit.py

@@ -11,7 +11,7 @@ from torch import nn
 from torch.nn import functional as F
 from .rotary_embedding_torch import RotaryEmbedding
 from einops import rearrange
-from .attention import SpatialAxialAttention, TemporalAxialAttention, MemTemporalAxialAttention, MemFullAttention
+from .attention import SpatialAxialAttention, TemporalAxialAttention
 from timm.models.vision_transformer import Mlp
 from timm.layers.helpers import to_2tuple
 import math
@@ -258,24 +258,6 @@ class SpatioTemporalDiTBlock(nn.Module):
         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,
@@ -284,11 +266,9 @@ class SpatioTemporalDiTBlock(nn.Module):
                 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.r_attn_anchor = FrameMemoryReferenceAttention(hidden_size, num_heads)
+            self.r_attn_dynamic = FrameMemoryReferenceAttention(hidden_size, num_heads)
+            self.r_attn_revisit = FrameMemoryReferenceAttention(hidden_size, num_heads)
 
         self.reference_length = reference_length
         self.relative_embedding = relative_embedding
@@ -299,11 +279,6 @@ class SpatioTemporalDiTBlock(nn.Module):
         if self.ref_mode == 'parallel':
             self.parallel_map = nn.Linear(hidden_size, hidden_size)
 
-        if self.use_memory_attention:
-            self.r_attn_anchor = FrameMemoryReferenceAttention(hidden_size, num_heads)
-            self.r_attn_dynamic = FrameMemoryReferenceAttention(hidden_size, num_heads)
-            self.r_attn_revisit = FrameMemoryReferenceAttention(hidden_size, num_heads)
-
     def _split_frame_memory(self, x, frame_memory_segments):
         target = int(frame_memory_segments.get("target", 0))
         anchor = int(frame_memory_segments.get("anchor", 0))
@@ -389,8 +364,6 @@ class SpatioTemporalDiTBlock(nn.Module):
             x = x_t
 
         # memory block
-        relative_embedding = self.relative_embedding # and mode == "training"
-
         if self.use_memory_attention and frame_memory_segments is not None:
             x = self._apply_frame_memory_reference_attention(
                 x,
@@ -399,74 +372,12 @@ class SpatioTemporalDiTBlock(nn.Module):
                 frame_memory_masks,
                 frame_memory_geometry,
             )
-        elif self.use_memory_attention:
-            reference_length = self.reference_length if reference_length is None else reference_length
-            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), 
-                                    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),
-                    reference_length=reference_length,
-                ), 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):
     """
@@ -589,10 +500,6 @@ class DiT(nn.Module):
             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)

tests/test_dememwm_latent_dataset.py

@@ -226,7 +226,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertEqual(preprocessed["image_hw"].tolist(), [[720, 1280], [720, 1280]])
 
     def test_training_step_uses_latent_dict_target_loss_and_memory_metadata(self):
-        from algorithms.dememwm.df_video import WorldMemMinecraft, _preprocess_dememwm_latent_batch
+        from algorithms.dememwm.df_video import DeMemWMMinecraft, _preprocess_dememwm_latent_batch
 
         class FakeDiffusion:
             stabilization_level = 7
@@ -281,7 +281,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         }
         harness = Harness()
 
-        output = WorldMemMinecraft.training_step(harness, batch, 0)
+        output = DeMemWMMinecraft.training_step(harness, batch, 0)
         call = harness.diffusion_model
 
         self.assertEqual(harness.noise_input_shape, (2, 1, 1, 1, 1))
@@ -301,7 +301,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertEqual(float(harness.logged[0][1]), 2.0)
 
     def test_validation_step_uses_latent_dict_target_sampling_and_metric(self):
-        from algorithms.dememwm.df_video import WorldMemMinecraft, _preprocess_dememwm_latent_batch
+        from algorithms.dememwm.df_video import DeMemWMMinecraft, _preprocess_dememwm_latent_batch
 
         class FakeDiffusion:
             def __init__(self):
@@ -372,7 +372,7 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         }
         harness = Harness()
 
-        loss = WorldMemMinecraft.validation_step(harness, batch, 0, namespace="test")
+        loss = DeMemWMMinecraft.validation_step(harness, batch, 0, namespace="test")
         calls = harness.diffusion_model.calls
 
         self.assertEqual(harness.horizons, [2])

tests/test_dememwm_temporal_attention.py

@@ -194,6 +194,46 @@ class DeMemWMTemporalAttentionTests(unittest.TestCase):
             )
         self.assertEqual(tuple(packed_out.shape), (1, 2, 1, 2, 2))
 
+    def test_dit_has_no_worldmem_reference_attention_fallback(self):
+        model = DiT(
+            input_h=2,
+            input_w=2,
+            patch_size=1,
+            in_channels=1,
+            hidden_size=8,
+            depth=1,
+            num_heads=1,
+            mlp_ratio=1.0,
+            action_cond_dim=3,
+            pose_cond_dim=0,
+            reference_length=1,
+            use_memory_attention=True,
+        )
+        block = model.blocks[0]
+
+        self.assertFalse(hasattr(block, "r_attn"))
+        self.assertFalse(hasattr(block, "pose_cond_mlp"))
+
+        calls = []
+
+        class SpyReferenceAttention(nn.Module):
+            def forward(self, target_hidden, memory_hidden, memory_mask=None, geometry_cache=None):
+                calls.append(memory_hidden)
+                return torch.zeros_like(target_hidden)
+
+        block.r_attn_anchor = SpyReferenceAttention()
+        block.r_attn_dynamic = SpyReferenceAttention()
+        block.r_attn_revisit = SpyReferenceAttention()
+
+        x = torch.randn((1, 3, 1, 2, 2))
+        t = torch.zeros((1, 3), dtype=torch.long)
+        action_cond = torch.zeros((1, 3, 3))
+        with torch.no_grad():
+            out = model(x, t, action_cond, reference_length=1)
+
+        self.assertEqual(tuple(out.shape), (1, 3, 1, 2, 2))
+        self.assertEqual(calls, [])
+
     def test_dit_builds_geometry_cache_once_and_threads_to_reference_attention(self):
         class ZeroAttention(nn.Module):
             def forward(self, x):