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