update

README.md

@@ -65,5 +65,5 @@ url={https://openreview.net/forum?id=2z1HT5lw5M}
 
 ## 👏 Acknowledgements
 - [Diffusion Forcing](https://github.com/buoyancy99/diffusion-forcing): Diffusion Forcing provides flexible training and inference strategies for our methods.
-- [Minedojo](https://github.com/MineDojo/MineDojo): We collect our minecraft dataset from Minedojo.
+- [Minedojo](https://github.com/MineDojo/MineDojo): We collect our Minecraft dataset from Minedojo.
 - [Open-oasis](https://github.com/etched-ai/open-oasis): Our model architecture is based on Open-oasis. We also use pretrained VAE and DiT weight from it.

algorithms/worldmem/df_video.py

@@ -333,7 +333,6 @@ class WorldMemMinecraft(DiffusionForcingBase):
         Args:
             cfg (DictConfig): Configuration object.
         """
-        # self.metrics = cfg.metrics
         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"):
@@ -352,13 +351,6 @@ class WorldMemMinecraft(DiffusionForcingBase):
         self.log_video = cfg.log_video
         self.self_consistency_eval = getattr(cfg, "self_consistency_eval", False)
 
-        self.is_interactive = cfg.get("is_interactive", False)
-        if self.is_interactive:
-            self_frames = None
-            self_poses = None
-            self_memory_c2w = None
-            self_frame_idx = None
-
         super().__init__(cfg)
             
     def _build_model(self):
@@ -910,11 +902,4 @@ class WorldMemMinecraft(DiffusionForcingBase):
         xs_pred = self.decode(xs_pred[n_context_frames:].to(device)).cpu()
 
         return xs_pred.cpu().numpy(), self_frames.cpu().numpy(), self_actions.cpu().numpy(), \
-            self_poses.cpu().numpy(), self_memory_c2w.cpu().numpy(), self_frame_idx.cpu().numpy()
-
-
-    def reset(self):
-        self_frames = None
-        self_poses = None
-        self_memory_c2w = None
-        self_frame_idx = None
+            self_poses.cpu().numpy(), self_memory_c2w.cpu().numpy(), self_frame_idx.cpu().numpy()

algorithms/worldmem/models/attention.py

@@ -11,15 +11,6 @@ from einops import rearrange
 from .rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb
 import numpy as np
 
-def create_attention_bias(f1, f2, device=None, dtype=torch.float32):
-    f = f1 + f2
-    mask = torch.zeros((f, f), dtype=dtype, device=device)
-    if f1 > 0:
-        mask[:f1, :f1] = float('-inf')
-    if f2 > 0:
-        mask[f1:, f1:] = float('-inf')
-    return mask
-
 class TemporalAxialAttention(nn.Module):
     def __init__(
         self,

app.py

@@ -21,57 +21,11 @@ import spaces
 from algorithms.worldmem import WorldMemMinecraft
 from huggingface_hub import hf_hub_download
 import tempfile
+import os
+import requests
 
 torch.set_float32_matmul_precision("high")
 
-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",
-    "cameraY",
-    "cameraX",
-    "jump",
-    "sneak",
-    "sprint",
-    "swapHands",
-    "attack",
-    "use",
-    "pickItem",
-    "drop",
-]
-
-# Mapping of input keys to action names
-KEY_TO_ACTION = {
-    "Q": ("forward", 1),
-    "E": ("back", 1),    
-    "W": ("cameraY", -1),
-    "S": ("cameraY", 1),
-    "A": ("cameraX", -1),
-    "D": ("cameraX", 1),
-    "U": ("drop", 1),
-    "N": ("noop", 1),
-    "1": ("hotbar.1", 1),
-}
-
-example_images = [
-    ["1", "assets/ice_plains.png", "turn rightgo backward→look up→turn left→look down→turn right→go forward→turn left", 20, 3, 8],
-    ["2", "assets/place.png", "put item→go backward→put item→go backward→go around", 20, 3, 8],
-    ["3", "assets/rain_sunflower_plains.png", "turn right→look up→turn right→look down→turn left→go backward→turn left", 20, 3, 8],
-    ["4", "assets/desert.png", "turn 360 degree→turn right→go forward→turn left", 20, 3, 8],
-]
-
 def load_custom_checkpoint(algo, checkpoint_path):
     hf_ckpt = str(checkpoint_path).split('/')
     repo_id = '/'.join(hf_ckpt[:2])
@@ -81,6 +35,26 @@ def load_custom_checkpoint(algo, checkpoint_path):
     ckpt = torch.load(model_path, map_location=torch.device('cpu'))
     algo.load_state_dict(ckpt['state_dict'], strict=False)
 
+def download_assets_if_needed():
+    ASSETS_URL_BASE = "https://huggingface.co/spaces/yslan/worldmem/resolve/main/assets/examples"
+    ASSETS_DIR = "assets/examples"
+    ASSETS = ['case1.npz', 'case2.npz', 'case3.npz', 'case4.npz']
+
+    if not os.path.exists(ASSETS_DIR):
+        os.makedirs(ASSETS_DIR)
+    
+    # Download assets if they don't exist (total 4 files)
+    for filename in ASSETS:
+        filepath = os.path.join(ASSETS_DIR, filename)
+        if not os.path.exists(filepath):
+            print(f"Downloading {filename}...")
+            url = f"{ASSETS_URL_BASE}/{filename}"
+            response = requests.get(url)
+            if response.status_code == 200:
+                with open(filepath, "wb") as f:
+                    f.write(response.content)
+            else:
+                print(f"Failed to download {filename}: {response.status_code}")
 
 def parse_input_to_tensor(input_str):
     """
@@ -137,6 +111,56 @@ def enable_amp(model, precision="16-mixed"):
     model.forward = amp_forward
     return model
 
+download_assets_if_needed()
+
+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",
+    "cameraY",
+    "cameraX",
+    "jump",
+    "sneak",
+    "sprint",
+    "swapHands",
+    "attack",
+    "use",
+    "pickItem",
+    "drop",
+]
+
+# Mapping of input keys to action names
+KEY_TO_ACTION = {
+    "Q": ("forward", 1),
+    "E": ("back", 1),    
+    "W": ("cameraY", -1),
+    "S": ("cameraY", 1),
+    "A": ("cameraX", -1),
+    "D": ("cameraX", 1),
+    "U": ("drop", 1),
+    "N": ("noop", 1),
+    "1": ("hotbar.1", 1),
+}
+
+example_images = [
+    ["1", "assets/ice_plains.png", "turn rightgo backward→look up→turn left→look down→turn right→go forward→turn left", 20, 3, 8],
+    ["2", "assets/place.png", "put item→go backward→put item→go backward→go around", 20, 3, 8],
+    ["3", "assets/rain_sunflower_plains.png", "turn right→look up→turn right→look down→turn left→go backward→turn left", 20, 3, 8],
+    ["4", "assets/desert.png", "turn 360 degree→turn right→go forward→turn left", 20, 3, 8],
+]
+
 memory_frames = []
 input_history = ""
 ICE_PLAINS_IMAGE = "assets/ice_plains.png"

app.sh

@@ -1,50 +0,0 @@
-wandb disabled
-# srun -p a6000_xgpan -w MICL-PanXGSvr2 --gres=gpu:1 --ntasks-per-node=1 --cpus-per-task=8 \
-export WANDB_API_KEY=a4f0741e80f509317597ad944a7292fabcb68bdf
-
-CHECKPOINT_PATH="checkpoints/diffusion_only.ckpt"
-
-python -m app +name=pumpkin \
-    algorithm=df_video_worldmemminecraft \
-    +checkpoint_path=$CHECKPOINT_PATH \
-    experiment.tasks=[interactive] \
-    dataset.validation_multiplier=1 \
-    dataset=video_minecraft \
-    +customized_load=true \
-    +dataset.n_frames_valid=100 \
-    +algorithm.n_tokens=8 \
-    +load_vae=false \
-    +load_t_to_r=false \
-    +zero_init_gate=false \
-    experiment.validation.batch_size=1 \
-    +algorithm.pose_cond_dim=5 \
-    +algorithm.condition_similar_length=8 \
-    +dataset.condition_similar_length=8 \
-    +algorithm.use_plucker=true \
-    +dataset.use_plucker=true \
-    +dataset.padding_pool=10 \
-    +dataset.focal_length=0.35 \
-    +algorithm.focal_length=0.35 \
-    +only_tune_refer=false \
-    +dataset.customized_validation=true \
-    +algorithm.customized_validation=true \
-    algorithm.context_frames=90 \
-    +algorithm.vis_gt=true \
-    +algorithm.relative_embedding=true \
-    dataset.save_dir=data/test_pumpkin  \
-    +algorithm.log_video=true \
-    experiment.training.data.num_workers=4 \
-    experiment.validation.data.num_workers=4 \
-    +dataset.angle_range=30 \
-    +dataset.pos_range=0.5 \
-    +algorithm.cond_only_on_qk=true \
-    +algorithm.add_pose_embed=false \
-    +algorithm.use_domain_adapter=false \
-    +algorithm.use_reference_attention=true \
-    +algorithm.add_frame_timestep_embedder=true \
-    +dataset.add_frame_timestep_embedder=true \
-    experiment.validation.limit_batch=1 \
-    algorithm.diffusion.sampling_timesteps=20 \
-    +algorithm.is_interactive=true \
-    +vae_path=checkpoints/vae_only.ckpt \
-    +pose_predictor_path=checkpoints/pose_prediction_model_only.ckpt

split_checkpoint.py

@@ -1,9 +0,0 @@
-import torch
-
-ckpt_path = "/mnt/xiaozeqi/diffusionforcing/outputs/2025-03-28/16-45-11/checkpoints/epoch0step595000.ckpt"
-checkpoint = torch.load(ckpt_path, map_location="cpu")  # map_location 可根据需要更换
-
-state_dict = checkpoint['state_dict']
-pose_prediction_model_dict = {k.replace('pose_prediction_model.', ''): v for k, v in state_dict.items() if k.startswith('pose_prediction_model.')}
-
-torch.save({'state_dict': pose_prediction_model_dict}, "pose_prediction_model_only.ckpt")

test.py

@@ -1,29 +0,0 @@
-import gradio as gr
-
-css = """
-@import url('https://fonts.googleapis.com/css2?family=Press+Start+2P&display=swap');
-
-body {
-    background: linear-gradient(to bottom, #79c152 0%, #79c152 60%, #5c432d 100%) !important;
-    font-family: 'Press Start 2P', cursive !important;
-    color: #ffffff;
-}
-
-.gr-button {
-    background-color: #3e8527 !important;
-    border: 2px solid #254d16 !important;
-    color: #ffffff !important;
-}
-
-.gr-button:hover {
-    background-color: #6fcf44 !important;
-    border-color: #4a6c2d !important;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    gr.Markdown("# 🌱 Minecraft 草地界面")
-    gr.Textbox(label="你想说啥")
-    gr.Button("点我")
-
-demo.launch()