"""Frozen Matrix diagnostics; raw-RGB PSNR is separate from benchmark scoring.""" import argparse import json import sys from pathlib import Path import torch import torch.nn.functional as F from PIL import Image ROOT = Path(__file__).resolve().parents[1] sys.path[:0] = [str(ROOT / 'shared'), str(ROOT / 'Matrix-Game/Matrix-Game-2')] from protocol import load_minecraft_case, select_cases from worldmem_adapter import MatrixGameAdapter, map_actions from demo_utils.constant import ZERO_VAE_CACHE MODES = ['wan_oracle', 'full_original', 'full_weighted', 'i2v_original', 'i2v_weighted'] parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--mode', choices=MODES, required=True) parser.add_argument('--data-root', type=Path, required=True) parser.add_argument('--output', type=Path, required=True) parser.add_argument('--dry-run', action='store_true') def weighted_norm(self, x): y = x.float() return (y * torch.rsqrt(y.square().mean(-1, keepdim=True) + self.eps)).to(x.dtype) * self.weight def save_result(output, name, prediction, reference): prediction = prediction.float().cpu() if prediction.shape != reference.shape or not torch.isfinite(prediction).all(): raise ValueError(f'Invalid Matrix prediction: {name}, shape={tuple(prediction.shape)}') prediction = prediction.clamp(0, 1) mse = (prediction - reference).square().mean((1, 2, 3)) report = {'case': name, 'raw_rgb_psnr': (-10 * mse.clamp_min(1e-12).log10()).tolist(), 'raw_rgb_mse': mse.tolist(), 'mean_rgb': prediction.mean((1, 2, 3)).tolist()} (output / f'{name}.json').write_text(json.dumps(report, indent=2)) grid = torch.cat([torch.cat([v[t] for t in [0, 1, 2, 5, 11, 23]], dim=2) for v in (reference, prediction)], dim=1) Image.fromarray(grid.mul(255).byte().permute(1, 2, 0).numpy()).save(output / f'{name}.jpg') print(json.dumps({'case': name, 'mean_psnr': sum(report['raw_rgb_psnr']) / 24, 'first_psnr': report['raw_rgb_psnr'][:3]}), flush=True) @torch.inference_mode() def main(): args = parser.parse_args() rows = select_cases(ROOT / 'deployment/manifests/minecraft_test_300.csv', limit=2) checkpoint = ROOT / 'checkpoints/matrix-game2' manifest = {'diagnostic_only': True, 'mode': args.mode, 'checkpoint': str(checkpoint), 'weight_source': 'released raw distilled', 'optimizer': None, 'lr': None, 'lora': None, 'activation_checkpointing': False, 'batch_size': 1, 'trainable_parameters': 0, 'denoising_steps': 3, 'shift': 5, 'local_attention_latents': 6, 'selector': 'native sliding cache', 'reference': 'raw RGB resized352x640, not benchmark Oasis references', 'targets': [700, 724], 'full_history': [100, 700], 'full_prefix_latents': 153, 'oracle': '9copies+RGB100:724 encodedonce;159latents; discard153latent prefix', 'i2v': {'anchor': 699, 'decoded_slice': [1, 25], 'controls': [700, 724], 'right_noops': 9}, 'normalization_change': 'weighted modes apply scales and accumulate RMS in fp32', 'cases': [{'index': r['global_test_index'], 'path': r['relative_path'], 'seed': 42000000 + int(r['global_test_index'])} for r in rows]} if args.dry_run: print(json.dumps(manifest, indent=2)) return args.output.mkdir(parents=True, exist_ok=True) (args.output / 'manifest.json').write_text(json.dumps(manifest, indent=2)) from wan.modules.action_module import WanRMSNorm adapter = MatrixGameAdapter(checkpoint) manifest['parameter_counts'] = adapter.parameter_counts manifest['action_norm_instances'] = sum(isinstance(m, WanRMSNorm) for m in adapter.pipeline.modules()) (args.output / 'manifest.json').write_text(json.dumps(manifest, indent=2)) if args.mode.endswith('weighted'): WanRMSNorm.forward = weighted_norm for row in rows: row.update(target_end='724', target_frames='24') case = load_minecraft_case(row, args.data_root) reference = F.interpolate(case['rgb'][600:624], (352, 640), mode='bilinear', align_corners=False, antialias=True) seed = 42000000 + int(row['global_test_index']) if args.mode == 'wan_oracle': rgb = F.interpolate(case['rgb'].cuda(), (352, 640), mode='bilinear', align_corners=False, antialias=True) rgb = torch.cat([rgb[:1].expand(9, -1, -1, -1), rgb]) latent = adapter.vae.encode(rgb.transpose(0, 1)[None].mul(2).sub(1).bfloat16(), device=adapter.device, tiled=False) cache, chunks = [None] * len(ZERO_VAE_CACHE), [] for start in range(0, latent.shape[2], 3): decoded, cache = adapter.pipeline.vae_decoder(latent[:, :, start:start + 3].transpose(1, 2).half(), *cache) if start >= 153: chunks.append(decoded.cpu()) prediction = torch.cat(chunks, dim=1)[0].float().add(1).div(2) del latent, rgb, cache, decoded, chunks elif args.mode.startswith('full'): prediction, _ = adapter.generate(case, target_frames=24, seed=seed) else: torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) anchor = F.interpolate(case['rgb'][599:600].cuda(), (352, 640), mode='bilinear', align_corners=False, antialias=True) anchor = anchor.transpose(0, 1)[None].mul(2).sub(1).bfloat16() image = anchor.new_zeros((1, 3, 33, 352, 640)) image[:, :, :1] = anchor image_latents = adapter.vae.encode(image, device=adapter.device, tiled=False) mask = torch.zeros_like(image_latents[:, :4]) mask[:, :, :1] = 1 actions = torch.cat([case['actions'][600:624], torch.zeros(9, 25)]).cuda().bfloat16() keyboard, mouse = map_actions(actions, 1.0) condition = {'cond_concat': torch.cat([mask, image_latents], dim=1), 'visual_context': adapter.vae.clip.encode_video(anchor).bfloat16(), 'keyboard_cond': keyboard[None], 'mouse_cond': mouse[None]} noise = torch.randn((1, 16, 9, 44, 80), device='cuda', dtype=torch.bfloat16) videos = adapter.pipeline.inference(noise, condition, initial_latent=None, mode='universal') prediction = torch.cat(videos, dim=1)[0, 1:25].float().add(1).div(2) del image, anchor, image_latents, condition, videos, noise save_result(args.output, f"case{row['global_test_index']}_{args.mode}", prediction, reference) del prediction, reference, case (args.output / 'complete.json').write_text(json.dumps({'mode': args.mode, 'cases': 2, 'frames': 24})) print('DIAGNOSTICS_COMPLETE', flush=True) if __name__ == '__main__': main()