Spaces:
Running on Zero
Running on Zero
| from .config import select_case_rows, case_output_id | |
| from .checkpoint import build_dit, load_complete_state_dict, stage2_options | |
| from .inference_args import parse_args | |
| import json | |
| import os | |
| import time | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| from PIL import Image | |
| from .utils import ModelConfig | |
| from .pipelines.wan_video_Nplayers import WanVideoNPlayersPipeline | |
| from .data.video import save_video | |
| from multics.models.wan_video_dit import prepare_hier_local_hub_state_dict_for_load | |
| from multics.data.dataset import ( | |
| PrecomputedLatentDataset, | |
| consecutive_precomputed_start_rows, | |
| ) | |
| from multics.data.raw_inputs import ( | |
| encode_raw_first_frames, | |
| is_raw_metadata, | |
| prepare_raw_inputs, | |
| ) | |
| from multics.data.utils import tensor_bchw_to_pil_list, pad_pil_list_to_length | |
| def _timed(section: str): | |
| """Context manager: print load-stage timing (flush for tee/log).""" | |
| class _Ctx: | |
| def __enter__(self): | |
| print(f"[timing] >> {section}", flush=True) | |
| self._t0 = time.perf_counter() | |
| return self | |
| def __exit__(self, exc_type, exc, tb): | |
| dt = time.perf_counter() - self._t0 | |
| if exc_type is None: | |
| print(f"[timing] << {section} {dt:.2f}s", flush=True) | |
| else: | |
| print( | |
| f"[timing] !! {section} failed after {dt:.2f}s: {exc}", flush=True | |
| ) | |
| return False | |
| return _Ctx() | |
| def _load_dit_checkpoint_state(ckpt_path: str): | |
| """Load a DiT checkpoint and fail fast when any saved weight is non-finite.""" | |
| partial_ckpt = False | |
| if os.path.isdir(ckpt_path): | |
| index_path = os.path.join(ckpt_path, "pytorch_model.bin.index.json") | |
| if not os.path.exists(index_path): | |
| raise FileNotFoundError( | |
| f"Checkpoint directory has no pytorch_model.bin.index.json: {ckpt_path}" | |
| ) | |
| with open(index_path, "r") as f: | |
| index = json.load(f) | |
| shards = sorted(set(index.get("weight_map", {}).values())) | |
| if not shards: | |
| raise RuntimeError(f"No shards listed in checkpoint index: {index_path}") | |
| state_dict = {} | |
| for shard in shards: | |
| shard_path = os.path.join(ckpt_path, shard) | |
| with _timed(f"torch.load DiT ckpt shard {shard}"): | |
| shard_state = torch.load( | |
| shard_path, map_location="cpu", weights_only=True | |
| ) | |
| if not isinstance(shard_state, dict): | |
| raise RuntimeError( | |
| f"Checkpoint shard is not a state_dict: {shard_path}" | |
| ) | |
| state_dict.update(shard_state) | |
| partial_ckpt = True | |
| else: | |
| with _timed(f"torch.load DiT ckpt {os.path.basename(ckpt_path)}"): | |
| state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=True) | |
| if ( | |
| isinstance(state_dict, dict) | |
| and "state_dict" in state_dict | |
| and isinstance(state_dict["state_dict"], dict) | |
| ): | |
| state_dict = state_dict["state_dict"] | |
| if not isinstance(state_dict, dict): | |
| raise RuntimeError(f"DiT checkpoint is not a state_dict: {ckpt_path}") | |
| if any((k.startswith("pipe.dit.") for k in state_dict.keys())): | |
| stripped = {} | |
| for key, value in state_dict.items(): | |
| if key.startswith("pipe.dit."): | |
| stripped[key[len("pipe.dit.") :]] = value | |
| state_dict = stripped | |
| partial_ckpt = True | |
| bad = [] | |
| for key, value in state_dict.items(): | |
| if not torch.is_tensor(value) or not ( | |
| value.is_floating_point() or value.is_complex() | |
| ): | |
| continue | |
| try: | |
| finite = bool(torch.isfinite(value).all().item()) | |
| except RuntimeError as exc: | |
| bad.append( | |
| f"{key}: shape={tuple(value.shape)} dtype={value.dtype} check_failed={exc}" | |
| ) | |
| continue | |
| if not finite: | |
| bad.append(f"{key}: shape={tuple(value.shape)} dtype={value.dtype}") | |
| if len(bad) >= 12: | |
| break | |
| if bad: | |
| raise RuntimeError( | |
| f"DiT checkpoint contains non-finite tensor(s); refusing to run inference because this will generate NaN/black video.\nckpt_path={ckpt_path}\n" | |
| + "\n".join((f" - {item}" for item in bad)) | |
| ) | |
| return (state_dict, partial_ckpt) | |
| def _patch_model_manager_load_timing(): | |
| """Log each Wan base weight file load inside from_pretrained.""" | |
| from multics.models import model_manager as mm | |
| if getattr(mm.ModelManager.load_model, "_infer_timing_patched", False): | |
| return | |
| _orig = mm.ModelManager.load_model | |
| def _load_model_timed( | |
| self, file_path, model_names=None, device=None, torch_dtype=None | |
| ): | |
| if isinstance(file_path, list): | |
| label = f"ModelManager.load_model [{len(file_path)} shards]" | |
| else: | |
| label = f"ModelManager.load_model {os.path.basename(str(file_path))}" | |
| with _timed(label): | |
| return _orig( | |
| self, | |
| file_path, | |
| model_names=model_names, | |
| device=device, | |
| torch_dtype=torch_dtype, | |
| ) | |
| _load_model_timed._infer_timing_patched = True | |
| mm.ModelManager.load_model = _load_model_timed | |
| NEGATIVE_PROMPT = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" | |
| def _batch_output_names(batch_idx: int, args) -> list: | |
| """Filenames written for one dataloader row. batch_idx is the CSV row index.""" | |
| names = [] | |
| if not args.no_save_separate_players: | |
| for pidx in range(args.num_players): | |
| names.append( | |
| f"video{batch_idx}_{_player_file_tag(pidx, args.num_players)}.mp4" | |
| ) | |
| if not args.no_save_grid: | |
| names.append(f"video{batch_idx}_grid.mp4") | |
| return names | |
| def _list_mp4_basenames(output_dir, s3_uri=None): | |
| return ( | |
| {name for name in os.listdir(output_dir) if name.endswith(".mp4")} | |
| if os.path.isdir(output_dir) | |
| else set() | |
| ) | |
| def _existing_batch_indices( | |
| n_planned: int, args, existing: set, output_ids=None | |
| ) -> set: | |
| done = set() | |
| for batch_idx in range(n_planned): | |
| expected = _batch_output_names( | |
| output_ids[batch_idx] if output_ids is not None else batch_idx, args | |
| ) | |
| if expected and all((name in existing for name in expected)): | |
| done.add(batch_idx) | |
| return done | |
| def _infer_rank_world(): | |
| rank = int(os.environ.get("RANK", "0")) | |
| world_size = int(os.environ.get("WORLD_SIZE", "1")) | |
| if world_size < 1: | |
| raise ValueError(f"WORLD_SIZE must be >= 1, got {world_size}") | |
| if rank < 0 or rank >= world_size: | |
| raise ValueError(f"RANK={rank} out of range for WORLD_SIZE={world_size}") | |
| return (rank, world_size) | |
| def _infer_shard_rows( | |
| n_planned: int, skip_batches: set, rank: int, world_size: int | |
| ) -> list: | |
| """CSV row ids this rank should generate. Names stay video{row}_*.mp4.""" | |
| return [ | |
| i for i in range(n_planned) if i % world_size == rank and i not in skip_batches | |
| ] | |
| def _infer_row_bounds(n_planned: int) -> tuple[int, int]: | |
| """Optional half-open row range for splitting one result repo across jobs.""" | |
| start = int(os.environ.get("INFER_ROW_START", "0")) | |
| end = int(os.environ.get("INFER_ROW_END", str(n_planned))) | |
| if start < 0 or end < start or end > n_planned: | |
| raise ValueError( | |
| f"invalid inference row range [{start}, {end}) for n_planned={n_planned}" | |
| ) | |
| return (start, end) | |
| TRAIN_HEIGHT = 352 | |
| TRAIN_WIDTH = 640 | |
| TILE_SIZE = (34, 34) | |
| TILE_STRIDE = (18, 16) | |
| def _resolve_metadata_path(dataset_path, metadata_file_name): | |
| if os.path.isabs(metadata_file_name): | |
| return metadata_file_name | |
| return os.path.join(dataset_path, metadata_file_name) | |
| def _precomputed_video_latents(batch) -> torch.Tensor: | |
| """[P,48,T_lat,h,w] from PrecomputedLatentDataset batch (batch_size=1).""" | |
| if "latents" in batch: | |
| if batch["latents"].shape[0] != 1: | |
| raise ValueError( | |
| "precomputed inference currently expects DataLoader batch_size=1" | |
| ) | |
| return batch["latents"].squeeze(0) | |
| return torch.cat([batch["latents1"], batch["latents2"]], dim=0) | |
| def _decode_latents_to_video( | |
| pipe, latents_b1: torch.Tensor, tiler_kwargs: dict | |
| ) -> torch.Tensor: | |
| """[1, 48, T_lat, h, w] -> [1, 3, T_rgb, H, W] for GT grid / parity.""" | |
| pipe.load_models_to_device(["vae"]) | |
| return pipe.decode_video( | |
| latents_b1.to(dtype=pipe.torch_dtype, device=pipe.device), **tiler_kwargs | |
| ) | |
| def _decode_latents_to_videos( | |
| pipe, latents_p: torch.Tensor, tiler_kwargs: dict | |
| ) -> torch.Tensor: | |
| """[P,48,T_lat,h,w] -> [P,3,T_rgb,H,W] for GT grids.""" | |
| pipe.load_models_to_device(["vae"]) | |
| return pipe.decode_video( | |
| latents_p.to(dtype=pipe.torch_dtype, device=pipe.device), **tiler_kwargs | |
| ) | |
| def _as_player_videos(outputs): | |
| """Normalize pipeline return to list length P; P=2 may be returned as a tuple for compatibility.""" | |
| if isinstance(outputs, tuple): | |
| return list(outputs) | |
| return outputs | |
| def _batch_prompt_strings(batch, num_players: int): | |
| """Return per-player prompt strings from a collated precomputed batch.""" | |
| if "prompt" not in batch: | |
| return None | |
| prompt = batch["prompt"] | |
| if isinstance(prompt, (list, tuple)): | |
| out = [ | |
| p[0] if isinstance(p, (list, tuple)) and len(p) > 0 else p for p in prompt | |
| ] | |
| else: | |
| out = [prompt] | |
| out = [str(p).strip() for p in out] | |
| if len(out) == 1 and num_players > 1: | |
| out = out * num_players | |
| if len(out) != num_players: | |
| raise ValueError( | |
| f"prompt count {len(out)} must match num_players={num_players}" | |
| ) | |
| return out | |
| def _prompt_context_from_batch(pipe, args, batch, prompt_strings): | |
| """Build T5 context [P,S,4096] or [1,S,4096] for full-sequence inference.""" | |
| prompt_emb = batch.get("prompt_emb", {}) | |
| context = prompt_emb.get("context") if isinstance(prompt_emb, dict) else None | |
| if context is not None: | |
| if context.dim() == 4: | |
| if context.shape[0] != 1: | |
| raise ValueError( | |
| f"prompt context batch must be 1, got {tuple(context.shape)}" | |
| ) | |
| context = context[0] | |
| return context.to(dtype=pipe.torch_dtype, device=pipe.device) | |
| pipe.load_models_to_device(["text_encoder"]) | |
| if prompt_strings is not None and any(prompt_strings): | |
| return pipe.encode_prompt(prompt_strings, positive=True)["context"].to( | |
| dtype=pipe.torch_dtype, device=pipe.device | |
| ) | |
| return pipe.encode_prompt("", positive=True)["context"].to( | |
| dtype=pipe.torch_dtype, device=pipe.device | |
| ) | |
| def _hier_full_sequence_inference( | |
| pipe, | |
| args, | |
| batch, | |
| video_latents: torch.Tensor, | |
| mouse: torch.Tensor, | |
| keyboard: torch.Tensor, | |
| camera_pose: torch.Tensor, | |
| agent_vertex_ids, | |
| ): | |
| """Hierarchical inference with a full-sequence block-causal forward, no KV cache. | |
| Key tensors: | |
| video_latents: [P,C,T,h,w], clean GT latents used only for frame-0 anchor. | |
| latents: [P,C,T,h,w], denoising state. | |
| timestep: [P,T], frame 0 pinned to t=0; other frames use current sigma. | |
| group_labels: [1,T,P] from dataloader, passed to WanModel.forward. | |
| active_group: [1,T,N] from dataloader, passed to WanModel.forward. | |
| """ | |
| if "group_labels" not in batch or "active_group" not in batch: | |
| raise KeyError( | |
| "hier_hub full-sequence inference requires group_labels and active_group" | |
| ) | |
| if video_latents is None: | |
| raise ValueError( | |
| "hier_hub full-sequence inference requires --precomputed_latents" | |
| ) | |
| video_latents = video_latents.to(dtype=pipe.torch_dtype, device=pipe.device) | |
| P, C, T, h, w = video_latents.shape | |
| if T == 1: | |
| T = int(batch["group_labels"].shape[1]) | |
| if P != args.num_players: | |
| raise ValueError( | |
| f"video_latents P={P} must match --num_players={args.num_players}" | |
| ) | |
| group_labels = batch["group_labels"].to(device=pipe.device, dtype=torch.long) | |
| active_group = batch["active_group"].to(device=pipe.device, dtype=torch.bool) | |
| membership = batch.get("membership") | |
| global_only = batch.get("global_only") | |
| if membership is not None: | |
| membership = membership.to(device=pipe.device, dtype=torch.bool) | |
| if global_only is None: | |
| raise KeyError( | |
| "hier_group_format=membership requires batch key global_only" | |
| ) | |
| global_only = global_only.to(device=pipe.device, dtype=torch.bool) | |
| if group_labels.shape[1] != T or group_labels.shape[2] != P: | |
| raise ValueError( | |
| f"group_labels must be [1,T={T},P={P}], got {tuple(group_labels.shape)}" | |
| ) | |
| if active_group.shape[1] != T or active_group.shape[2] != args.num_groups: | |
| raise ValueError( | |
| f"active_group must be [1,T={T},N={args.num_groups}], got {tuple(active_group.shape)}" | |
| ) | |
| mouse = mouse.to(dtype=pipe.torch_dtype, device=pipe.device) | |
| keyboard = keyboard.to(dtype=pipe.torch_dtype, device=pipe.device) | |
| if camera_pose is not None: | |
| camera_pose = camera_pose.to(dtype=pipe.torch_dtype, device=pipe.device) | |
| pipe.scheduler.set_timesteps( | |
| args.num_inference_steps, denoising_strength=1.0, shift=5.0 | |
| ) | |
| noise = pipe.generate_noise( | |
| (P, C, T, h, w), seed=args.seed, device="cpu", dtype=torch.float32 | |
| ) | |
| latents = noise.to(dtype=pipe.torch_dtype, device=pipe.device) | |
| latents[:, :, 0:1, ...] = video_latents[:, :, 0:1, ...].clone() | |
| prompt_strings = _batch_prompt_strings(batch, args.num_players) | |
| prompt_context = _prompt_context_from_batch(pipe, args, batch, prompt_strings) | |
| has_prompt = bool(prompt_strings is not None and any(prompt_strings)) | |
| has_prompt_emb = "prompt_emb" in batch and batch["prompt_emb"] is not None | |
| cfg_scale = args.cfg_scale if has_prompt else 1.0 | |
| negative_context = None | |
| if cfg_scale != 1.0: | |
| pipe.load_models_to_device(["text_encoder"]) | |
| negative_context = pipe.encode_prompt([NEGATIVE_PROMPT] * P, positive=False)[ | |
| "context" | |
| ].to(dtype=pipe.torch_dtype, device=pipe.device) | |
| print( | |
| f"[full_sequence] P={P} latent_T={T} cfg_scale={cfg_scale} prompt={('caption/emb' if has_prompt or has_prompt_emb else 'blank')}", | |
| flush=True, | |
| ) | |
| pipe.load_models_to_device(["dit"]) | |
| for progress_id, step_t in enumerate(pipe.scheduler.timesteps): | |
| timestep = torch.full((P, T), step_t, dtype=torch.float32, device=pipe.device) | |
| timestep[:, 0] = 0.0 | |
| noise_pred_pos = pipe.dit( | |
| latents, | |
| timestep=timestep, | |
| context=prompt_context, | |
| mouse=mouse, | |
| keyboard=keyboard, | |
| loss_type="diffusion_forcing", | |
| agent_vertex_ids=agent_vertex_ids, | |
| camera_pose=camera_pose, | |
| group_labels=group_labels, | |
| active_group=active_group, | |
| membership=membership, | |
| global_only=global_only, | |
| ) | |
| if cfg_scale != 1.0: | |
| noise_pred_neg = pipe.dit( | |
| latents, | |
| timestep=timestep, | |
| context=negative_context, | |
| mouse=mouse, | |
| keyboard=keyboard, | |
| loss_type="diffusion_forcing", | |
| agent_vertex_ids=agent_vertex_ids, | |
| camera_pose=camera_pose, | |
| group_labels=group_labels, | |
| active_group=active_group, | |
| membership=membership, | |
| global_only=global_only, | |
| ) | |
| noise_pred = noise_pred_neg + cfg_scale * (noise_pred_pos - noise_pred_neg) | |
| else: | |
| noise_pred = noise_pred_pos | |
| latents = pipe.scheduler.step( | |
| noise_pred, pipe.scheduler.timesteps[progress_id], latents | |
| ) | |
| latents[:, :, 0:1, ...] = video_latents[:, :, 0:1, ...].clone() | |
| player_videos = pipe._player_videos_from_latents( | |
| latents, {"tiled": True, "tile_size": TILE_SIZE, "tile_stride": TILE_STRIDE} | |
| ) | |
| pipe.load_models_to_device([]) | |
| return _as_player_videos(pipe._return_player_videos(player_videos)) | |
| def _player_file_tag(player_idx: int, num_players: int) -> str: | |
| if num_players == 2: | |
| return "playerA" if player_idx == 0 else "playerB" | |
| return f"player{player_idx + 1:02d}" | |
| def _build_player_rows_grid_video(rows, t_grid: int, border: int = 1): | |
| """Build a generic rows x P grid. | |
| Tensor/PIL convention at caller: | |
| rows: List[R][P][T] of PIL frames. | |
| Current N-player inference uses R=2: | |
| row 0 = generated players [P,T] | |
| row 1 = GT players [P,T] | |
| """ | |
| if not rows or not rows[0]: | |
| raise ValueError("_build_player_rows_grid_video: empty rows") | |
| num_rows = len(rows) | |
| num_players = len(rows[0]) | |
| for ridx, row in enumerate(rows): | |
| if len(row) != num_players: | |
| raise ValueError( | |
| f"_build_player_rows_grid_video: row {ridx} has {len(row)} players, expected {num_players}" | |
| ) | |
| for pidx, frames in enumerate(row): | |
| if len(frames) < t_grid: | |
| raise ValueError( | |
| f"_build_player_rows_grid_video: row={ridx} player={pidx} has {len(frames)} frames, expected >= {t_grid}" | |
| ) | |
| w0, h0 = rows[0][0][0].size | |
| total_w = num_players * w0 + (num_players - 1) * border | |
| total_h = num_rows * h0 + (num_rows - 1) * border | |
| grid_frames = [] | |
| for tidx in range(t_grid): | |
| canvas = Image.new("RGB", (total_w, total_h), color=(0, 0, 0)) | |
| for ridx, row in enumerate(rows): | |
| for pidx, frames in enumerate(row): | |
| x = pidx * (w0 + border) | |
| y = ridx * (h0 + border) | |
| canvas.paste(frames[tidx], (x, y)) | |
| grid_frames.append(canvas) | |
| return grid_frames | |
| def main(): | |
| _t_main = time.perf_counter() | |
| print("[timing] main() start", flush=True) | |
| with _timed("parse_args"): | |
| args = parse_args() | |
| if args.num_players < 2: | |
| raise ValueError(f"--num_players must be >= 2, got {args.num_players}") | |
| agent_vertex_ids = args.agent_vertex_ids | |
| if agent_vertex_ids is None: | |
| agent_vertex_ids = list(range(args.num_players)) | |
| elif len(agent_vertex_ids) != args.num_players: | |
| raise ValueError( | |
| f"--agent_vertex_ids length ({len(agent_vertex_ids)}) must match --num_players ({args.num_players})" | |
| ) | |
| args.loss_type = "diffusion_forcing" | |
| with _timed("read metadata csv"): | |
| meta_path = _resolve_metadata_path(args.dataset_path, args.metadata_file_name) | |
| metadata_table = pd.read_csv(meta_path) | |
| n_rows = len(metadata_table) | |
| if args.precomputed_consecutive_rows < 1: | |
| raise ValueError("--precomputed_consecutive_rows must be >= 1") | |
| precomputed_start_rows = None | |
| precomputed_start_rows = consecutive_precomputed_start_rows( | |
| metadata_table, args.precomputed_consecutive_rows, args.num_players | |
| ) | |
| available_rows = len(precomputed_start_rows) | |
| if available_rows < 1: | |
| raise ValueError( | |
| f"metadata has no valid window for precomputed_consecutive_rows={args.precomputed_consecutive_rows} (source rows={n_rows})" | |
| ) | |
| n_planned = ( | |
| available_rows | |
| if args.max_batches is None | |
| else min(available_rows, args.max_batches) | |
| ) | |
| sample_metadata = metadata_table.iloc[precomputed_start_rows].reset_index(drop=True) | |
| output_ids = [case_output_id(metadata_table, row) for row in precomputed_start_rows] | |
| sample_metadata["case_id"] = output_ids | |
| rank, world_size = _infer_rank_world() | |
| row_start, row_end = _infer_row_bounds(n_planned) | |
| eligible = select_case_rows(args, sample_metadata, list(range(row_start, row_end))) | |
| selected = set(eligible) | |
| skip_batches = set() | |
| if os.environ.get("FORCE_INFER", "0") != "1": | |
| existing = _list_mp4_basenames(args.output_dir, None) | |
| skip_batches = _existing_batch_indices(n_planned, args, existing, output_ids) | |
| remaining = n_planned - len(skip_batches) | |
| print( | |
| f"[skip] {len(skip_batches)}/{n_planned} batches already have mp4 (local); remaining={remaining}; FORCE_INFER=1 to overwrite", | |
| flush=True, | |
| ) | |
| if skip_batches and remaining == 0: | |
| print("[skip] all planned videos exist; nothing to generate", flush=True) | |
| return | |
| if skip_batches: | |
| print(f"[skip] done batches={sorted(skip_batches)}", flush=True) | |
| my_rows = [ | |
| i | |
| for i in _infer_shard_rows(n_planned, skip_batches, rank, world_size) | |
| if i in selected | |
| ] | |
| print( | |
| f"[shard] rank={rank}/{world_size} n_planned={n_planned} range=[{row_start},{row_end}) skip={len(skip_batches)} my_rows={my_rows}", | |
| flush=True, | |
| ) | |
| if os.environ.get("INFER_SHARD_DRY", "0") == "1": | |
| print("[shard] INFER_SHARD_DRY=1, exit before model load", flush=True) | |
| return | |
| if not my_rows: | |
| print("[shard] nothing left on this rank; skip model load", flush=True) | |
| return | |
| raw_inputs = is_raw_metadata(metadata_table) | |
| if raw_inputs: | |
| if args.precomputed_consecutive_rows != 1: | |
| raise ValueError("Raw-input CSVs require precomputed_consecutive_rows=1") | |
| with _timed("prepare actions, cameras, and grouping from tick logs"): | |
| meta_path = prepare_raw_inputs( | |
| metadata_table, | |
| args.dataset_path, | |
| my_rows, | |
| args.num_players, | |
| args.target_rgb_frames, | |
| args.prepared_inputs_dir, | |
| ) | |
| print(f"[infer] prepared raw inputs: {meta_path}", flush=True) | |
| with _timed("dataset init"): | |
| dataset = PrecomputedLatentDataset( | |
| args.dataset_path, | |
| meta_path, | |
| steps_per_epoch=available_rows, | |
| num_players=args.num_players, | |
| target_rgb_frames=args.target_rgb_frames, | |
| blank_prompt_emb_path=args.blank_prompt_emb_path, | |
| vae_latent_channels=args.vae_latent_channels, | |
| debug_no_retry=os.environ.get("INFER_SKIP_BAD_ROWS") != "1", | |
| use_camera_pose=args.use_camera_pose, | |
| use_group_labels=True, | |
| num_groups=args.num_groups, | |
| group_path_column=args.group_path_column, | |
| group_label_player_ids=args.group_label_player_ids, | |
| hier_group_mode=args.hier_group_mode, | |
| hier_group_seed=args.hier_group_seed, | |
| hier_random_preserve_sizes=bool(args.hier_random_preserve_sizes), | |
| hier_group_format=args.hier_group_format, | |
| num_frame_per_block=args.num_frame_per_block, | |
| consecutive_rows=args.precomputed_consecutive_rows, | |
| grouping_sample_indices=my_rows, | |
| grouping_cache_dir=args.grouping_cache_dir, | |
| grouping_model_profile=(args.num_local_hub_tokens, args.num_global_hub_tokens, | |
| args.hier_direct_frac, args.num_frame_per_block), | |
| allow_first_frame_latents=raw_inputs, | |
| ) | |
| local_rank = int(os.environ.get("LOCAL_RANK", "0")) | |
| if torch.cuda.is_available(): | |
| torch.cuda.set_device(local_rank) | |
| wan_root = args.wan_model_root.rstrip("/") | |
| print(f"[timing] wan_model_root={wan_root}", flush=True) | |
| print(f"[timing] ckpt_path={args.ckpt_path}", flush=True) | |
| _patch_model_manager_load_timing() | |
| with _timed("WanVideoNPlayersPipeline.from_pretrained (T5+VAE+tokenizer)"): | |
| pipe = WanVideoNPlayersPipeline.from_pretrained( | |
| torch_dtype=torch.bfloat16, | |
| device="cuda", | |
| tokenizer_config=ModelConfig( | |
| path=f"{wan_root}/Wan2.1-T2V-1.3B/google/umt5-xxl" | |
| ), | |
| model_configs=[ | |
| ModelConfig( | |
| path=f"{wan_root}/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth", | |
| offload_device="cpu", | |
| ), | |
| ModelConfig( | |
| path=f"{wan_root}/Wan2.2-TI2V-5B/Wan2.2_VAE.pth", | |
| offload_device="cpu", | |
| ), | |
| ], | |
| ) | |
| pipe.num_players = int(args.num_players) | |
| latent_t = (args.target_rgb_frames - 1) // 4 + 1 | |
| if latent_t % args.num_frame_per_block != 0: | |
| raise ValueError( | |
| f"latent_T={latent_t} must divide num_frame_per_block={args.num_frame_per_block}" | |
| ) | |
| if args.local_attn_size != -1 and args.local_attn_size < args.num_frame_per_block: | |
| raise ValueError( | |
| f"local_attn_size must be -1 or >= num_frame_per_block; got local={args.local_attn_size}, n={args.num_frame_per_block}" | |
| ) | |
| kg = args.num_global_hub_tokens | |
| if kg is None: | |
| kg = args.num_hub_tokens - args.num_groups * args.num_local_hub_tokens | |
| if args.num_groups * args.num_local_hub_tokens + kg != args.num_hub_tokens: | |
| raise ValueError( | |
| f"hier_hub requires num_groups*num_local_hub_tokens + num_global_hub_tokens == num_hub_tokens; got {args.num_groups}*{args.num_local_hub_tokens}+{kg} != {args.num_hub_tokens}" | |
| ) | |
| if args.hier_no_local_hub: | |
| if args.hier_local_hub_mode != "shared": | |
| raise ValueError( | |
| "--hier_no_local_hub requires --hier_local_hub_mode shared" | |
| ) | |
| if args.num_groups <= 0 or args.num_local_hub_tokens != 0 or kg <= 0: | |
| raise ValueError( | |
| f"--hier_no_local_hub requires num_groups>0, K_L=0, and K_G>0; got N={args.num_groups}, K_L={args.num_local_hub_tokens}, K_G={kg}" | |
| ) | |
| infer_kernel = os.environ.get("WAN_HIER_INFER_KERNEL", "sdpa").strip().lower() | |
| if infer_kernel not in ("sdpa", "packed_varlen", "paged_kv"): | |
| raise ValueError( | |
| "--hier_no_local_hub supports WAN_HIER_INFER_KERNEL=sdpa, packed_varlen, or paged_kv; got " | |
| + infer_kernel | |
| ) | |
| elif args.num_local_hub_tokens <= 0: | |
| raise ValueError("ordinary hier_hub requires num_local_hub_tokens>0") | |
| # Set pipeline runtime options before attaching the DiT so its parameters | |
| # are configured exactly once, before loading the complete checkpoint. | |
| pipe.enable_stage2(**stage2_options(args)) | |
| with _timed("construct Multics DiT without base weights"): | |
| pipe.dit = build_dit(args) | |
| state_dict, _ = _load_dit_checkpoint_state(args.ckpt_path) | |
| if not args.hier_no_local_hub: | |
| state_dict, migration = prepare_hier_local_hub_state_dict_for_load( | |
| state_dict, | |
| target_mode=args.hier_local_hub_mode, | |
| target_num_groups=pipe.dit.num_groups, | |
| target_num_local_hub_tokens=pipe.dit.num_local_hub_tokens, | |
| model_dim=pipe.dit.dim, | |
| ) | |
| if migration is not None: | |
| print( | |
| f"[hier-local-hub] inference migration: {migration['source_mode']} {migration['source_shape']} -> {migration['target_mode']} {migration['target_shape']}", | |
| flush=True, | |
| ) | |
| if args.hier_no_local_hub: | |
| if "local_hub_embed" in state_dict: | |
| raise RuntimeError( | |
| "No-local inference refuses to migrate an old hierarchical checkpoint. Warm-start it through training with --hier_no_local_hub, save a new checkpoint, then infer from that checkpoint." | |
| ) | |
| global_seed = state_dict.get("global_hub_embed") | |
| expected = tuple(pipe.dit.global_hub_embed.shape) | |
| if not torch.is_tensor(global_seed) or tuple(global_seed.shape) != expected: | |
| raise RuntimeError( | |
| f"No-local checkpoint has an incompatible global_hub_embed: expected {expected}, got {(None if global_seed is None else tuple(global_seed.shape))}" | |
| ) | |
| with _timed("load complete Multics DiT checkpoint"): | |
| load_complete_state_dict(pipe.dit, state_dict) | |
| del state_dict | |
| with _timed("pipe.to(cuda) + bfloat16"): | |
| pipe.to("cuda") | |
| pipe.to(dtype=torch.bfloat16) | |
| if raw_inputs: | |
| with _timed("encode first frames with the VAE"): | |
| encode_raw_first_frames(pipe, meta_path, my_rows, args.num_players) | |
| os.makedirs(args.output_dir, exist_ok=True) | |
| if args.precomputed_consecutive_rows > 1: | |
| manifest = { | |
| "metadata_file": meta_path, | |
| "target_rgb_frames": int(args.target_rgb_frames), | |
| "consecutive_rows": int(args.precomputed_consecutive_rows), | |
| "samples": [ | |
| { | |
| "sample_index": sample_index, | |
| "case_id": output_ids[sample_index], | |
| "output_index": output_ids[sample_index], | |
| "source_rows": list( | |
| range( | |
| int(precomputed_start_rows[sample_index]), | |
| int(precomputed_start_rows[sample_index]) | |
| + args.precomputed_consecutive_rows, | |
| ) | |
| ), | |
| } | |
| for sample_index in range(n_planned) | |
| ], | |
| } | |
| with open( | |
| os.path.join(args.output_dir, "consecutive_rows_manifest.json"), | |
| "w", | |
| encoding="utf-8", | |
| ) as file: | |
| json.dump(manifest, file, indent=2, ensure_ascii=False) | |
| print( | |
| f"[timing] model load total {time.perf_counter() - _t_main:.2f}s", | |
| flush=True, | |
| ) | |
| loader = torch.utils.data.DataLoader( | |
| torch.utils.data.Subset(dataset, my_rows), | |
| shuffle=False, | |
| batch_size=1, | |
| num_workers=args.dataloader_num_workers, | |
| ) | |
| tiler_kwargs = {"tiled": True, "tile_size": TILE_SIZE, "tile_stride": TILE_STRIDE} | |
| target_text = "" | |
| _t_first_batch = time.perf_counter() | |
| for local_i, batch in enumerate(loader): | |
| batch_idx = output_ids[my_rows[local_i]] | |
| if local_i == 0: | |
| src = "precomputed latent/action" | |
| print( | |
| f"[timing] first dataloader batch {time.perf_counter() - _t_first_batch:.2f}s ({src})", | |
| flush=True, | |
| ) | |
| video_latents = None | |
| gt_all = None | |
| keyboard = batch["keyboard"].squeeze(0)[:, :-1, :] | |
| mouse = batch["mouse"].squeeze(0)[:, :-1, :] | |
| video_latents = _precomputed_video_latents(batch).to( | |
| dtype=pipe.torch_dtype, device=pipe.device | |
| ) | |
| if video_latents.shape[0] != args.num_players: | |
| raise ValueError( | |
| f"batch video_latents P={video_latents.shape[0]} does not match --num_players={args.num_players}" | |
| ) | |
| # First-frame-only latents (raw inputs) have no ground-truth clip. | |
| if video_latents.shape[2] > 1: | |
| gt_all = _decode_latents_to_videos(pipe, video_latents, tiler_kwargs) | |
| video1 = gt_all[0:1] | |
| video2 = gt_all[1:2] | |
| else: | |
| video1 = video2 = None | |
| camera_pose = None | |
| if args.use_camera_pose: | |
| if "camera_pose" not in batch: | |
| raise KeyError( | |
| "--use_camera_pose 需要 batch 含 camera_pose（CSV camera_pose_path* / camera12_path*）" | |
| ) | |
| camera_pose = ( | |
| batch["camera_pose"] | |
| .squeeze(0) | |
| .to(dtype=pipe.torch_dtype, device=pipe.device) | |
| ) | |
| prompt = target_text | |
| negative_prompt = NEGATIVE_PROMPT | |
| has_caption_prompt = False | |
| if "prompt" in batch: | |
| prompt = batch["prompt"] | |
| if isinstance(prompt, (list, tuple)): | |
| prompt = [ | |
| p[0] if isinstance(p, (list, tuple)) and len(p) > 0 else p | |
| for p in prompt | |
| ] | |
| else: | |
| prompt = [prompt] | |
| prompt = [str(p).strip() for p in prompt] | |
| has_caption_prompt = any(prompt) | |
| if has_caption_prompt: | |
| negative_prompt = [ | |
| "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" | |
| ] * len(prompt) | |
| else: | |
| prompt = target_text | |
| causal_cfg_scale = args.cfg_scale if has_caption_prompt else 1.0 | |
| print( | |
| f"[infer] prompt={('caption' if has_caption_prompt else 'blank')} cfg_scale={causal_cfg_scale}", | |
| flush=True, | |
| ) | |
| if args.hier_infer_mode == "full_sequence": | |
| gen_players = _hier_full_sequence_inference( | |
| pipe, | |
| args, | |
| batch, | |
| video_latents, | |
| mouse, | |
| keyboard, | |
| camera_pose, | |
| agent_vertex_ids, | |
| ) | |
| else: | |
| if "group_labels" not in batch or "active_group" not in batch: | |
| raise KeyError( | |
| "hier_hub cached inference requires group_labels and active_group" | |
| ) | |
| gen_players = _as_player_videos( | |
| pipe.causal_inference( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| video1=video1, | |
| video2=video2, | |
| mouse=mouse, | |
| keyboard=keyboard, | |
| num_frames=args.causal_num_frames, | |
| height=TRAIN_HEIGHT, | |
| width=TRAIN_WIDTH, | |
| cfg_scale=causal_cfg_scale, | |
| num_inference_steps=args.num_inference_steps, | |
| seed=args.seed, | |
| tiled=True, | |
| tile_size=TILE_SIZE, | |
| tile_stride=TILE_STRIDE, | |
| loss_type="diffusion_forcing", | |
| return_both_players=True, | |
| agent_vertex_ids=agent_vertex_ids, | |
| video_latents=video_latents, | |
| camera_pose=camera_pose, | |
| overlap_path=None, | |
| group_labels=batch["group_labels"], | |
| active_group=batch["active_group"], | |
| membership=batch.get("membership"), | |
| global_only=batch.get("global_only"), | |
| ) | |
| ) | |
| if len(gen_players) != args.num_players: | |
| raise RuntimeError( | |
| f"Generated {len(gen_players)} player videos, expected P={args.num_players}" | |
| ) | |
| fps = args.grid_fps | |
| if not args.no_save_separate_players: | |
| for pidx, gen_p in enumerate(gen_players): | |
| save_video( | |
| gen_p, | |
| os.path.join( | |
| args.output_dir, | |
| f"video{batch_idx}_{_player_file_tag(pidx, args.num_players)}.mp4", | |
| ), | |
| fps=fps, | |
| quality=5, | |
| ) | |
| if not args.no_save_grid: | |
| T_gen = len(gen_players[0]) | |
| for pidx, gen_p in enumerate(gen_players): | |
| if len(gen_p) != T_gen: | |
| raise RuntimeError( | |
| f"Player outputs length mismatch: player0={T_gen} vs player{pidx}={len(gen_p)}" | |
| ) | |
| grid_rows = [gen_players] | |
| if gt_all is not None: | |
| grid_rows.append( | |
| [ | |
| tensor_bchw_to_pil_list(gt_all[pidx]) | |
| for pidx in range(args.num_players) | |
| ] | |
| ) | |
| T_grid = max(len(frames) for row in grid_rows for frames in row) | |
| grid_frames = _build_player_rows_grid_video( | |
| [ | |
| [pad_pil_list_to_length(frames, T_grid) for frames in row] | |
| for row in grid_rows | |
| ], | |
| T_grid, | |
| ) | |
| save_video( | |
| grid_frames, | |
| os.path.join(args.output_dir, f"video{batch_idx}_grid.mp4"), | |
| fps=fps, | |
| quality=5, | |
| ) | |