Spaces:

hugging-apps
/

multics-multiplayer-video

Running on Zero

Upload app.py with huggingface_hub

a1e0c09 verified 7 days ago

21.1 kB

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

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