Spaces:

ForgeWM
/

forgewm-world-model

Running on Zero

Size the ZeroGPU reservation from measured runtimes; expose the example path on the API

fee621f verified 3 months ago

20.3 kB

"""ForgeWM — few-step, action-conditioned Minecraft world model.
Gradio demo for `ForgeWM/ForgeWM` (paper: "ForgeWM: Progressive Causal Training
for Few-Step Action-Conditioned Video World Models").
Faithful to the repo's own `inference.py` / `pipeline/causal_inference.py`:
same 352x640 resolution, 3-latent-frame causal blocks, sliding window of 6
latent frames, the released `warp_denoising_step` schedules, and First-Frame
Enhancement for the 1-/2-step students.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import spaces  # noqa: E402  (must precede torch)

import gc  # noqa: E402
import shutil  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402
from typing import List, Tuple  # noqa: E402

import gradio as gr  # noqa: E402
import imageio.v2 as imageio  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
import torch.nn.functional as F  # noqa: E402
from huggingface_hub import hf_hub_download  # noqa: E402
from omegaconf import OmegaConf  # noqa: E402
from PIL import Image  # noqa: E402

from pipeline import CausalInferencePipeline
from utils.wan_wrapper import WanVAEWrapper

# ────────────────────────────── constants ──────────────────────────────
HEIGHT, WIDTH = 352, 640
FPS = 12
VAE_TCR = 4                 # VAE temporal compression ratio
FRAMES_PER_BLOCK = 3        # latent frames per causal block
CAM_VALUE = 0.10            # camera delta magnitude, from inference.py

MINECRAFT_ACTIONS = [
    "forward", "back", "left", "right",
    "turn_right", "turn_left", "look_up", "look_down",
    "forward_turn_right", "random", "no_action",
]

BASE_REPO = "Skywork/Matrix-Game-2.0"
FORGEWM_REPO = "ForgeWM/ForgeWM"

VARIANTS = {
    "ForgeWM-4 · 4 steps": ("configs/stage3_dmd.yaml", "stage3/model.pt"),
    "ForgeWM-2 · 2 steps": ("configs/stage3_dmd_2step.yaml", "2step/model.pt"),
    "ForgeWM-1 · 1 step": ("configs/stage3_dmd_1step.yaml", "1step/model.pt"),
}
DEFAULT_VARIANT = "ForgeWM-4 · 4 steps"

# Measured on the ZeroGPU Blackwell card (12-block rollouts): 1.04 s/block for
# the 4-step student, 0.40 s/block for the 1-step one, plus ~4 s of CLIP + VAE
# encode/decode and mp4 muxing.  Used only to size the @spaces.GPU reservation.
SECONDS_PER_CHUNK = {
    "ForgeWM-4 · 4 steps": 1.20,
    "ForgeWM-2 · 2 steps": 0.75,
    "ForgeWM-1 · 1 step": 0.50,
}
FIXED_OVERHEAD = 13.0       # CLIP + VAE + mux + first-call weight streaming

CKPT_DIR = os.path.join(os.getcwd(), "ckpts", "MG2-base")


# ─────────────────────────── weight preparation ───────────────────────────
def _link(src: str, dst: str) -> None:
    if os.path.lexists(dst):
        os.remove(dst)
    os.symlink(src, dst)


def _purge(path: str) -> None:
    """Drop a hub blob from local disk once its tensors are in memory."""
    try:
        real = os.path.realpath(path)
        if os.path.isfile(real):
            os.remove(real)
        if os.path.islink(path):
            os.remove(path)
    except OSError as exc:  # pragma: no cover
        print(f"[disk] could not purge {path}: {exc}")


def _disk() -> str:
    total, used, free = shutil.disk_usage("/")
    return f"disk {used / 2**30:.1f}G used / {free / 2**30:.1f}G free"


os.makedirs(os.path.join(CKPT_DIR, "xlm-roberta-large"), exist_ok=True)

print(f"[setup] fetching Matrix-Game-2.0 base weights … ({_disk()})", flush=True)
_dit = hf_hub_download(BASE_REPO, "base_model/diffusion_pytorch_model.safetensors")
_link(_dit, os.path.join(CKPT_DIR, "diffusion_pytorch_model.safetensors"))
_link(hf_hub_download(BASE_REPO, "base_model/base_config.json"),
      os.path.join(CKPT_DIR, "base_config.json"))
_vae_path = hf_hub_download(BASE_REPO, "Wan2.1_VAE.pth")
_clip_path = hf_hub_download(
    BASE_REPO, "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth")
for _tok in ("sentencepiece.bpe.model", "special_tokens_map.json",
             "tokenizer.json", "tokenizer_config.json"):
    _link(hf_hub_download(BASE_REPO, f"xlm-roberta-large/{_tok}"),
          os.path.join(CKPT_DIR, "xlm-roberta-large", _tok))
print(f"[setup] base weights ready ({_disk()})", flush=True)

# ─────────────────────────────── models ───────────────────────────────
torch.set_grad_enabled(False)
DTYPE = torch.bfloat16

print("[setup] building VAE + CLIP …", flush=True)
VAE = WanVAEWrapper(
    vae_path=_vae_path,
    clip_checkpoint_path=_clip_path,
    clip_tokenizer_path=os.path.join(CKPT_DIR, "xlm-roberta-large"),
).eval()
VAE = VAE.to(device="cuda", dtype=DTYPE)
# `WanVAEWrapper.clip` is a plain object, not a submodule, so `.to()` above does
# not reach it.  Move it explicitly (in fp32, as the repo does) so ZeroGPU packs
# it with everything else instead of paying a 4.8 GB host copy per request.
VAE.clip.model = VAE.clip.model.to("cuda")
gc.collect()
print(f"[setup] VAE + CLIP on GPU ({_disk()})", flush=True)


def _build_variant(name: str) -> CausalInferencePipeline:
    cfg_path, ckpt_file = VARIANTS[name]
    config = OmegaConf.merge(OmegaConf.load("configs/default.yaml"),
                             OmegaConf.load(cfg_path))
    config.model_kwargs.model_name = CKPT_DIR

    pipe = CausalInferencePipeline(config, device=torch.device("cuda"), vae=VAE)

    ckpt = hf_hub_download(FORGEWM_REPO, ckpt_file)
    try:
        state = torch.load(ckpt, map_location="cpu", weights_only=True)
    except Exception:
        state = torch.load(ckpt, map_location="cpu", weights_only=False)
    gen_sd = state.get("generator", state.get("generator_ema", state)) \
        if isinstance(state, dict) else state
    fixed = {k.replace("._fsdp_wrapped_module.", ".")
              .replace("._checkpoint_wrapped_module.", "."): v
             for k, v in gen_sd.items()}
    missing, unexpected = pipe.generator.load_state_dict(fixed, strict=False)
    print(f"[{name}] loaded: missing={len(missing)} unexpected={len(unexpected)}",
          flush=True)
    del state, gen_sd, fixed
    gc.collect()
    _purge(ckpt)

    pipe.generator = pipe.generator.to(device="cuda", dtype=DTYPE).eval()
    pipe.vae = VAE
    gc.collect()
    print(f"[{name}] ready ({_disk()})", flush=True)
    return pipe


PIPELINES = {}
for _name in VARIANTS:
    try:
        PIPELINES[_name] = _build_variant(_name)
    except Exception as exc:  # keep the Space usable if one student fails
        print(f"[setup] FAILED to load {_name}: {exc!r}", flush=True)

if not PIPELINES:
    raise RuntimeError("No ForgeWM student could be loaded.")
if DEFAULT_VARIANT not in PIPELINES:
    DEFAULT_VARIANT = next(iter(PIPELINES))

# The base DiT safetensors were only needed to instantiate the architecture;
# every parameter has since been overwritten by the ForgeWM checkpoints.
_purge(os.path.join(CKPT_DIR, "diffusion_pytorch_model.safetensors"))
print(f"[setup] {len(PIPELINES)} variant(s) ready ({_disk()})", flush=True)


# ──────────────────────────── action scripting ────────────────────────────
def make_action(action_type: str, num_raw_frames: int) -> Tuple[torch.Tensor, torch.Tensor]:
    """Minecraft action palette (mouse_dim_in=2), verbatim from inference.py."""
    mouse = torch.zeros(1, num_raw_frames, 2)
    keyboard = torch.zeros(1, num_raw_frames, 6)

    if action_type == "forward":
        keyboard[:, :, 0] = 1.0
    elif action_type == "back":
        keyboard[:, :, 1] = 1.0
    elif action_type == "left":
        keyboard[:, :, 2] = 1.0
    elif action_type == "right":
        keyboard[:, :, 3] = 1.0
    elif action_type == "turn_right":
        mouse[:, :, 1] = CAM_VALUE
    elif action_type == "turn_left":
        mouse[:, :, 1] = -CAM_VALUE
    elif action_type == "look_up":
        mouse[:, :, 0] = CAM_VALUE
    elif action_type == "look_down":
        mouse[:, :, 0] = -CAM_VALUE
    elif action_type == "forward_turn_right":
        keyboard[:, :, 0] = 1.0
        mouse[:, :, 1] = CAM_VALUE
    elif action_type == "random":
        torch.manual_seed(42)
        mouse = (torch.rand(1, num_raw_frames, 2) - 0.5) * (2 * CAM_VALUE)
        keyboard[:, :, :4] = (torch.rand(1, num_raw_frames, 4) > 0.5).float()
    elif action_type == "no_action":
        pass
    else:
        raise gr.Error(f"Unknown action '{action_type}'.")
    return mouse, keyboard


def _chunk_start(chunk_index: int) -> int:
    """First raw (pixel) frame owned by causal block `chunk_index`.
    Latent frame 0 maps to raw frame 0; latent frame f>=1 maps to raw frames
    4f-3 … 4f.  A block spans 3 latent frames, so block c covers raw frames
    [12c-3, 12c+9) — and [0, 9) for c == 0.
    """
    return 0 if chunk_index == 0 else 12 * chunk_index - 3


def build_action_track(segments: List[Tuple[str, int]]) -> Tuple[torch.Tensor, torch.Tensor]:
    total_chunks = sum(n for _, n in segments)
    num_raw = _chunk_start(total_chunks)
    mouse = torch.zeros(1, num_raw, 2)
    keyboard = torch.zeros(1, num_raw, 6)
    chunk = 0
    for action, count in segments:
        for _ in range(count):
            lo, hi = _chunk_start(chunk), _chunk_start(chunk + 1)
            m, k = make_action(action, hi - lo)
            mouse[:, lo:hi] = m
            keyboard[:, lo:hi] = k
            chunk += 1
    return mouse, keyboard


# ───────────────────────────── preprocessing ─────────────────────────────
def load_reference_frame(image_path: str) -> torch.Tensor:
    """Aspect-preserving resize + centre crop to 352x640, scaled to [-1, 1]."""
    image = Image.open(image_path).convert("RGB")
    arr = torch.from_numpy(np.asarray(image)).permute(2, 0, 1)[None].float() / 255.0
    _, _, h, w = arr.shape
    if h / w > HEIGHT / WIDTH:
        new_w, new_h = WIDTH, max(HEIGHT, int(round(h * WIDTH / w)))
    else:
        new_h, new_w = HEIGHT, max(WIDTH, int(round(w * HEIGHT / h)))
    arr = F.interpolate(arr, size=(new_h, new_w), mode="bilinear", align_corners=False)
    top, left = (new_h - HEIGHT) // 2, (new_w - WIDTH) // 2
    arr = arr[:, :, top:top + HEIGHT, left:left + WIDTH]
    arr = (arr - 0.5) / 0.5
    return arr.unsqueeze(0).to(device="cuda", dtype=DTYPE)  # [1, 1, 3, H, W]


def build_conditional_dict(pipe, pixel, num_frames, mouse_cond, keyboard_cond):
    """MG2-style conditioning: CLIP context + first-frame latent + mask."""
    num_pixel_frames = (num_frames - 1) * VAE_TCR + 1
    visual_context = pipe.vae.encode_visual_context_from_pixels(pixel).to(DTYPE)

    first_frame = pixel[:, 0:1]
    pad = torch.zeros(1, num_pixel_frames - 1, 3, pixel.shape[3], pixel.shape[4],
                      device=pixel.device, dtype=DTYPE)
    padded = torch.cat([first_frame, pad], dim=1).permute(0, 2, 1, 3, 4)
    img_cond = pipe.vae.encode_to_latent(padded).to(DTYPE)

    _, _, _, h_lat, w_lat = img_cond.shape
    mask = torch.zeros(1, num_frames, 4, h_lat, w_lat,
                       device=pixel.device, dtype=DTYPE)
    mask[:, 0:1] = 1
    return {
        "visual_context": visual_context,
        "cond_concat": torch.cat([mask, img_cond], dim=2),
        "mouse_condition": mouse_cond.to(device="cuda", dtype=DTYPE),
        "keyboard_condition": keyboard_cond.to(device="cuda", dtype=DTYPE),
    }


# ──────────────────────────────── core ────────────────────────────────
@torch.no_grad()
def _rollout(image_path, variant, segments, seed):
    if image_path is None:
        raise gr.Error("Please provide a reference frame.")
    if variant not in PIPELINES:
        variant = DEFAULT_VARIANT
    segments = [(a, int(n)) for a, n in segments if int(n) > 0]
    total_chunks = sum(n for _, n in segments)
    if total_chunks < 1:
        raise gr.Error("Give at least one action segment a non-zero length.")
    if total_chunks > 12:
        raise gr.Error("Keep the total rollout at 12 chunks or fewer.")

    pipe = PIPELINES[variant]
    num_frames = total_chunks * FRAMES_PER_BLOCK

    pixel = load_reference_frame(image_path)
    mouse, keyboard = build_action_track(segments)
    cond = build_conditional_dict(pipe, pixel, num_frames, mouse, keyboard)

    torch.manual_seed(int(seed))
    noise = torch.randn([1, num_frames, 16, HEIGHT // 8, WIDTH // 8],
                        device="cuda", dtype=DTYPE)

    start = time.time()
    video = pipe.inference(noise=noise, conditional_dict=cond, return_latents=False)
    torch.cuda.synchronize()
    elapsed = time.time() - start
    pipe.vae.model.clear_cache()

    frames = (video[0].permute(0, 2, 3, 1).float().cpu().numpy() * 255)
    frames = frames.clip(0, 255).astype(np.uint8)
    out_path = os.path.join(tempfile.mkdtemp(), "forgewm.mp4")
    writer = imageio.get_writer(out_path, fps=FPS, codec="libx264",
                                quality=8, macro_block_size=None)
    for frame in frames:
        writer.append_data(frame)
    writer.close()

    script = " → ".join(f"`{a}` ×{n}" for a, n in segments)
    info = (
        f"**{variant}** · {total_chunks} causal blocks · {num_frames} latent "
        f"frames → {len(frames)} pixel frames ({len(frames) / FPS:.1f}s @ {FPS} fps)\n\n"
        f"Action script: {script}\n\n"
        f"Rollout + VAE decode: **{elapsed:.2f}s** "
        f"({1000 * elapsed / total_chunks:.0f} ms per block, decode included) · seed `{int(seed)}`"
    )
    return out_path, info


def _estimate(variant, chunks) -> int:
    per = SECONDS_PER_CHUNK.get(variant, 1.2)
    return int(FIXED_OVERHEAD + per * max(1, min(int(chunks), 12)))


def _duration_main(*args):
    variant = args[1] if len(args) > 1 else DEFAULT_VARIANT
    chunks = sum(int(args[i]) for i in (3, 5, 7) if len(args) > i)
    return _estimate(variant, chunks)


def _duration_example(*args):
    chunks = sum(int(args[i]) for i in (2, 4, 6) if len(args) > i)
    return _estimate(DEFAULT_VARIANT, chunks)


@spaces.GPU(duration=_duration_main)
def generate(
    image: str,
    variant: str,
    action_1: str,
    chunks_1: int,
    action_2: str,
    chunks_2: int,
    action_3: str,
    chunks_3: int,
    seed: int = 0,
    progress=gr.Progress(track_tqdm=True),
):
    """Roll out an action-conditioned Minecraft video from one reference frame.
    Args:
        image: Path to the reference frame that anchors the world.
        variant: Which few-step ForgeWM student to sample with.
        action_1: Action held during the first segment of the rollout.
        chunks_1: Length of the first segment, in 3-latent-frame blocks (~1s each).
        action_2: Action held during the second segment.
        chunks_2: Length of the second segment, in blocks.
        action_3: Action held during the third segment.
        chunks_3: Length of the third segment, in blocks.
        seed: Random seed for the initial noise.
    Returns:
        An mp4 of the generated rollout and a markdown summary of the run.
    """
    return _rollout(
        image, variant,
        [(action_1, chunks_1), (action_2, chunks_2), (action_3, chunks_3)],
        seed,
    )


@spaces.GPU(duration=_duration_example)
def generate_example(
    image: str,
    action_1: str,
    chunks_1: int,
    action_2: str,
    chunks_2: int,
    action_3: str,
    chunks_3: int,
    progress=gr.Progress(track_tqdm=True),
):
    """Run a bundled example with the default student and seed 0."""
    return _rollout(
        image, DEFAULT_VARIANT,
        [(action_1, chunks_1), (action_2, chunks_2), (action_3, chunks_3)],
        0,
    )


# ───────────────────────────────── UI ─────────────────────────────────
CSS = """
.gradio-container { max-width: 1200px !important; }
"""

with gr.Blocks(title="ForgeWM") as demo:
    gr.Markdown(
        "# 🎮 ForgeWM — few-step action-conditioned world model\n"
        "Give it **one Minecraft frame** and a short **action script**; the "
        "block-causal diffusion transformer rolls the world forward at "
        "**1, 2 or 4 denoising steps** per 3-frame block.\n\n"
        "[Paper](https://huggingface.co/papers/2608.14022) · "
        "[Model](https://huggingface.co/ForgeWM/ForgeWM) · "
        "[Code](https://github.com/asdfo123/ForgeWM) · "
        "[Project page](https://asdfo123.github.io/ForgeWM/)"
    )

    with gr.Row():
        with gr.Column(scale=1):
            image = gr.Image(label="Reference frame", type="filepath",
                             height=280, sources=["upload", "clipboard"])
            variant = gr.Radio(
                choices=list(PIPELINES.keys()),
                value=DEFAULT_VARIANT,
                label="Student",
                info="Fewer steps = faster. 1-/2-step use First-Frame Enhancement.",
            )
            gr.Markdown("### Action script  \nEach block is 12 frames ≈ 1 second.")
            with gr.Row():
                action_1 = gr.Dropdown(MINECRAFT_ACTIONS, value="forward",
                                       label="Segment 1", scale=2)
                chunks_1 = gr.Slider(0, 8, value=3, step=1, label="blocks", scale=1)
            with gr.Row():
                action_2 = gr.Dropdown(MINECRAFT_ACTIONS, value="turn_right",
                                       label="Segment 2", scale=2)
                chunks_2 = gr.Slider(0, 8, value=2, step=1, label="blocks", scale=1)
            with gr.Row():
                action_3 = gr.Dropdown(MINECRAFT_ACTIONS, value="forward",
                                       label="Segment 3", scale=2)
                chunks_3 = gr.Slider(0, 8, value=2, step=1, label="blocks", scale=1)
            with gr.Accordion("Advanced", open=False):
                seed = gr.Slider(0, 2**31 - 1, value=0, step=1, label="Seed")
            run = gr.Button("Roll out the world", variant="primary")

        with gr.Column(scale=1):
            video = gr.Video(label="Generated rollout", autoplay=True, loop=True)
            info = gr.Markdown()

    gr.Examples(
        examples=[
            ["examples/forest.png", "forward", 3, "turn_right", 2, "forward", 2],
            ["examples/plains.png", "forward", 3, "look_up", 1, "forward_turn_right", 3],
            ["examples/cave.png", "forward", 2, "turn_left", 2, "forward", 3],
        ],
        inputs=[image, action_1, chunks_1, action_2, chunks_2, action_3, chunks_3],
        outputs=[video, info],
        fn=generate_example,
        cache_examples=True,
        cache_mode="lazy",
        label="Examples (reference frames from the ForgeWM repo)",
    )

    gr.Markdown(
        "Rollouts are 352×640 at 12 fps. Attention uses a 6-latent-frame sliding "
        "window with a block-causal KV cache, exactly as the released students "
        "were trained. Base weights: "
        "[Skywork/Matrix-Game-2.0](https://huggingface.co/Skywork/Matrix-Game-2.0)."
    )

    # Same entry point the bundled examples use: seven inputs, default student
    # and seed.  Hidden in the UI, exposed on the API / MCP server.
    example_trigger = gr.Button("Run example", visible=False)
    example_trigger.click(
        fn=generate_example,
        inputs=[image, action_1, chunks_1, action_2, chunks_2, action_3, chunks_3],
        outputs=[video, info],
        concurrency_limit=1,
        api_name="generate_default",
    )

    run.click(
        fn=generate,
        inputs=[image, variant, action_1, chunks_1, action_2, chunks_2,
                action_3, chunks_3, seed],
        outputs=[video, info],
        concurrency_limit=1,
        api_name="generate",
    )

demo.queue(max_size=12).launch(
    theme=gr.themes.Citrus(), css=CSS, mcp_server=True, show_error=True)

	"""ForgeWM — few-step, action-conditioned Minecraft world model.

	Gradio demo for `ForgeWM/ForgeWM` (paper: "ForgeWM: Progressive Causal Training
	for Few-Step Action-Conditioned Video World Models").

	Faithful to the repo's own `inference.py` / `pipeline/causal_inference.py`:
	same 352x640 resolution, 3-latent-frame causal blocks, sliding window of 6
	latent frames, the released `warp_denoising_step` schedules, and First-Frame
	Enhancement for the 1-/2-step students.
	"""

	import os

	os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
	os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

	import spaces # noqa: E402 (must precede torch)

	import gc # noqa: E402
	import shutil # noqa: E402
	import tempfile # noqa: E402
	import time # noqa: E402
	from typing import List, Tuple # noqa: E402

	import gradio as gr # noqa: E402
	import imageio.v2 as imageio # noqa: E402
	import numpy as np # noqa: E402
	import torch # noqa: E402
	import torch.nn.functional as F # noqa: E402
	from huggingface_hub import hf_hub_download # noqa: E402
	from omegaconf import OmegaConf # noqa: E402
	from PIL import Image # noqa: E402

	from pipeline import CausalInferencePipeline
	from utils.wan_wrapper import WanVAEWrapper

	# ────────────────────────────── constants ──────────────────────────────
	HEIGHT, WIDTH = 352, 640
	FPS = 12
	VAE_TCR = 4 # VAE temporal compression ratio
	FRAMES_PER_BLOCK = 3 # latent frames per causal block
	CAM_VALUE = 0.10 # camera delta magnitude, from inference.py

	MINECRAFT_ACTIONS = [
	"forward", "back", "left", "right",
	"turn_right", "turn_left", "look_up", "look_down",
	"forward_turn_right", "random", "no_action",
	]

	BASE_REPO = "Skywork/Matrix-Game-2.0"
	FORGEWM_REPO = "ForgeWM/ForgeWM"

	VARIANTS = {
	"ForgeWM-4 · 4 steps": ("configs/stage3_dmd.yaml", "stage3/model.pt"),
	"ForgeWM-2 · 2 steps": ("configs/stage3_dmd_2step.yaml", "2step/model.pt"),
	"ForgeWM-1 · 1 step": ("configs/stage3_dmd_1step.yaml", "1step/model.pt"),
	}
	DEFAULT_VARIANT = "ForgeWM-4 · 4 steps"

	# Measured on the ZeroGPU Blackwell card (12-block rollouts): 1.04 s/block for
	# the 4-step student, 0.40 s/block for the 1-step one, plus ~4 s of CLIP + VAE
	# encode/decode and mp4 muxing. Used only to size the @spaces.GPU reservation.
	SECONDS_PER_CHUNK = {
	"ForgeWM-4 · 4 steps": 1.20,
	"ForgeWM-2 · 2 steps": 0.75,
	"ForgeWM-1 · 1 step": 0.50,
	}
	FIXED_OVERHEAD = 13.0 # CLIP + VAE + mux + first-call weight streaming

	CKPT_DIR = os.path.join(os.getcwd(), "ckpts", "MG2-base")


	# ─────────────────────────── weight preparation ───────────────────────────
	def _link(src: str, dst: str) -> None:
	if os.path.lexists(dst):
	os.remove(dst)
	os.symlink(src, dst)


	def _purge(path: str) -> None:
	"""Drop a hub blob from local disk once its tensors are in memory."""
	try:
	real = os.path.realpath(path)
	if os.path.isfile(real):
	os.remove(real)
	if os.path.islink(path):
	os.remove(path)
	except OSError as exc: # pragma: no cover
	print(f"[disk] could not purge {path}: {exc}")


	def _disk() -> str:
	total, used, free = shutil.disk_usage("/")
	return f"disk {used / 230:.1f}G used / {free / 230:.1f}G free"


	os.makedirs(os.path.join(CKPT_DIR, "xlm-roberta-large"), exist_ok=True)

	print(f"[setup] fetching Matrix-Game-2.0 base weights … ({_disk()})", flush=True)
	_dit = hf_hub_download(BASE_REPO, "base_model/diffusion_pytorch_model.safetensors")
	_link(_dit, os.path.join(CKPT_DIR, "diffusion_pytorch_model.safetensors"))
	_link(hf_hub_download(BASE_REPO, "base_model/base_config.json"),
	os.path.join(CKPT_DIR, "base_config.json"))
	_vae_path = hf_hub_download(BASE_REPO, "Wan2.1_VAE.pth")
	_clip_path = hf_hub_download(
	BASE_REPO, "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth")
	for _tok in ("sentencepiece.bpe.model", "special_tokens_map.json",
	"tokenizer.json", "tokenizer_config.json"):
	_link(hf_hub_download(BASE_REPO, f"xlm-roberta-large/{_tok}"),
	os.path.join(CKPT_DIR, "xlm-roberta-large", _tok))
	print(f"[setup] base weights ready ({_disk()})", flush=True)

	# ─────────────────────────────── models ───────────────────────────────
	torch.set_grad_enabled(False)
	DTYPE = torch.bfloat16

	print("[setup] building VAE + CLIP …", flush=True)
	VAE = WanVAEWrapper(
	vae_path=_vae_path,
	clip_checkpoint_path=_clip_path,
	clip_tokenizer_path=os.path.join(CKPT_DIR, "xlm-roberta-large"),
	).eval()
	VAE = VAE.to(device="cuda", dtype=DTYPE)
	# `WanVAEWrapper.clip` is a plain object, not a submodule, so `.to()` above does
	# not reach it. Move it explicitly (in fp32, as the repo does) so ZeroGPU packs
	# it with everything else instead of paying a 4.8 GB host copy per request.
	VAE.clip.model = VAE.clip.model.to("cuda")
	gc.collect()
	print(f"[setup] VAE + CLIP on GPU ({_disk()})", flush=True)


	def _build_variant(name: str) -> CausalInferencePipeline:
	cfg_path, ckpt_file = VARIANTS[name]
	config = OmegaConf.merge(OmegaConf.load("configs/default.yaml"),
	OmegaConf.load(cfg_path))
	config.model_kwargs.model_name = CKPT_DIR

	pipe = CausalInferencePipeline(config, device=torch.device("cuda"), vae=VAE)

	ckpt = hf_hub_download(FORGEWM_REPO, ckpt_file)
	try:
	state = torch.load(ckpt, map_location="cpu", weights_only=True)
	except Exception:
	state = torch.load(ckpt, map_location="cpu", weights_only=False)
	gen_sd = state.get("generator", state.get("generator_ema", state)) \
	if isinstance(state, dict) else state
	fixed = {k.replace("._fsdp_wrapped_module.", ".")
	.replace("._checkpoint_wrapped_module.", "."): v
	for k, v in gen_sd.items()}
	missing, unexpected = pipe.generator.load_state_dict(fixed, strict=False)
	print(f"[{name}] loaded: missing={len(missing)} unexpected={len(unexpected)}",
	flush=True)
	del state, gen_sd, fixed
	gc.collect()
	_purge(ckpt)

	pipe.generator = pipe.generator.to(device="cuda", dtype=DTYPE).eval()
	pipe.vae = VAE
	gc.collect()
	print(f"[{name}] ready ({_disk()})", flush=True)
	return pipe


	PIPELINES = {}
	for _name in VARIANTS:
	try:
	PIPELINES[_name] = _build_variant(_name)
	except Exception as exc: # keep the Space usable if one student fails
	print(f"[setup] FAILED to load {_name}: {exc!r}", flush=True)

	if not PIPELINES:
	raise RuntimeError("No ForgeWM student could be loaded.")
	if DEFAULT_VARIANT not in PIPELINES:
	DEFAULT_VARIANT = next(iter(PIPELINES))

	# The base DiT safetensors were only needed to instantiate the architecture;
	# every parameter has since been overwritten by the ForgeWM checkpoints.
	_purge(os.path.join(CKPT_DIR, "diffusion_pytorch_model.safetensors"))
	print(f"[setup] {len(PIPELINES)} variant(s) ready ({_disk()})", flush=True)


	# ──────────────────────────── action scripting ────────────────────────────
	def make_action(action_type: str, num_raw_frames: int) -> Tuple[torch.Tensor, torch.Tensor]:
	"""Minecraft action palette (mouse_dim_in=2), verbatim from inference.py."""
	mouse = torch.zeros(1, num_raw_frames, 2)
	keyboard = torch.zeros(1, num_raw_frames, 6)

	if action_type == "forward":
	keyboard[:, :, 0] = 1.0
	elif action_type == "back":
	keyboard[:, :, 1] = 1.0
	elif action_type == "left":
	keyboard[:, :, 2] = 1.0
	elif action_type == "right":
	keyboard[:, :, 3] = 1.0
	elif action_type == "turn_right":
	mouse[:, :, 1] = CAM_VALUE
	elif action_type == "turn_left":
	mouse[:, :, 1] = -CAM_VALUE
	elif action_type == "look_up":
	mouse[:, :, 0] = CAM_VALUE
	elif action_type == "look_down":
	mouse[:, :, 0] = -CAM_VALUE
	elif action_type == "forward_turn_right":
	keyboard[:, :, 0] = 1.0
	mouse[:, :, 1] = CAM_VALUE
	elif action_type == "random":
	torch.manual_seed(42)
	mouse = (torch.rand(1, num_raw_frames, 2) - 0.5) * (2 * CAM_VALUE)
	keyboard[:, :, :4] = (torch.rand(1, num_raw_frames, 4) > 0.5).float()
	elif action_type == "no_action":
	pass
	else:
	raise gr.Error(f"Unknown action '{action_type}'.")
	return mouse, keyboard


	def _chunk_start(chunk_index: int) -> int:
	"""First raw (pixel) frame owned by causal block `chunk_index`.

	Latent frame 0 maps to raw frame 0; latent frame f>=1 maps to raw frames
	4f-3 … 4f. A block spans 3 latent frames, so block c covers raw frames
	[12c-3, 12c+9) — and [0, 9) for c == 0.
	"""
	return 0 if chunk_index == 0 else 12 * chunk_index - 3


	def build_action_track(segments: List[Tuple[str, int]]) -> Tuple[torch.Tensor, torch.Tensor]:
	total_chunks = sum(n for _, n in segments)
	num_raw = _chunk_start(total_chunks)
	mouse = torch.zeros(1, num_raw, 2)
	keyboard = torch.zeros(1, num_raw, 6)
	chunk = 0
	for action, count in segments:
	for _ in range(count):
	lo, hi = _chunk_start(chunk), _chunk_start(chunk + 1)
	m, k = make_action(action, hi - lo)
	mouse[:, lo:hi] = m
	keyboard[:, lo:hi] = k
	chunk += 1
	return mouse, keyboard


	# ───────────────────────────── preprocessing ─────────────────────────────
	def load_reference_frame(image_path: str) -> torch.Tensor:
	"""Aspect-preserving resize + centre crop to 352x640, scaled to [-1, 1]."""
	image = Image.open(image_path).convert("RGB")
	arr = torch.from_numpy(np.asarray(image)).permute(2, 0, 1)[None].float() / 255.0
	_, _, h, w = arr.shape
	if h / w > HEIGHT / WIDTH:
	new_w, new_h = WIDTH, max(HEIGHT, int(round(h * WIDTH / w)))
	else:
	new_h, new_w = HEIGHT, max(WIDTH, int(round(w * HEIGHT / h)))
	arr = F.interpolate(arr, size=(new_h, new_w), mode="bilinear", align_corners=False)
	top, left = (new_h - HEIGHT) // 2, (new_w - WIDTH) // 2
	arr = arr[:, :, top:top + HEIGHT, left:left + WIDTH]
	arr = (arr - 0.5) / 0.5
	return arr.unsqueeze(0).to(device="cuda", dtype=DTYPE) # [1, 1, 3, H, W]


	def build_conditional_dict(pipe, pixel, num_frames, mouse_cond, keyboard_cond):
	"""MG2-style conditioning: CLIP context + first-frame latent + mask."""
	num_pixel_frames = (num_frames - 1) * VAE_TCR + 1
	visual_context = pipe.vae.encode_visual_context_from_pixels(pixel).to(DTYPE)

	first_frame = pixel[:, 0:1]
	pad = torch.zeros(1, num_pixel_frames - 1, 3, pixel.shape[3], pixel.shape[4],
	device=pixel.device, dtype=DTYPE)
	padded = torch.cat([first_frame, pad], dim=1).permute(0, 2, 1, 3, 4)
	img_cond = pipe.vae.encode_to_latent(padded).to(DTYPE)

	_, _, _, h_lat, w_lat = img_cond.shape
	mask = torch.zeros(1, num_frames, 4, h_lat, w_lat,
	device=pixel.device, dtype=DTYPE)
	mask[:, 0:1] = 1
	return {
	"visual_context": visual_context,
	"cond_concat": torch.cat([mask, img_cond], dim=2),
	"mouse_condition": mouse_cond.to(device="cuda", dtype=DTYPE),
	"keyboard_condition": keyboard_cond.to(device="cuda", dtype=DTYPE),
	}


	# ──────────────────────────────── core ────────────────────────────────
	@torch.no_grad()
	def _rollout(image_path, variant, segments, seed):
	if image_path is None:
	raise gr.Error("Please provide a reference frame.")
	if variant not in PIPELINES:
	variant = DEFAULT_VARIANT
	segments = [(a, int(n)) for a, n in segments if int(n) > 0]
	total_chunks = sum(n for _, n in segments)
	if total_chunks < 1:
	raise gr.Error("Give at least one action segment a non-zero length.")
	if total_chunks > 12:
	raise gr.Error("Keep the total rollout at 12 chunks or fewer.")

	pipe = PIPELINES[variant]
	num_frames = total_chunks * FRAMES_PER_BLOCK

	pixel = load_reference_frame(image_path)
	mouse, keyboard = build_action_track(segments)
	cond = build_conditional_dict(pipe, pixel, num_frames, mouse, keyboard)

	torch.manual_seed(int(seed))
	noise = torch.randn([1, num_frames, 16, HEIGHT // 8, WIDTH // 8],
	device="cuda", dtype=DTYPE)

	start = time.time()
	video = pipe.inference(noise=noise, conditional_dict=cond, return_latents=False)
	torch.cuda.synchronize()
	elapsed = time.time() - start
	pipe.vae.model.clear_cache()

	frames = (video[0].permute(0, 2, 3, 1).float().cpu().numpy() * 255)
	frames = frames.clip(0, 255).astype(np.uint8)
	out_path = os.path.join(tempfile.mkdtemp(), "forgewm.mp4")
	writer = imageio.get_writer(out_path, fps=FPS, codec="libx264",
	quality=8, macro_block_size=None)
	for frame in frames:
	writer.append_data(frame)
	writer.close()

	script = " → ".join(f"`{a}` ×{n}" for a, n in segments)
	info = (
	f"{variant} · {total_chunks} causal blocks · {num_frames} latent "
	f"frames → {len(frames)} pixel frames ({len(frames) / FPS:.1f}s @ {FPS} fps)\n\n"
	f"Action script: {script}\n\n"
	f"Rollout + VAE decode: {elapsed:.2f}s "
	f"({1000 * elapsed / total_chunks:.0f} ms per block, decode included) · seed `{int(seed)}`"
	)
	return out_path, info


	def _estimate(variant, chunks) -> int:
	per = SECONDS_PER_CHUNK.get(variant, 1.2)
	return int(FIXED_OVERHEAD + per * max(1, min(int(chunks), 12)))


	def _duration_main(*args):
	variant = args[1] if len(args) > 1 else DEFAULT_VARIANT
	chunks = sum(int(args[i]) for i in (3, 5, 7) if len(args) > i)
	return _estimate(variant, chunks)


	def _duration_example(*args):
	chunks = sum(int(args[i]) for i in (2, 4, 6) if len(args) > i)
	return _estimate(DEFAULT_VARIANT, chunks)


	@spaces.GPU(duration=_duration_main)
	def generate(
	image: str,
	variant: str,
	action_1: str,
	chunks_1: int,
	action_2: str,
	chunks_2: int,
	action_3: str,
	chunks_3: int,
	seed: int = 0,
	progress=gr.Progress(track_tqdm=True),
	):
	"""Roll out an action-conditioned Minecraft video from one reference frame.

	Args:
	image: Path to the reference frame that anchors the world.
	variant: Which few-step ForgeWM student to sample with.
	action_1: Action held during the first segment of the rollout.
	chunks_1: Length of the first segment, in 3-latent-frame blocks (~1s each).
	action_2: Action held during the second segment.
	chunks_2: Length of the second segment, in blocks.
	action_3: Action held during the third segment.
	chunks_3: Length of the third segment, in blocks.
	seed: Random seed for the initial noise.

	Returns:
	An mp4 of the generated rollout and a markdown summary of the run.
	"""
	return _rollout(
	image, variant,
	[(action_1, chunks_1), (action_2, chunks_2), (action_3, chunks_3)],
	seed,
	)


	@spaces.GPU(duration=_duration_example)
	def generate_example(
	image: str,
	action_1: str,
	chunks_1: int,
	action_2: str,
	chunks_2: int,
	action_3: str,
	chunks_3: int,
	progress=gr.Progress(track_tqdm=True),
	):
	"""Run a bundled example with the default student and seed 0."""
	return _rollout(
	image, DEFAULT_VARIANT,
	[(action_1, chunks_1), (action_2, chunks_2), (action_3, chunks_3)],
	0,
	)


	# ───────────────────────────────── UI ─────────────────────────────────
	CSS = """
	.gradio-container { max-width: 1200px !important; }
	"""

	with gr.Blocks(title="ForgeWM") as demo:
	gr.Markdown(
	"# 🎮 ForgeWM — few-step action-conditioned world model\n"
	"Give it one Minecraft frame and a short action script; the "
	"block-causal diffusion transformer rolls the world forward at "
	"1, 2 or 4 denoising steps per 3-frame block.\n\n"
	"[Paper](https://huggingface.co/papers/2608.14022) · "
	"[Model](https://huggingface.co/ForgeWM/ForgeWM) · "
	"[Code](https://github.com/asdfo123/ForgeWM) · "
	"[Project page](https://asdfo123.github.io/ForgeWM/)"
	)

	with gr.Row():
	with gr.Column(scale=1):
	image = gr.Image(label="Reference frame", type="filepath",
	height=280, sources=["upload", "clipboard"])
	variant = gr.Radio(
	choices=list(PIPELINES.keys()),
	value=DEFAULT_VARIANT,
	label="Student",
	info="Fewer steps = faster. 1-/2-step use First-Frame Enhancement.",
	)
	gr.Markdown("### Action script \nEach block is 12 frames ≈ 1 second.")
	with gr.Row():
	action_1 = gr.Dropdown(MINECRAFT_ACTIONS, value="forward",
	label="Segment 1", scale=2)
	chunks_1 = gr.Slider(0, 8, value=3, step=1, label="blocks", scale=1)
	with gr.Row():
	action_2 = gr.Dropdown(MINECRAFT_ACTIONS, value="turn_right",
	label="Segment 2", scale=2)
	chunks_2 = gr.Slider(0, 8, value=2, step=1, label="blocks", scale=1)
	with gr.Row():
	action_3 = gr.Dropdown(MINECRAFT_ACTIONS, value="forward",
	label="Segment 3", scale=2)
	chunks_3 = gr.Slider(0, 8, value=2, step=1, label="blocks", scale=1)
	with gr.Accordion("Advanced", open=False):
	seed = gr.Slider(0, 2**31 - 1, value=0, step=1, label="Seed")
	run = gr.Button("Roll out the world", variant="primary")

	with gr.Column(scale=1):
	video = gr.Video(label="Generated rollout", autoplay=True, loop=True)
	info = gr.Markdown()

	gr.Examples(
	examples=[
	["examples/forest.png", "forward", 3, "turn_right", 2, "forward", 2],
	["examples/plains.png", "forward", 3, "look_up", 1, "forward_turn_right", 3],
	["examples/cave.png", "forward", 2, "turn_left", 2, "forward", 3],
	],
	inputs=[image, action_1, chunks_1, action_2, chunks_2, action_3, chunks_3],
	outputs=[video, info],
	fn=generate_example,
	cache_examples=True,
	cache_mode="lazy",
	label="Examples (reference frames from the ForgeWM repo)",
	)

	gr.Markdown(
	"Rollouts are 352×640 at 12 fps. Attention uses a 6-latent-frame sliding "
	"window with a block-causal KV cache, exactly as the released students "
	"were trained. Base weights: "
	"[Skywork/Matrix-Game-2.0](https://huggingface.co/Skywork/Matrix-Game-2.0)."
	)

	# Same entry point the bundled examples use: seven inputs, default student
	# and seed. Hidden in the UI, exposed on the API / MCP server.
	example_trigger = gr.Button("Run example", visible=False)
	example_trigger.click(
	fn=generate_example,
	inputs=[image, action_1, chunks_1, action_2, chunks_2, action_3, chunks_3],
	outputs=[video, info],
	concurrency_limit=1,
	api_name="generate_default",
	)

	run.click(
	fn=generate,
	inputs=[image, variant, action_1, chunks_1, action_2, chunks_2,
	action_3, chunks_3, seed],
	outputs=[video, info],
	concurrency_limit=1,
	api_name="generate",
	)

	demo.queue(max_size=12).launch(
	theme=gr.themes.Citrus(), css=CSS, mcp_server=True, show_error=True)