DecMem/evaluation/worldmem_adapter.py

"""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()