Spaces:
Running on Zero
Running on Zero
| """Construct the complete Multics DiT without loading Wan backbone weights.""" | |
| import copy | |
| import torch | |
| from torch import nn | |
| from .models.utils import init_weights_on_device | |
| from .models.wan_video_dit import WanModel | |
| # Architecture of the Wan2.2 TI2V 5B backbone used to train Multics. | |
| DIT_CONFIG = { | |
| "has_image_input": False, | |
| "patch_size": [1, 2, 2], | |
| "in_dim": 48, | |
| "dim": 3072, | |
| "ffn_dim": 14336, | |
| "freq_dim": 256, | |
| "text_dim": 4096, | |
| "out_dim": 48, | |
| "num_heads": 24, | |
| "num_layers": 30, | |
| "eps": 1e-6, | |
| "seperated_timestep": True, | |
| "require_clip_embedding": False, | |
| "require_vae_embedding": False, | |
| "fuse_vae_embedding_in_latents": True, | |
| } | |
| def stage2_options(args): | |
| return { | |
| name: getattr(args, name) | |
| for name in ( | |
| "num_hub_tokens", "local_attn_size", "num_frame_per_block", | |
| "attention_mode", "num_groups", "num_local_hub_tokens", | |
| "num_global_hub_tokens", "hier_direct_schedule", "hier_direct_frac", | |
| "hier_direct_layers", "attn_impl", "hier_no_local_hub", | |
| "hier_local_hub_mode", | |
| ) | |
| } | |
| def build_dit(args): | |
| """Allocate parameters on meta; keep computed RoPE buffers on CPU.""" | |
| with init_weights_on_device(): | |
| model = WanModel(**DIT_CONFIG) | |
| dim = model.dim | |
| model.keyboard_embedding = nn.Linear(23, dim) | |
| for block in model.blocks: | |
| block.action_mlp = nn.Sequential( | |
| nn.LayerNorm(dim + 24, eps=1e-5, elementwise_affine=True), | |
| nn.Linear(dim + 24, dim + 24), | |
| nn.GELU(), | |
| nn.Linear(dim + 24, dim), | |
| nn.LayerNorm(dim, eps=1e-5, elementwise_affine=True), | |
| ) | |
| block.action_temporal_attn = copy.deepcopy(block.self_attn) | |
| block.action_cross_attn = copy.deepcopy(block.cross_attn) | |
| model.enable_simplex_rope( | |
| pool_size=args.simplex_pool_size, gamma=args.simplex_gamma | |
| ) | |
| model.enable_stage2(**stage2_options(args)) | |
| if args.use_camera_pose: | |
| model.enable_camera_pose_alignment( | |
| cam_dim=args.camera_dim, dropout_p=args.camera_dropout_p | |
| ) | |
| return model | |
| def load_complete_state_dict(model, state_dict, dtype=torch.bfloat16): | |
| """Reject incomplete/incompatible checkpoints before materializing parameters.""" | |
| expected = model.state_dict() | |
| missing = sorted(expected.keys() - state_dict.keys()) | |
| unexpected = sorted(state_dict.keys() - expected.keys()) | |
| mismatched = [ | |
| key for key in expected.keys() & state_dict.keys() | |
| if not isinstance(state_dict[key], torch.Tensor) | |
| or expected[key].shape != state_dict[key].shape | |
| ] | |
| if missing or unexpected or mismatched: | |
| raise RuntimeError( | |
| "Inference requires a complete Multics DiT checkpoint. " | |
| f"Missing keys ({len(missing)}): {missing[:8]}; " | |
| f"unexpected keys ({len(unexpected)}): {unexpected[:8]}; " | |
| f"incompatible tensors ({len(mismatched)}): {sorted(mismatched)[:8]}" | |
| ) | |
| model.load_state_dict(state_dict, strict=True, assign=True) | |
| if any(p.is_meta for p in model.parameters()): | |
| raise RuntimeError("Checkpoint left uninitialized DiT parameters") | |
| return model.to(dtype=dtype).eval() | |