"""Run the released DecMem memory model on the fixed WorldMem Minecraft split.""" import argparse import json import math import os from pathlib import Path import sys import torch import torch.nn.functional as F from omegaconf import OmegaConf REPO = Path(__file__).resolve().parents[1] sys.path.insert(0, str(REPO)) sys.path.insert(0, str(REPO.parent / "shared")) from protocol import case_seed, load_minecraft_case, save_case, select_cases from utils.pose_utils import poses_to_viewmats def timeline(history_frames=600, target_frames=500): prefix_padding = (1 - history_frames) % 4 context_latents = (history_frames + prefix_padding - 1) // 4 + 1 total_latents = math.ceil((context_latents + math.ceil(target_frames / 4)) / 4) * 4 generated_frames = (total_latents - context_latents) * 4 return context_latents, prefix_padding, total_latents, generated_frames def prepare_controls(case, target_frames): history = case["history"] n_cond, padding, total_latents, generated_frames = timeline(history, target_frames) poses = case["poses"][:history + target_frames].clone().float() # Match extract_pose_action_from_npz: translate xyz, negate yaw, then wrap angles. poses[:, :3] -= poses[:1, :3].clone() poses[:, -1].neg_() poses[:, 3:].remainder_(360) actions = case["actions"][:history + target_frames].float() tail = generated_frames - target_frames poses = torch.cat([poses[:1].expand(padding, -1), poses, poses[-1:].expand(tail, -1)]) actions = torch.cat([actions.new_zeros((padding, 25)), actions, actions.new_zeros((tail, 25))]) source_indices = case["source_indices"][:history + target_frames] source_indices = torch.cat([source_indices[:1].expand(padding), source_indices, source_indices[-1:].expand(tail)]) return actions, poses[::4], source_indices[::4], (n_cond, padding, total_latents, generated_frames) class DecMemAdapter: def __init__(self, checkpoint, device="cuda"): self.checkpoint = str(Path(checkpoint).resolve()) # Upstream model and VAE loading use paths relative to the DecMem checkout. os.chdir(REPO) from pipeline import WorldMemCausalDiffusionInferencePipeline self.device = torch.device(device) self.dtype = torch.bfloat16 self.config = OmegaConf.merge( OmegaConf.load(REPO / "configs/default_config.yaml"), OmegaConf.load(REPO / "configs/decmem.yaml"), ) self.config.use_text_encoder = False self.config.image_or_video_shape[1] = timeline()[2] self.config.num_frames = self.config.image_or_video_shape[1] self.config.prope_max_time = 60 self.config.rope_scale_pattern = "rope" self.pipeline = WorldMemCausalDiffusionInferencePipeline(self.config, device=self.device) checkpoint_state = torch.load(self.checkpoint, map_location="cpu", weights_only=True, mmap=True) self.pipeline.generator.load_state_dict(checkpoint_state["generator"], strict=True) self.pipeline.to(device=self.device, dtype=self.dtype).eval().requires_grad_(False) self.parameter_counts = { "generator_total": sum(p.numel() for p in self.pipeline.generator.parameters()), "generator_trainable": sum(p.numel() for p in self.pipeline.generator.parameters() if p.requires_grad), "action_encoder": sum(p.numel() for p in self.pipeline.generator.action_encoder.parameters()), "memory_projection": sum(p.numel() for n, p in self.pipeline.generator.named_parameters() if ".mem_proj." in n), "vae_total": sum(p.numel() for p in self.pipeline.vae.parameters()), "vae_trainable": sum(p.numel() for p in self.pipeline.vae.parameters() if p.requires_grad), } print(json.dumps({"parameter_counts": self.parameter_counts, "checkpoint_strict_load": "passed"}), flush=True) @torch.inference_mode() def generate(self, case, target_frames=500, seed=0): torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) history = case["history"] actions, poses, pose_indices, layout = prepare_controls(case, target_frames) n_cond, padding, total_latents, generated_frames = layout # Future RGB never reaches the VAE; only the complete observed prefix is encoded. observed = case["rgb"][:history].to(self.device) observed = F.interpolate(observed, (352, 640), mode="bilinear", align_corners=False, antialias=True) observed = torch.cat([observed[:1].expand(padding, -1, -1, -1), observed]) observed = observed.transpose(0, 1).unsqueeze(0).mul(2).sub(1).to(self.dtype) context_latent = self.pipeline.vae.encode_to_latent(observed).to(self.device, self.dtype) del observed initial_latent = context_latent.new_zeros((1, total_latents, 16, 44, 80)) initial_latent[:, :n_cond] = context_latent del context_latent actions = actions.unsqueeze(0).to(self.device, self.dtype) viewmats = poses_to_viewmats(poses.unsqueeze(0).to(self.device)).to(self.dtype) noise = torch.randn(initial_latent.shape, device=self.device, dtype=self.dtype) print(json.dumps({"observed_frames": history, "prefix_padding": padding, "condition_latents": n_cond, "total_latents": total_latents}), flush=True) _, output_latents = self.pipeline.inference( noise=noise, actions=actions, viewmats=viewmats, initial_latent=initial_latent, num_condition_frames=n_cond, guidance_scale=5.0, sampling_steps=20, return_latents=True, skip_vae_decode=True, ) # Denoising is finished. Release global KV before decoding the complete causal timeline. self.pipeline.kv_cache = None self.pipeline.crossattn_cache = None torch.cuda.empty_cache() video = self.pipeline.vae.decode_to_pixel(output_latents) prediction = video[0, history + padding:history + padding + target_frames].float().mul(0.5).add(0.5).clamp(0, 1).cpu() metadata = { "method": "decmem", "implementation": "official public Wan2.1-T2V-1.3B reference", "checkpoint": self.checkpoint, "weight_source": "released generator state; no EMA substitution", "parameter_counts": self.parameter_counts, "dtype": "bfloat16", "batch_size": 1, "cfg_batch_size": 2, "sampling": {"solver": "UniPC", "steps": 20, "guidance_scale": 5.0, "timestep_shift": 5.0, "source": "scripts/infer_example.sh"}, "memory": {"use_ltm_attn": True, "ltm_topk": 80, "ltm_tile_size": [1, 8, 8], "stm_sliding_window": 8, "num_frame_per_block": 4}, "history_frames": history, "history_source_indices": case["source_indices"][:history].tolist(), "history_left_padding": padding, "condition_latents": n_cond, "total_latents": total_latents, "generated_frames": generated_frames, "scored_frames": target_frames, "future_padding": generated_frames - target_frames, "action_mapping": "official WorldMem 25 channels, no scale changes; prepend one noop", "pose_mapping": "translate xyz relative to source100; negate yaw; angles modulo360; padded source indices[::4]", "latent_pose_source_indices": pose_indices.tolist(), "pose_matrix_precision": "float32 construction, then bfloat16 model input", "prope_max_time": 60, "prope_time_semantics": "native absolute latent indices;60 is not a sequence length cap", "native_resolution": [352, 640], "preprocessing": "full-frame bilinear resize with antialias; no crop", "target_contract": "shared fixed Oasis VAE reconstruction; no method VAE target substitution", "seed": seed, } return prediction, metadata def main(): parser = argparse.ArgumentParser() parser.add_argument("--manifest", required=True) parser.add_argument("--data-root", required=True) parser.add_argument("--checkpoint", default=str(REPO.parent / "checkpoints/decmem/decmem.pt")) parser.add_argument("--output", required=True) parser.add_argument("--limit", type=int) parser.add_argument("--max-generated", type=int, default=500) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--rank", type=int, default=int(os.environ.get("SLURM_PROCID", 0))) parser.add_argument("--world-size", type=int, default=int(os.environ.get("SLURM_NTASKS", 1))) parser.add_argument("--preflight-only", action="store_true") args = parser.parse_args() if not 1 <= args.max_generated <= 500: parser.error("--max-generated must be between 1 and the fixed 500-frame evaluation horizon") torch.set_num_threads(4) rows = select_cases(args.manifest, rank=args.rank, world_size=args.world_size, limit=args.limit) if args.preflight_only: case = load_minecraft_case(rows[0], args.data_root) actions, poses, indices, layout = prepare_controls(case, args.max_generated) report = {"selected_cases": len(rows), "actions_shape": list(actions.shape), "poses_shape": list(poses.shape), "latent_pose_source_indices": indices.tolist(), "layout": layout, "gpu_model_import": "not attempted"} output = Path(args.output) output.mkdir(parents=True, exist_ok=True) (output / f"preflight_rank{args.rank}.json").write_text(json.dumps(report, indent=2) + "\n") print(json.dumps(report), flush=True) return model = DecMemAdapter(args.checkpoint) revision = "f355a1809444a9cc6746ff99d5255614348b20ac" for row in rows: case = load_minecraft_case(row, args.data_root) prediction, metadata = model.generate(case, args.max_generated, case_seed(row, args.seed)) metadata.update({"git_revision": revision, "arguments": vars(args)}) directory = save_case(args.output, row, prediction, metadata) print(f"Saved {tuple(prediction.shape)} to {directory}", flush=True) if __name__ == "__main__": main()