"""Gradio demo for Multics (TencentARC/Multics): 4-player generative video. Multics is a hierarchical multi-player video world model built on the Wan2.2-TI2V-5B backbone. Given per-player first frames, text captions, and per-player action/camera tick logs, it renders one synchronous Counter-Strike-style clip seen from all four players at once. This demo ships the authors' five bundled example scenes (with their original tick logs) and also accepts four custom first frames with new captions. """ import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces # noqa: E402 (must be imported before torch) # Hierarchical cached-attention kernel. The reference config uses # flashinfer's paged_kv; packed_varlen implements the same hierarchical # packing on the native torch varlen flash-attention op (no extra # dependencies, Blackwell-safe). Set WAN_HIER_INFER_KERNEL=sdpa for a # pure-SDPA fallback. os.environ.setdefault("WAN_HIER_INFER_KERNEL", "packed_varlen") import csv # noqa: E402 import gc # noqa: E402 import tempfile # noqa: E402 import time # noqa: E402 from pathlib import Path # noqa: E402 import gradio as gr # noqa: E402 import numpy as np # noqa: E402 import torch # noqa: E402 from huggingface_hub import snapshot_download # noqa: E402 from PIL import Image # noqa: E402 from multics.checkpoint import ( # noqa: E402 build_dit, load_complete_state_dict, stage2_options, ) from multics.data.dataset import _camera12_from_absolute_c2w # noqa: E402 from multics.data.raw_inputs import ( # noqa: E402 KEYBOARD_DIM, clip_actions, clip_camera, encode_first_frame, load_first_frame, ) from multics.grouping.labels import generate_labels # noqa: E402 from multics.grouping.state import GroupingState # noqa: E402 from multics.models.hier_masks import ( # noqa: E402 prepare_hier_local_hub_state_dict_for_load, ) from multics.pipelines.wan_video_Nplayers import ( # noqa: E402 WanVideoNPlayersPipeline, ) from multics.utils import ModelConfig # noqa: E402 # ----------------------------------------------------------------- constants NUM_PLAYERS = 4 RUN_ID = "4players" TRAIN_HEIGHT = 352 TRAIN_WIDTH = 640 TARGET_RGB_FRAMES = 165 LATENT_T = (TARGET_RGB_FRAMES - 1) // 4 + 1 # 42 ACTION_FRAMES = TARGET_RGB_FRAMES - 1 # 164 frames drive motion KEYBOARD_DIMS = KEYBOARD_DIM # 23 GRID_FPS = 20 TILE_SIZE = (34, 34) TILE_STRIDE = (18, 16) NUM_GROUPS = 2 POS_SCALE_M = 32.0 MODEL_REPO = "TencentARC/Multics" # Reference negative prompt from the author's inference code. NEGATIVE_PROMPT = ( "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，" "JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，" "手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" ) EXAMPLES_DIR = Path(__file__).parent / "examples" / "4players" _WAN_PATTERNS = [ "Wan-AI/Wan2.1-T2V-1.3B/google/umt5-xxl/*", "Wan-AI/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth", "Wan-AI/Wan2.2-TI2V-5B/Wan2.2_VAE.pth", ] _CKP_PATTERN = f"multics/{RUN_ID}/checkpoints/model.ckpt" def _download_weights() -> dict: """Download T5 / VAE / tokenizer / DiT from TencentARC/Multics to /data.""" root = Path("/data/multics-weights") root.mkdir(parents=True, exist_ok=True) t0 = time.perf_counter() snapshot_download(MODEL_REPO, local_dir=str(root), allow_patterns=_WAN_PATTERNS) snapshot_download(MODEL_REPO, local_dir=str(root), allow_patterns=[_CKP_PATTERN]) print(f"[multics] weights ready in {time.perf_counter() - t0:.1f}s", flush=True) paths = { "t5": root / "Wan-AI/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth", "vae": root / "Wan-AI/Wan2.2-TI2V-5B/Wan2.2_VAE.pth", "tokenizer": root / "Wan-AI/Wan2.1-T2V-1.3B/google/umt5-xxl", "ckpt": root / "multics" / RUN_ID / "checkpoints" / "model.ckpt", } missing = [name for name, p in paths.items() if not p.exists()] if missing: raise FileNotFoundError(f"missing weight files: {missing}") return paths class _Args: """Reference inference config (configs/inference.yaml + hparams.yaml).""" num_players = NUM_PLAYERS simplex_pool_size = 10 simplex_gamma = 1.0 num_local_hub_tokens = 5 num_global_hub_tokens = 6 num_hub_tokens = 16 hier_direct_schedule = "uniform" hier_direct_frac = 1.0 hier_direct_layers = None hier_group_format = "membership" use_camera_pose = True camera_dim = 12 camera_dropout_p = 0.1 target_rgb_frames = TARGET_RGB_FRAMES num_frame_per_block = 3 local_attn_size = 21 attention_mode = "hier_hub" num_groups = NUM_GROUPS hier_group_mode = "csv" hier_local_hub_mode = "shared" hier_no_local_hub = False # sdpa keeps the DiT off torch.compile paths; the reference "flex" # block-mask is bypassed during cached inference anyway. attn_impl = "sdpa" def _load_pipeline() -> WanVideoNPlayersPipeline: """Build the pipeline and DiT at module scope (ZeroGPU packs weights here).""" t0 = time.perf_counter() paths = _download_weights() print(f"[multics] loading pipeline, t={time.perf_counter() - t0:.1f}s", flush=True) pipe = WanVideoNPlayersPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", tokenizer_config=ModelConfig(path=str(paths["tokenizer"])), model_configs=[ ModelConfig(path=str(paths["t5"]), offload_device="cpu"), ModelConfig(path=str(paths["vae"]), offload_device="cpu"), ], ) pipe.num_players = NUM_PLAYERS args = _Args() pipe.enable_stage2(**stage2_options(args)) print(f"[multics] building DiT, t={time.perf_counter() - t0:.1f}s", flush=True) pipe.dit = build_dit(args) print(f"[multics] loading ckpt, t={time.perf_counter() - t0:.1f}s", flush=True) state_dict = torch.load( str(paths["ckpt"]), map_location="cpu", weights_only=True, mmap=True ) if isinstance(state_dict, dict) and isinstance( state_dict.get("state_dict"), dict ): state_dict = state_dict["state_dict"] if any(k.startswith("pipe.dit.") for k in state_dict): state_dict = { k[len("pipe.dit."):]: v for k, v in state_dict.items() if k.startswith("pipe.dit.") } print(f"[multics] migrating local hub, t={time.perf_counter() - t0:.1f}s", flush=True) state_dict, migration = prepare_hier_local_hub_state_dict_for_load( state_dict, target_mode=pipe.dit.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, ) print(f"[multics] strict-loading DiT, t={time.perf_counter() - t0:.1f}s", flush=True) if migration is not None: print(f"[multics] local hub migration: {migration}", flush=True) load_complete_state_dict(pipe.dit, state_dict) del state_dict gc.collect() print(f"[multics] moving to cuda, t={time.perf_counter() - t0:.1f}s", flush=True) pipe.to("cuda") pipe.to(dtype=torch.bfloat16) print(f"[multics] pipeline ready, t={time.perf_counter() - t0:.1f}s", flush=True) return pipe _pipe = _load_pipeline() # --------------------------------------------------------------- input prep def _case_metadata() -> dict: """Map case_id -> {captions, images, ticks} for the bundled scenes.""" meta = {} with open(EXAMPLES_DIR / "inputs.csv") as f: for row in csv.DictReader(f): case_dir = f"4_{int(row['case_id']):02d}" meta[case_dir] = { "captions": [ (row.get(f"caption{i}") or "").strip() for i in range(1, NUM_PLAYERS + 1) ], "images": [ EXAMPLES_DIR / case_dir / f"player={p:02d}" / "first_frame.png" for p in range(NUM_PLAYERS) ], "ticks": [ EXAMPLES_DIR / case_dir / f"player={p:02d}" / "ticks.parquet" for p in range(NUM_PLAYERS) ], } return meta CASE_META = _case_metadata() CASE_DIRS = sorted(CASE_META.keys()) def _prepare_case(case: str) -> dict: """Assemble actions / cameras / grouping for one bundled scene. The bundled tick logs are pre-cropped to 48 frames of history plus the 165-frame clip, so the log-relative clip start is derived from the grouping state's round-tick timeline rather than copied from ``clip_start_frame`` (which indexes the full round). """ info = CASE_META[case] state = GroupingState.load(str(EXAMPLES_DIR / case / "grouping_state.npz")) payload = generate_labels(state) membership = torch.from_numpy(np.ascontiguousarray(payload["membership"])).to(torch.bool) global_only = torch.from_numpy(np.ascontiguousarray(payload["global_only"])).to(torch.bool) start = int( (state.round_tick_start + state.clip_start_frame * 2 - int(state.ticks[0])) // 2 ) keyboards, mice, cams = [], [], [] for p in range(NUM_PLAYERS): path = str(info["ticks"][p]) act = clip_actions(path, start, TARGET_RGB_FRAMES) keyboards.append(act["keyboard"]) mice.append(act["mouse"]) cam = clip_camera(path, start, TARGET_RGB_FRAMES, p) cams.append(cam["c2w_abs"].float()) keyboard = torch.stack(keyboards)[:, :-1, :] # [P, 164, 23] mouse = torch.stack(mice)[:, :-1, :] # [P, 164, 2] camera12 = _camera12_from_absolute_c2w(torch.stack(cams), POS_SCALE_M) return { "keyboard": keyboard, "mouse": mouse, "camera12": camera12, "membership": membership[:LATENT_T], "global_only": global_only[:LATENT_T], "captions": info["captions"], "images": [str(p) for p in info["images"]], } def _prepare_custom(image_paths, captions) -> dict: """Static conditioning for custom uploads: no actions, fixed 2+2 grouping. The 2+2 topology keeps the communication budget well inside the cost limit for 4 players, matches how the bundled scenes are grouped, and needs no tick logs. """ keyboard = torch.zeros(NUM_PLAYERS, ACTION_FRAMES, KEYBOARD_DIMS) mouse = torch.zeros(NUM_PLAYERS, ACTION_FRAMES, 2) c2w = torch.eye(4).unsqueeze(0).unsqueeze(0).repeat(NUM_PLAYERS, LATENT_T, 1, 1) camera12 = _camera12_from_absolute_c2w(c2w, POS_SCALE_M) membership = torch.zeros(LATENT_T, NUM_GROUPS, NUM_PLAYERS, dtype=torch.bool) membership[:, 0, 0] = True membership[:, 0, 1] = True membership[:, 1, 2] = True membership[:, 1, 3] = True global_only = torch.zeros(LATENT_T, NUM_PLAYERS, dtype=torch.bool) return { "keyboard": keyboard, "mouse": mouse, "camera12": camera12, "membership": membership, "global_only": global_only, "captions": [str(c) for c in captions], "images": [str(p) for p in image_paths], } def _encode_frames(image_paths) -> torch.Tensor: """VAE-encode the four first frames -> [P, 48, 1, h, w] on the CPU.""" latents = [] for path in image_paths: frame = load_first_frame(path, TRAIN_HEIGHT, TRAIN_WIDTH) latents.append(encode_first_frame(_pipe, frame)) return torch.stack(latents) def _save_grid_mp4(player_frames, out_path: str) -> None: """Tile the per-player clips into one N-across grid video at 20 fps. Follows the authors' ``_build_player_rows_grid_video`` convention (multics/inference.py): players are laid out left-to-right in a single row with a 1px border, each cell is the full player frame size. """ import imageio # PIL .size is (width, height) -- keep that order. w0, h0 = player_frames[0][0].size border = 1 grid_w = NUM_PLAYERS * w0 + (NUM_PLAYERS - 1) * border grid_h = h0 # imageio/ffmpeg silently rescales anything not divisible by 16; pad the # canvas ourselves so the geometry we build is the geometry we ship. pad_w = -grid_w % 16 pad_h = -grid_h % 16 canvas_size = (grid_w + pad_w, grid_h + pad_h) writer = imageio.get_writer(out_path, fps=GRID_FPS, quality=5, macro_block_size=1) try: for t in range(len(player_frames[0])): canvas = Image.new("RGB", canvas_size, color=(0, 0, 0)) for p, frames in enumerate(player_frames): canvas.paste(frames[t], (p * (w0 + border), 0)) writer.append_data(np.asarray(canvas)) finally: writer.close() def _estimate_duration(*args, **kwargs) -> int: """ZeroGPU reservation: full-quality runs need several GPU-minutes. 50-step runs cost ~14 latent blocks x (50 steps x 2 CFG passes + 2 context recaches) of DiT forwards; fewer steps scale roughly linearly. Gradio passes handler inputs positionally, so ``num_inference_steps`` is either a keyword or the 12th positional argument (after case, use_bundled, 4 images, 4 captions, and seed). """ raw = kwargs.get("num_inference_steps") if raw is None and len(args) >= 12: raw = args[11] try: steps = int(raw) except (TypeError, ValueError): steps = 50 return 600 if steps >= 50 else 420 @spaces.GPU(duration=_estimate_duration, size="xlarge") def generate( case="4_01", use_bundled=True, image1=None, image2=None, image3=None, image4=None, caption1="", caption2="", caption3="", caption4="", seed=0, num_inference_steps=50, cfg_scale=5.0, progress=gr.Progress(track_tqdm=True), ): """Generate a 165-frame 4-player Counter-Strike-style video grid. Args: case: bundled scene id, e.g. "4_01" (used when use_bundled is True). use_bundled: True -> replay the bundled scene's first frames plus its real action / camera tick logs and grouping; False -> use the four uploaded first frames and the captions below (no actions, static cameras, fixed 2+2 grouping). image1..image4: per-player custom first frames (352x640 auto-fit). caption1..caption4: per-player text prompts; in bundled mode non-empty captions override the scene's own captions. seed: RNG seed. num_inference_steps: denoising steps (reference: 50). cfg_scale: classifier-free guidance scale (reference: 5.0). Returns: Path to the generated 4-up grid video (.mp4, ~8 s at 20 fps). """ t0 = time.perf_counter() pipe = _pipe if case not in CASE_DIRS: raise gr.Error(f"Unknown bundled scene {case!r}") if use_bundled: prep = _prepare_case(case) captions = list(prep["captions"]) overrides = [caption1, caption2, caption3, caption4] captions = [ override if override.strip() else captions[i] for i, override in enumerate(overrides) ] else: images = [image1, image2, image3, image4] if any(img is None for img in images): raise gr.Error("Please provide first frames for all four players.") captions = [caption1, caption2, caption3, caption4] if not any(c.strip() for c in captions): raise gr.Error("Please provide at least one caption (prompt).") prep = _prepare_custom(images, captions) latents = _encode_frames(prep["images"]) print(f"[multics] input prep done in {time.perf_counter() - t0:.1f}s", flush=True) videos = pipe.causal_inference( prompt=list(captions), negative_prompt=[NEGATIVE_PROMPT] * NUM_PLAYERS, video1=None, video2=None, mouse=prep["mouse"].to(dtype=pipe.torch_dtype, device="cuda"), keyboard=prep["keyboard"].to(dtype=pipe.torch_dtype, device="cuda"), num_frames=TARGET_RGB_FRAMES, height=TRAIN_HEIGHT, width=TRAIN_WIDTH, cfg_scale=float(cfg_scale), num_inference_steps=int(num_inference_steps), seed=int(seed), tiled=True, tile_size=TILE_SIZE, tile_stride=TILE_STRIDE, loss_type="diffusion_forcing", return_both_players=True, agent_vertex_ids=list(range(NUM_PLAYERS)), video_latents=latents.to(dtype=pipe.torch_dtype, device="cuda"), camera_pose=prep["camera12"].to(dtype=pipe.torch_dtype, device="cuda"), group_labels=None, active_group=None, membership=prep["membership"].to(device="cuda"), global_only=prep["global_only"].to(device="cuda"), ) if isinstance(videos, tuple): videos = list(videos) out_path = str( Path(tempfile.gettempdir()) / f"multics_grid_{int(time.time())}_{os.getpid()}.mp4" ) _save_grid_mp4(videos, out_path) print(f"[multics] generate done in {time.perf_counter() - t0:.1f}s", flush=True) return out_path # ----------------------------------------------------------------------- UI CASE_LABELS = { "4_01": "4_01 · Ancient ruins firefight", "4_02": "4_02 · Colorful knife alley", "4_03": "4_03 · banana throw in the alleys", "4_04": "4_04 · Future-gun temple fight", "4_05": "4_05 · Smoke push", } with gr.Blocks() as demo: gr.Markdown( """ # 🎮 Multics · Generative Multi-Player Video World Model **Multics** (TencentARC) renders one synchronous 4-player Counter-Strike-style clip from per-player first frames, captions, and real in-game action/camera tick logs (Wan2.2-TI2V-5B backbone). [Model card](https://huggingface.co/TencentARC/Multics) · [Code](https://github.com/TencentARC/Multics) > ⚠️ **Heavy job** — one 165-frame (~8 s) 4-player clip takes several minutes of GPU time. > Not playable on the free ZeroGPU daily quota; a PRO account is needed. > > 🎬 The five bundled scenes (frames + tick logs) come from the > [Multics repo](https://github.com/TencentARC/Multics) (© Multics authors; demo content > derived from the OpenCS2 dataset). """ ) with gr.Row(): with gr.Column(scale=1): use_bundled = gr.Checkbox( label="Use a bundled scene (real tick logs: actions, cameras, grouping)", value=True, ) case_dd = gr.Dropdown( choices=[(CASE_LABELS.get(c, c), c) for c in CASE_DIRS], value="4_01", label="Bundled scene", ) caption1 = gr.Textbox(label="Caption · Player 1", lines=2) caption2 = gr.Textbox(label="Caption · Player 2", lines=2) caption3 = gr.Textbox(label="Caption · Player 3", lines=2) caption4 = gr.Textbox(label="Caption · Player 4", lines=2) with gr.Accordion( "Custom first frames (used when the checkbox above is off)", open=False, ): image1 = gr.Image( label="Player 1 first frame", type="filepath", height=140 ) image2 = gr.Image( label="Player 2 first frame", type="filepath", height=140 ) image3 = gr.Image( label="Player 3 first frame", type="filepath", height=140 ) image4 = gr.Image( label="Player 4 first frame", type="filepath", height=140 ) with gr.Accordion("Advanced", open=False): seed = gr.Slider(0, 2**31 - 1, value=0, step=1, label="Seed") num_inference_steps = gr.Slider( 20, 50, value=50, step=1, label="Inference steps" ) cfg_scale = gr.Slider( 1.0, 10.0, value=5.0, step=0.1, label="Guidance scale (CFG)" ) run = gr.Button("Generate 4-player video", variant="primary") with gr.Column(scale=2): video_out = gr.Video(label="4-player grid video") run.click( generate, inputs=[ case_dd, use_bundled, image1, image2, image3, image4, caption1, caption2, caption3, caption4, seed, num_inference_steps, cfg_scale, ], outputs=[video_out], api_name="generate", ) gr.Examples( examples=[[case] for case in CASE_DIRS], inputs=[case_dd], fn=generate, outputs=[video_out], cache_examples=True, cache_mode="lazy", label="Bundled scenes (click a row to auto-fill and run)", ) demo.launch(mcp_server=True, theme=gr.themes.Citrus())