Spaces:

hugging-apps
/

ppmem

Running on Zero

ppmem / app.py

Upload app.py with huggingface_hub

8af6a7d verified 1 day ago

20.2 kB

"""
PPMem: Persistent Planar Memory for Video World Models — interactive demo.
Single-image, camera-posed video generation with the authors' released checkpoint
(Robbyant-Research/PPMem): a starting image, a camera trajectory (or a synthesized
gentle dolly), optional geometry (depth + intrinsics) — and out comes a 25-frame
1024x576 RGB video plus the planar geometry the model predicts along that camera:
depth maps, normal maps and a world-space planar mesh (PLY).
This is a 1:1 port of the reference rollout (ppmem.rollout.step0 + VideoGenerator.
generate_video) restricted to one AR chunk. The one hardware-forced deviation: the
reference encodes the prompt at build time and frees the text encoder; on ZeroGPU no
compute is allowed at module scope and the prompt arrives per request, so tokenizer +
UMT5 text encoder stay resident and prompts are encoded inside the GPU call.
"""

import os

# Cut VRAM fragmentation on the large per-request allocations (DiT + VAE + TE).
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  (must precede torch / CUDA-touching imports)

import gc
import tempfile
import time
from pathlib import Path

import cv2
import gradio as gr
import numpy as np
import torch
from einops import rearrange
from PIL import Image

from huggingface_hub import snapshot_download

from ppmem._logging import logger
from ppmem.config import MODEL, JobConfig
from ppmem.engine import depth_latent
from ppmem.engine import depth_latent as _depth_latent
from ppmem.engine.determinism import set_determinism
from ppmem.model.dit import WanTransformerPlana3DModel
from ppmem.model.loaders import (
    load_text_encoder,
    load_tokenizer,
    load_vae,
    read_checkpoint,
)
from ppmem.model.pipeline import PPMemImageToVideoPipeline, t5_prompt_embeds
from ppmem.model.planes import (
    planes_to_mesh,
    save_depth_maps,
    save_normal_maps,
    save_planes_as_ply,
    transform_planes_to_world,
)
from ppmem.model.scheduler import FlowUniPCMultistepScheduler
from ppmem.pool.init_geometry import load_init_geometry
from ppmem.pool.segment_math import chunk_seed
from ppmem.rollout.latents import LatentEncoder
from ppmem.rollout.moge import MoGeEstimator
from ppmem.rollout.step0 import step0

try:
    from diffusers.utils import export_to_video
except ImportError:  # diffuters moved it around between versions
    from diffusers.video_utils import export_to_video

SEGMENT_SIZE = 25           # frames per AR chunk (the released rollout's chunk size)
VIDEO_H, VIDEO_W = 576, 1024
FPS = 15
NUM_CONTEXT = 5
DEFAULT_STEPS = 50
DEFAULT_GUIDANCE = 5.0
DEFAULT_SEED = 42

DTYPE = torch.bfloat16
DEVICE = "cuda"

_job = JobConfig()  # run defaults: steps=50, guidance=5.0, flow_shift=5.0

# ----------------------------------------------------------------------
# Module-scope model load (eager; ZeroGPU intercepts .to("cuda") and streams
# the weights into VRAM on the first @spaces.GPU entry).
# ----------------------------------------------------------------------

ckpt_dir = snapshot_download(repo_id="Robbyant-Research/PPMem")
logger.info(f"[demo] checkpoint at {ckpt_dir}")

set_determinism(_job.run.deterministic)

tokenizer = load_tokenizer(ckpt_dir, subfolder="tokenizer")
text_encoder = load_text_encoder(
    ckpt_dir, subfolder="text_encoder", torch_dtype=DTYPE, torch_device=DEVICE,
)
text_encoder.requires_grad_(False)

with torch.device("meta"):
    transformer = WanTransformerPlana3DModel(MODEL)
transformer.cast_tensors(DTYPE)
# The reference builder materializes via to_empty(device); on ZeroGPU the plain
# .to("cuda") path is the one the hijack intercepts (to_empty allocates tensors
# directly with torch.empty(device=...), which may bypass the monkey-patch), so
# materialize on CPU, load the checkpoint, then hand the result to .to("cuda")
# like a normal load.
transformer.to_empty(device="cpu")
with torch.no_grad():
    transformer.init_weights()

_state_dict, _read_as = read_checkpoint(ckpt_dir)
_missing, _unexpected = transformer.load_state_dict(_state_dict, strict=False)
del _state_dict
gc.collect()
_missing = [k for k in _missing if not k.startswith("vae.")]
if _missing or _unexpected:
    raise RuntimeError(
        f"checkpoint mismatch: missing={_missing[:5]} unexpected={_unexpected[:5]}"
    )
transformer.cast_for_inference(DTYPE)
transformer.to(DEVICE)  # ZeroGPU-hijacked: weights packed to disk, streamed into VRAM per call

vae = load_vae(ckpt_dir, subfolder="vae", torch_dtype=DTYPE, torch_device=DEVICE)
vae.requires_grad_(False)
transformer.vae = vae  # dit.py reads vae.config latents_mean/std during denoising

scheduler = FlowUniPCMultistepScheduler(shift=_job.run.flow_shift)
pipeline = PPMemImageToVideoPipeline(
    tokenizer=None, text_encoder=None, transformer=transformer, vae=vae,
    scheduler=scheduler, pipe_dtype=DTYPE,
)

encoder = LatentEncoder(vae, device=DEVICE, dtype=DTYPE)


# The repo binds InferenceEngine.encode_depth_to_latent AND normalize_depth to the
# plain `self`-taking functions (engine/builder.py:115-116: `InferenceEngine.normalize_depth
# = depth_latent.normalize_depth`); encode_depth_to_latent calls self.normalize_depth
# internally (depth_latent.py:88). Mirror that assignment on a host object exactly:
# plain-function class attributes bind as methods, so `self` resolves correctly.
class _DepthCodec:
    normalize_depth = _depth_latent.normalize_depth
    encode_depth_to_latent = _depth_latent.encode_depth_to_latent


codec = _DepthCodec()
codec.vae = vae
codec.model_dtype = DTYPE
codec.device = DEVICE
encode_depth_to_latent = codec.encode_depth_to_latent

moge = None  # lazy: bundled examples ship geometry.npz and never touch MoGe


# ----------------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------------


def synthesize_dolly(dist: float = 1.5, n: int = SEGMENT_SIZE) -> np.ndarray:
    """A gentle forward dolly (authors' `traj_dolly` shape), (n, 4, 4) c2w float32.
    Frame 0 is the identity pose; the camera translates forward along its optical
    axis, within the per-frame motion guidelines the repo documents (<=0.20 m/frame).
    """
    half = (n - 1) / 2.0
    zs = [dist * (i / half) if i <= half else dist * (2.0 - i / half) for i in range(n)]
    poses = np.stack([np.eye(4, dtype=np.float64) for _ in range(n)])
    for i, z in enumerate(zs):
        poses[i, 2, 3] = z
    return poses.astype(np.float32)


def _load_camera(camera_path) -> np.ndarray:
    """Read a user-supplied camera .npy and validate its shape; or synthesize a dolly."""
    if not camera_path:
        return synthesize_dolly()
    arr = np.load(camera_path)
    if arr.ndim != 3 or arr.shape[1:] != (4, 4):
        raise ValueError(
            f"camera file must hold an (N, 4, 4) array of c2w poses, got {arr.shape}"
        )
    return arr


def _depth_video_file(out_dir: Path, prefix: str = "depth") -> str:
    """Stitch the turbo-colormapped depth PNGs into an mp4 (10 evenly spaced frames)."""
    import glob as _glob

    frames = sorted(_glob.glob(str(out_dir / f"{prefix}_*.png")))
    if not frames:
        return None
    pickup = frames[:: max(1, len(frames) // 10)][:10]
    arr = np.stack([np.asarray(Image.open(f).convert("RGB")) for f in pickup])
    path = out_dir / f"{prefix}_video.mp4"
    try:
        import imageio

        imageio.mimwrite(str(path), arr, fps=5, quality=9, macro_block_size=1)
    except Exception as exc:  # ffmpeg backend problems: fall back to gr.Video-free PNGs
        logger.warning(f"depth video write failed ({exc}); returning first PNG")
        return frames[0]
    return str(path)


def _gpu_duration(
    image=None, camera_file=None, geometry_file=None, prompt=None,
    seed=DEFAULT_SEED, steps=DEFAULT_STEPS, guidance=DEFAULT_GUIDANCE,
    *args, **kwargs,
):
    """Reserve GPU seconds from the measured 50-step run (scene01, geometry supplied:
    ~290 s GPU time end-to-end, cold worker included). Denoise scales ~linearly with
    steps, so the request shrinks with the slider and visitors with partial quota can
    still run at fewer steps. Without a geometry file MoGe-2 depth estimation runs
    first, which adds a fixed cost."""
    t = 60 + float(steps) * 4.8
    if not geometry_file:
        t += 35
    return int(min(t, 360))


@spaces.GPU(duration=_gpu_duration)
def generate(
    image,
    camera_file,
    geometry_file,
    prompt: str,
    seed: int = DEFAULT_SEED,
    steps: int = DEFAULT_STEPS,
    guidance: float = DEFAULT_GUIDANCE,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a 25-frame camera-posed video plus planar geometry from one image.
    Args:
        image: the starting frame (RGB image).
        camera_file: optional .npy of (N, 4, 4) camera-to-world poses. When empty,
            a gentle forward dolly is synthesized.
        geometry_file: optional .npz with depth/fx/fy/cx/cy for frame 0 (as the
            authors' examples ship). When empty, MoGe-2 estimates it from the image.
        prompt: scene description for the text conditioning.
        seed: generation seed (chunk seed = seed + 0, as the reference rollout at step 0).
        steps: number of UniPC inference steps.
        guidance: classifier-free guidance scale.
        progress: Gradio progress.
    Returns:
        (video_path, depth_video or first depth map, [normal map paths], ply_path)
    """
    global moge
    t0 = time.time()
    progress((1, 12), desc="prep")
    work = Path(tempfile.mkdtemp(prefix="ppmem_"))

    if image is None:
        raise gr.Error("Please provide a starting image.")
    src_image = Path(work) / "input.png"
    if isinstance(image, np.ndarray):  # gr.Image numpy (H, W, 3)
        Image.fromarray(image.astype(np.uint8)).save(src_image)
    else:
        Image.open(image).convert("RGB").save(src_image)

    # ---- geometry bootstrap (reference: ppmem.cli --init_geometry) ----
    init_depth = init_K = None
    if geometry_file:
        with Image.open(src_image) as im:
            hw_image = (im.size[1], im.size[0])
        init_depth, init_K = load_init_geometry(geometry_file, (VIDEO_H, VIDEO_W), hw_image)

    extrinsic_c2w = _load_camera(camera_file)

    # ---- step0: projection + conditioning latents (1:1 reference path) ----
    if init_depth is None:
        if moge is None:
            moge = MoGeEstimator(model_type="moge-2-vitl-normal", device=DEVICE)
    progress((2, 12), desc="step0: project + encode")
    meta = step0(
        image_path=str(src_image),
        extrinsic_c2w=extrinsic_c2w,
        prompt=prompt,
        incremental_dir=str(work / "pool"),
        h_tgt=VIDEO_H,
        w_tgt=VIDEO_W,
        device=DEVICE,
        moge=moge,
        encoder=encoder,
        num_context=NUM_CONTEXT,
        segment_size=SEGMENT_SIZE,
        init_depth=init_depth,
        init_K=init_K,
    )
    meta.pop("proj_rgb_frames", None)

    # ---- prompt embeds (the reference pre-encodes; ZeroGPU forces per-call) ----
    progress((3, 12), desc="encode prompt")
    with torch.no_grad():
        prompt_embeds = t5_prompt_embeds(tokenizer, text_encoder, prompt, device=DEVICE)
        negative_prompt_embeds = t5_prompt_embeds(tokenizer, text_encoder, "", device=DEVICE)

    # ---- func_input, exactly as VideoGenerator.generate_video assembles it ----
    latent_h = meta["latent_height"]
    latent_w = meta["latent_width"]
    latent_f = meta["latent_num_frames"]

    generator = torch.manual_seed(chunk_seed(seed, 0))

    func_input = {
        "prompt_embeds": prompt_embeds,
        "negative_prompt_embeds": negative_prompt_embeds,
        "height": meta["video_height"],
        "width": meta["video_width"],
        "num_frames": meta["video_num_frames"],
        "num_inference_steps": steps,
        "guidance_scale": guidance,
        "generator": generator,
        "image": meta["first_frame_pil_images"].resize((meta["video_width"], meta["video_height"])),
    }

    num_ctx = meta.get("context_num_frames", 0)
    if num_ctx > 0 and "latent_context_frames" in meta:
        func_input["latents_context"] = rearrange(
            meta["latent_context_frames"], "(f h w) c -> c f h w",
            f=num_ctx, h=latent_h, w=latent_w,
        ).to(DEVICE)

    vf_data = meta["latent_very_first_frame"]
    n_vf = vf_data.shape[0] // (latent_h * latent_w)
    func_input["latents_very_first"] = rearrange(
        vf_data, "(f h w) c -> c f h w", f=n_vf, h=latent_h, w=latent_w,
    ).to(DEVICE)

    for suffix in ("1", "2"):
        for prefix in ("proj_rgb", "proj_mask"):
            key = f"latent_{prefix}_{suffix}"
            func_input[f"latent_{prefix}_{suffix}"] = rearrange(
                meta[key], "(f h w) c -> c f h w",
                f=latent_f, h=latent_h, w=latent_w,
            ).to(DEVICE)

    depth_norm_method = MODEL.depth_norm_method  # 'log'
    depth_range = (MODEL.depth_norm_range_min, MODEL.depth_norm_range_max)  # (0.2, 99.0)

    next_ar_depth = meta["next_ar_depth"].to(DEVICE)
    func_input["latents_depth"] = encode_depth_to_latent(
        next_ar_depth, norm_method=depth_norm_method, depth_range=depth_range,
    ).squeeze(0)
    # hv = history (none at step 0) + very first (frame 0's depth)
    hv_vf = meta["next_ar_depth"][0].cpu().numpy().astype(np.float32)
    hv_depth = torch.from_numpy(hv_vf[None, ...]).to(DEVICE)
    func_input["latents_hv_depth"] = encode_depth_to_latent(
        hv_depth, norm_method=depth_norm_method, depth_range=depth_range,
    ).squeeze(0)
    func_input["raw_next_ar_depth"] = [next_ar_depth]
    if num_ctx > 0 and meta.get("context_depth", None) is not None and meta["context_depth"].numel() > 0:
        func_input["latents_context_depth"] = encode_depth_to_latent(
            meta["context_depth"].to(DEVICE),
            norm_method=depth_norm_method, depth_range=depth_range,
        ).squeeze(0)

    K_list = []
    for intr in meta["intrinsics"]:
        K_list.append(torch.tensor([
            [intr["fx"], 0.0, intr["cx"]],
            [0.0, intr["fy"], intr["cy"]],
            [0.0, 0.0, 1.0],
        ], dtype=torch.float32))
    intrinsics = torch.stack(K_list).to(DEVICE)

    extrinsics_c2w_full = torch.tensor(meta["extrinsics_c2w"], dtype=torch.float32).to(DEVICE)
    if extrinsics_c2w_full.shape[0] < SEGMENT_SIZE:  # reference pads short sequences
        pad = SEGMENT_SIZE - extrinsics_c2w_full.shape[0]
        extrinsics_c2w_full = torch.cat(
            [extrinsics_c2w_full, extrinsics_c2w_full[-1:].repeat(pad, 1, 1)], dim=0
        )
    intrinsics = intrinsics[:SEGMENT_SIZE]
    while intrinsics.shape[0] < SEGMENT_SIZE:
        intrinsics = torch.cat([intrinsics, intrinsics[-1:]], dim=0)
    func_input["intrinsics"] = intrinsics
    func_input["extrinsics_c2w"] = extrinsics_c2w_full[:SEGMENT_SIZE]

    # This VAE maps frames 1:1; pass the count the pipeline maps back to
    # segment_size latent frames (Wan temporal convention: (n-1)*4+1).
    func_input["num_frames"] = (func_input["num_frames"] - 1) * 4 + 1

    progress((4, 12), desc="denoising (UniPC)")
    with torch.no_grad():
        pipeline_output = pipeline(**func_input)
    videos = pipeline_output.frames
    plane_params = getattr(pipeline_output, "plane_params", None)

    progress((10, 12), desc="export")
    out_dir = work / "out"
    out_dir.mkdir(parents=True, exist_ok=True)
    video_path = out_dir / "generated_video.mp4"
    export_to_video(videos[0], str(video_path), fps=FPS)

    depth_out = None
    normal_paths = []
    ply_path = None
    if plane_params is not None:
        extr = plane_params.get("extrinsics_c2w", None)
        pps = plane_params.get("plane_params_per_video", [])
        if extr is not None and len(pps):
            pp = pps[0]
            extr_t = extr if torch.is_tensor(extr) else torch.tensor(np.asarray(extr), dtype=torch.float32)
            verts_w, faces_w, colors_w, vnormals_w = planes_to_mesh(pp, extr_t.to(pp["plane_center_camera"].device))
            ply_path = out_dir / "planes_world.ply"
            save_planes_as_ply(str(ply_path), verts_w, faces_w, colors_w, vnormals_w)
            ply_path = str(ply_path)

            d_dir = out_dir / "cam_0_depth"
            save_depth_maps(str(d_dir), pp)
            depth_out = _depth_video_file(Path(d_dir))

            n_dir = out_dir / "cam_0_normals"
            save_normal_maps(str(n_dir), pp)
            import glob as _g

            normal_paths = sorted(_g.glob(str(n_dir / "*.png")))[:: max(1, len(sorted(_g.glob(str(n_dir / "*.png")))) // 4)][:4]

    logger.info(f"[demo] generate done in {time.time()-t0:.1f}s -> {video_path}")

    if depth_out is None:
        depth_out = str(video_path)  # keep the slot filled if geometry was skipped
    return str(video_path), depth_out, normal_paths, ply_path


# ----------------------------------------------------------------------
# UI
# ----------------------------------------------------------------------

examples_dir = Path(__file__).parent / "examples"
_scenes = sorted(p.name for p in examples_dir.glob("scene*") if p.is_dir())
EXAMPLES = []
for scene in _scenes:
    base = examples_dir / scene
    prompt_text = (base / "prompt.txt").read_text().strip() if (base / "prompt.txt").exists() else "A video of an indoor scene."
    EXAMPLES.append([
        str(base / "input.png"),
        str(base / "camera.npy"),
        str(base / "geometry.npz"),
        prompt_text,
    ])

with gr.Blocks(title="PPMem: Persistent Planar Memory for Video World Models") as demo:
    gr.Markdown(
        """
        # PPMem: Persistent Planar Memory for Video World Models
        **Paper** [2610.23795](https://huggingface.co/papers/2610.23795) · **Code** [robbyant-research/PPMem](https://github.com/robbyant-research/PPMem) · **Weights** [Robbyant-Research/PPMem](https://huggingface.co/Robbyant-Research/PPMem)
        An interactive video world model: give it a **starting image** and a **camera trajectory**, and it generates a 25-frame 1024×576 video from the first pose — plus the **3D planar geometry** it predicts along that camera: depth maps, normal maps, and a world-space planar mesh.
        """
    )
    with gr.Row():
        with gr.Column():
            image_in = gr.Image(label="Starting image", type="filepath", height=280)
            camera_in = gr.File(label="Camera trajectory (.npy, (N,4,4) c2w poses) — optional; a gentle dolly is used if empty", file_types=[".npy"])
            geometry_in = gr.File(label="Geometry .npz (depth, fx, fy, cx, cy) — optional; MoGe-2 estimates it if empty", file_types=[".npz"])
            prompt_in = gr.Textbox(label="Prompt", lines=3, placeholder="Describe the scene…")
            with gr.Accordion("Advanced", open=False):
                seed_in = gr.Number(label="Seed", value=DEFAULT_SEED, precision=0)
                steps_in = gr.Slider(label="Inference steps", minimum=10, maximum=50, value=DEFAULT_STEPS, step=1)
                guidance_in = gr.Slider(label="Guidance scale", minimum=1.0, maximum=10.0, value=DEFAULT_GUIDANCE, step=0.5)
            run_btn = gr.Button("Generate video", variant="primary")
        with gr.Column():
            video_out = gr.Video(label="Generated RGB video (1024×576, 25 frames @ 15fps)")
            depth_out = gr.Video(label="Predicted depth (turbo, from planar geometry)")
            normals_out = gr.Gallery(label="Predicted normals (sample frames)", columns=4, height=180)
            planes_out = gr.Model3D(label="World-space planar mesh (PLY)")

    run_btn.click(
        generate,
        inputs=[image_in, camera_in, geometry_in, prompt_in, seed_in, steps_in, guidance_in],
        outputs=[video_out, depth_out, normals_out, planes_out],
        api_name="generate",
    )

    if EXAMPLES:
        gr.Examples(
            examples=EXAMPLES,
            inputs=[image_in, camera_in, geometry_in, prompt_in],
            fn=generate,
            outputs=[video_out, depth_out, normals_out, planes_out],
            cache_examples=True,
            cache_mode="lazy",
        )

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus())

	"""
	PPMem: Persistent Planar Memory for Video World Models — interactive demo.

	Single-image, camera-posed video generation with the authors' released checkpoint
	(Robbyant-Research/PPMem): a starting image, a camera trajectory (or a synthesized
	gentle dolly), optional geometry (depth + intrinsics) — and out comes a 25-frame
	1024x576 RGB video plus the planar geometry the model predicts along that camera:
	depth maps, normal maps and a world-space planar mesh (PLY).

	This is a 1:1 port of the reference rollout (ppmem.rollout.step0 + VideoGenerator.
	generate_video) restricted to one AR chunk. The one hardware-forced deviation: the
	reference encodes the prompt at build time and frees the text encoder; on ZeroGPU no
	compute is allowed at module scope and the prompt arrives per request, so tokenizer +
	UMT5 text encoder stay resident and prompts are encoded inside the GPU call.
	"""

	import os

	# Cut VRAM fragmentation on the large per-request allocations (DiT + VAE + TE).
	os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

	import spaces # noqa: E402 (must precede torch / CUDA-touching imports)

	import gc
	import tempfile
	import time
	from pathlib import Path

	import cv2
	import gradio as gr
	import numpy as np
	import torch
	from einops import rearrange
	from PIL import Image

	from huggingface_hub import snapshot_download

	from ppmem._logging import logger
	from ppmem.config import MODEL, JobConfig
	from ppmem.engine import depth_latent
	from ppmem.engine import depth_latent as _depth_latent
	from ppmem.engine.determinism import set_determinism
	from ppmem.model.dit import WanTransformerPlana3DModel
	from ppmem.model.loaders import (
	load_text_encoder,
	load_tokenizer,
	load_vae,
	read_checkpoint,
	)
	from ppmem.model.pipeline import PPMemImageToVideoPipeline, t5_prompt_embeds
	from ppmem.model.planes import (
	planes_to_mesh,
	save_depth_maps,
	save_normal_maps,
	save_planes_as_ply,
	transform_planes_to_world,
	)
	from ppmem.model.scheduler import FlowUniPCMultistepScheduler
	from ppmem.pool.init_geometry import load_init_geometry
	from ppmem.pool.segment_math import chunk_seed
	from ppmem.rollout.latents import LatentEncoder
	from ppmem.rollout.moge import MoGeEstimator
	from ppmem.rollout.step0 import step0

	try:
	from diffusers.utils import export_to_video
	except ImportError: # diffuters moved it around between versions
	from diffusers.video_utils import export_to_video

	SEGMENT_SIZE = 25 # frames per AR chunk (the released rollout's chunk size)
	VIDEO_H, VIDEO_W = 576, 1024
	FPS = 15
	NUM_CONTEXT = 5
	DEFAULT_STEPS = 50
	DEFAULT_GUIDANCE = 5.0
	DEFAULT_SEED = 42

	DTYPE = torch.bfloat16
	DEVICE = "cuda"

	_job = JobConfig() # run defaults: steps=50, guidance=5.0, flow_shift=5.0

	# ----------------------------------------------------------------------
	# Module-scope model load (eager; ZeroGPU intercepts .to("cuda") and streams
	# the weights into VRAM on the first @spaces.GPU entry).
	# ----------------------------------------------------------------------

	ckpt_dir = snapshot_download(repo_id="Robbyant-Research/PPMem")
	logger.info(f"[demo] checkpoint at {ckpt_dir}")

	set_determinism(_job.run.deterministic)

	tokenizer = load_tokenizer(ckpt_dir, subfolder="tokenizer")
	text_encoder = load_text_encoder(
	ckpt_dir, subfolder="text_encoder", torch_dtype=DTYPE, torch_device=DEVICE,
	)
	text_encoder.requires_grad_(False)

	with torch.device("meta"):
	transformer = WanTransformerPlana3DModel(MODEL)
	transformer.cast_tensors(DTYPE)
	# The reference builder materializes via to_empty(device); on ZeroGPU the plain
	# .to("cuda") path is the one the hijack intercepts (to_empty allocates tensors
	# directly with torch.empty(device=...), which may bypass the monkey-patch), so
	# materialize on CPU, load the checkpoint, then hand the result to .to("cuda")
	# like a normal load.
	transformer.to_empty(device="cpu")
	with torch.no_grad():
	transformer.init_weights()

	_state_dict, _read_as = read_checkpoint(ckpt_dir)
	_missing, _unexpected = transformer.load_state_dict(_state_dict, strict=False)
	del _state_dict
	gc.collect()
	_missing = [k for k in _missing if not k.startswith("vae.")]
	if _missing or _unexpected:
	raise RuntimeError(
	f"checkpoint mismatch: missing={_missing[:5]} unexpected={_unexpected[:5]}"
	)
	transformer.cast_for_inference(DTYPE)
	transformer.to(DEVICE) # ZeroGPU-hijacked: weights packed to disk, streamed into VRAM per call

	vae = load_vae(ckpt_dir, subfolder="vae", torch_dtype=DTYPE, torch_device=DEVICE)
	vae.requires_grad_(False)
	transformer.vae = vae # dit.py reads vae.config latents_mean/std during denoising

	scheduler = FlowUniPCMultistepScheduler(shift=_job.run.flow_shift)
	pipeline = PPMemImageToVideoPipeline(
	tokenizer=None, text_encoder=None, transformer=transformer, vae=vae,
	scheduler=scheduler, pipe_dtype=DTYPE,
	)

	encoder = LatentEncoder(vae, device=DEVICE, dtype=DTYPE)


	# The repo binds InferenceEngine.encode_depth_to_latent AND normalize_depth to the
	# plain `self`-taking functions (engine/builder.py:115-116: `InferenceEngine.normalize_depth
	# = depth_latent.normalize_depth`); encode_depth_to_latent calls self.normalize_depth
	# internally (depth_latent.py:88). Mirror that assignment on a host object exactly:
	# plain-function class attributes bind as methods, so `self` resolves correctly.
	class _DepthCodec:
	normalize_depth = _depth_latent.normalize_depth
	encode_depth_to_latent = _depth_latent.encode_depth_to_latent


	codec = _DepthCodec()
	codec.vae = vae
	codec.model_dtype = DTYPE
	codec.device = DEVICE
	encode_depth_to_latent = codec.encode_depth_to_latent

	moge = None # lazy: bundled examples ship geometry.npz and never touch MoGe


	# ----------------------------------------------------------------------
	# Helpers
	# ----------------------------------------------------------------------


	def synthesize_dolly(dist: float = 1.5, n: int = SEGMENT_SIZE) -> np.ndarray:
	"""A gentle forward dolly (authors' `traj_dolly` shape), (n, 4, 4) c2w float32.

	Frame 0 is the identity pose; the camera translates forward along its optical
	axis, within the per-frame motion guidelines the repo documents (<=0.20 m/frame).
	"""
	half = (n - 1) / 2.0
	zs = [dist * (i / half) if i <= half else dist * (2.0 - i / half) for i in range(n)]
	poses = np.stack([np.eye(4, dtype=np.float64) for _ in range(n)])
	for i, z in enumerate(zs):
	poses[i, 2, 3] = z
	return poses.astype(np.float32)


	def _load_camera(camera_path) -> np.ndarray:
	"""Read a user-supplied camera .npy and validate its shape; or synthesize a dolly."""
	if not camera_path:
	return synthesize_dolly()
	arr = np.load(camera_path)
	if arr.ndim != 3 or arr.shape[1:] != (4, 4):
	raise ValueError(
	f"camera file must hold an (N, 4, 4) array of c2w poses, got {arr.shape}"
	)
	return arr


	def _depth_video_file(out_dir: Path, prefix: str = "depth") -> str:
	"""Stitch the turbo-colormapped depth PNGs into an mp4 (10 evenly spaced frames)."""
	import glob as _glob

	frames = sorted(_glob.glob(str(out_dir / f"{prefix}_*.png")))
	if not frames:
	return None
	pickup = frames[:: max(1, len(frames) // 10)][:10]
	arr = np.stack([np.asarray(Image.open(f).convert("RGB")) for f in pickup])
	path = out_dir / f"{prefix}_video.mp4"
	try:
	import imageio

	imageio.mimwrite(str(path), arr, fps=5, quality=9, macro_block_size=1)
	except Exception as exc: # ffmpeg backend problems: fall back to gr.Video-free PNGs
	logger.warning(f"depth video write failed ({exc}); returning first PNG")
	return frames[0]
	return str(path)


	def _gpu_duration(
	image=None, camera_file=None, geometry_file=None, prompt=None,
	seed=DEFAULT_SEED, steps=DEFAULT_STEPS, guidance=DEFAULT_GUIDANCE,
	args, *kwargs,
	):
	"""Reserve GPU seconds from the measured 50-step run (scene01, geometry supplied:
	~290 s GPU time end-to-end, cold worker included). Denoise scales ~linearly with
	steps, so the request shrinks with the slider and visitors with partial quota can
	still run at fewer steps. Without a geometry file MoGe-2 depth estimation runs
	first, which adds a fixed cost."""
	t = 60 + float(steps) * 4.8
	if not geometry_file:
	t += 35
	return int(min(t, 360))


	@spaces.GPU(duration=_gpu_duration)
	def generate(
	image,
	camera_file,
	geometry_file,
	prompt: str,
	seed: int = DEFAULT_SEED,
	steps: int = DEFAULT_STEPS,
	guidance: float = DEFAULT_GUIDANCE,
	progress=gr.Progress(track_tqdm=True),
	):
	"""Generate a 25-frame camera-posed video plus planar geometry from one image.

	Args:
	image: the starting frame (RGB image).
	camera_file: optional .npy of (N, 4, 4) camera-to-world poses. When empty,
	a gentle forward dolly is synthesized.
	geometry_file: optional .npz with depth/fx/fy/cx/cy for frame 0 (as the
	authors' examples ship). When empty, MoGe-2 estimates it from the image.
	prompt: scene description for the text conditioning.
	seed: generation seed (chunk seed = seed + 0, as the reference rollout at step 0).
	steps: number of UniPC inference steps.
	guidance: classifier-free guidance scale.
	progress: Gradio progress.

	Returns:
	(video_path, depth_video or first depth map, [normal map paths], ply_path)
	"""
	global moge
	t0 = time.time()
	progress((1, 12), desc="prep")
	work = Path(tempfile.mkdtemp(prefix="ppmem_"))

	if image is None:
	raise gr.Error("Please provide a starting image.")
	src_image = Path(work) / "input.png"
	if isinstance(image, np.ndarray): # gr.Image numpy (H, W, 3)
	Image.fromarray(image.astype(np.uint8)).save(src_image)
	else:
	Image.open(image).convert("RGB").save(src_image)

	# ---- geometry bootstrap (reference: ppmem.cli --init_geometry) ----
	init_depth = init_K = None
	if geometry_file:
	with Image.open(src_image) as im:
	hw_image = (im.size[1], im.size[0])
	init_depth, init_K = load_init_geometry(geometry_file, (VIDEO_H, VIDEO_W), hw_image)

	extrinsic_c2w = _load_camera(camera_file)

	# ---- step0: projection + conditioning latents (1:1 reference path) ----
	if init_depth is None:
	if moge is None:
	moge = MoGeEstimator(model_type="moge-2-vitl-normal", device=DEVICE)
	progress((2, 12), desc="step0: project + encode")
	meta = step0(
	image_path=str(src_image),
	extrinsic_c2w=extrinsic_c2w,
	prompt=prompt,
	incremental_dir=str(work / "pool"),
	h_tgt=VIDEO_H,
	w_tgt=VIDEO_W,
	device=DEVICE,
	moge=moge,
	encoder=encoder,
	num_context=NUM_CONTEXT,
	segment_size=SEGMENT_SIZE,
	init_depth=init_depth,
	init_K=init_K,
	)
	meta.pop("proj_rgb_frames", None)

	# ---- prompt embeds (the reference pre-encodes; ZeroGPU forces per-call) ----
	progress((3, 12), desc="encode prompt")
	with torch.no_grad():
	prompt_embeds = t5_prompt_embeds(tokenizer, text_encoder, prompt, device=DEVICE)
	negative_prompt_embeds = t5_prompt_embeds(tokenizer, text_encoder, "", device=DEVICE)

	# ---- func_input, exactly as VideoGenerator.generate_video assembles it ----
	latent_h = meta["latent_height"]
	latent_w = meta["latent_width"]
	latent_f = meta["latent_num_frames"]

	generator = torch.manual_seed(chunk_seed(seed, 0))

	func_input = {
	"prompt_embeds": prompt_embeds,
	"negative_prompt_embeds": negative_prompt_embeds,
	"height": meta["video_height"],
	"width": meta["video_width"],
	"num_frames": meta["video_num_frames"],
	"num_inference_steps": steps,
	"guidance_scale": guidance,
	"generator": generator,
	"image": meta["first_frame_pil_images"].resize((meta["video_width"], meta["video_height"])),
	}

	num_ctx = meta.get("context_num_frames", 0)
	if num_ctx > 0 and "latent_context_frames" in meta:
	func_input["latents_context"] = rearrange(
	meta["latent_context_frames"], "(f h w) c -> c f h w",
	f=num_ctx, h=latent_h, w=latent_w,
	).to(DEVICE)

	vf_data = meta["latent_very_first_frame"]
	n_vf = vf_data.shape[0] // (latent_h * latent_w)
	func_input["latents_very_first"] = rearrange(
	vf_data, "(f h w) c -> c f h w", f=n_vf, h=latent_h, w=latent_w,
	).to(DEVICE)

	for suffix in ("1", "2"):
	for prefix in ("proj_rgb", "proj_mask"):
	key = f"latent_{prefix}_{suffix}"
	func_input[f"latent_{prefix}_{suffix}"] = rearrange(
	meta[key], "(f h w) c -> c f h w",
	f=latent_f, h=latent_h, w=latent_w,
	).to(DEVICE)

	depth_norm_method = MODEL.depth_norm_method # 'log'
	depth_range = (MODEL.depth_norm_range_min, MODEL.depth_norm_range_max) # (0.2, 99.0)

	next_ar_depth = meta["next_ar_depth"].to(DEVICE)
	func_input["latents_depth"] = encode_depth_to_latent(
	next_ar_depth, norm_method=depth_norm_method, depth_range=depth_range,
	).squeeze(0)
	# hv = history (none at step 0) + very first (frame 0's depth)
	hv_vf = meta["next_ar_depth"][0].cpu().numpy().astype(np.float32)
	hv_depth = torch.from_numpy(hv_vf[None, ...]).to(DEVICE)
	func_input["latents_hv_depth"] = encode_depth_to_latent(
	hv_depth, norm_method=depth_norm_method, depth_range=depth_range,
	).squeeze(0)
	func_input["raw_next_ar_depth"] = [next_ar_depth]
	if num_ctx > 0 and meta.get("context_depth", None) is not None and meta["context_depth"].numel() > 0:
	func_input["latents_context_depth"] = encode_depth_to_latent(
	meta["context_depth"].to(DEVICE),
	norm_method=depth_norm_method, depth_range=depth_range,
	).squeeze(0)

	K_list = []
	for intr in meta["intrinsics"]:
	K_list.append(torch.tensor([
	[intr["fx"], 0.0, intr["cx"]],
	[0.0, intr["fy"], intr["cy"]],
	[0.0, 0.0, 1.0],
	], dtype=torch.float32))
	intrinsics = torch.stack(K_list).to(DEVICE)

	extrinsics_c2w_full = torch.tensor(meta["extrinsics_c2w"], dtype=torch.float32).to(DEVICE)
	if extrinsics_c2w_full.shape[0] < SEGMENT_SIZE: # reference pads short sequences
	pad = SEGMENT_SIZE - extrinsics_c2w_full.shape[0]
	extrinsics_c2w_full = torch.cat(
	[extrinsics_c2w_full, extrinsics_c2w_full[-1:].repeat(pad, 1, 1)], dim=0
	)
	intrinsics = intrinsics[:SEGMENT_SIZE]
	while intrinsics.shape[0] < SEGMENT_SIZE:
	intrinsics = torch.cat([intrinsics, intrinsics[-1:]], dim=0)
	func_input["intrinsics"] = intrinsics
	func_input["extrinsics_c2w"] = extrinsics_c2w_full[:SEGMENT_SIZE]

	# This VAE maps frames 1:1; pass the count the pipeline maps back to
	# segment_size latent frames (Wan temporal convention: (n-1)*4+1).
	func_input["num_frames"] = (func_input["num_frames"] - 1) * 4 + 1

	progress((4, 12), desc="denoising (UniPC)")
	with torch.no_grad():
	pipeline_output = pipeline(**func_input)
	videos = pipeline_output.frames
	plane_params = getattr(pipeline_output, "plane_params", None)

	progress((10, 12), desc="export")
	out_dir = work / "out"
	out_dir.mkdir(parents=True, exist_ok=True)
	video_path = out_dir / "generated_video.mp4"
	export_to_video(videos[0], str(video_path), fps=FPS)

	depth_out = None
	normal_paths = []
	ply_path = None
	if plane_params is not None:
	extr = plane_params.get("extrinsics_c2w", None)
	pps = plane_params.get("plane_params_per_video", [])
	if extr is not None and len(pps):
	pp = pps[0]
	extr_t = extr if torch.is_tensor(extr) else torch.tensor(np.asarray(extr), dtype=torch.float32)
	verts_w, faces_w, colors_w, vnormals_w = planes_to_mesh(pp, extr_t.to(pp["plane_center_camera"].device))
	ply_path = out_dir / "planes_world.ply"
	save_planes_as_ply(str(ply_path), verts_w, faces_w, colors_w, vnormals_w)
	ply_path = str(ply_path)

	d_dir = out_dir / "cam_0_depth"
	save_depth_maps(str(d_dir), pp)
	depth_out = _depth_video_file(Path(d_dir))

	n_dir = out_dir / "cam_0_normals"
	save_normal_maps(str(n_dir), pp)
	import glob as _g

	normal_paths = sorted(_g.glob(str(n_dir / ".png")))[:: max(1, len(sorted(_g.glob(str(n_dir / ".png")))) // 4)][:4]

	logger.info(f"[demo] generate done in {time.time()-t0:.1f}s -> {video_path}")

	if depth_out is None:
	depth_out = str(video_path) # keep the slot filled if geometry was skipped
	return str(video_path), depth_out, normal_paths, ply_path


	# ----------------------------------------------------------------------
	# UI
	# ----------------------------------------------------------------------

	examples_dir = Path(__file__).parent / "examples"
	_scenes = sorted(p.name for p in examples_dir.glob("scene*") if p.is_dir())
	EXAMPLES = []
	for scene in _scenes:
	base = examples_dir / scene
	prompt_text = (base / "prompt.txt").read_text().strip() if (base / "prompt.txt").exists() else "A video of an indoor scene."
	EXAMPLES.append([
	str(base / "input.png"),
	str(base / "camera.npy"),
	str(base / "geometry.npz"),
	prompt_text,
	])

	with gr.Blocks(title="PPMem: Persistent Planar Memory for Video World Models") as demo:
	gr.Markdown(
	"""
	# PPMem: Persistent Planar Memory for Video World Models
	Paper [2610.23795](https://huggingface.co/papers/2610.23795) · Code [robbyant-research/PPMem](https://github.com/robbyant-research/PPMem) · Weights [Robbyant-Research/PPMem](https://huggingface.co/Robbyant-Research/PPMem)

	An interactive video world model: give it a starting image and a camera trajectory, and it generates a 25-frame 1024×576 video from the first pose — plus the 3D planar geometry it predicts along that camera: depth maps, normal maps, and a world-space planar mesh.
	"""
	)
	with gr.Row():
	with gr.Column():
	image_in = gr.Image(label="Starting image", type="filepath", height=280)
	camera_in = gr.File(label="Camera trajectory (.npy, (N,4,4) c2w poses) — optional; a gentle dolly is used if empty", file_types=[".npy"])
	geometry_in = gr.File(label="Geometry .npz (depth, fx, fy, cx, cy) — optional; MoGe-2 estimates it if empty", file_types=[".npz"])
	prompt_in = gr.Textbox(label="Prompt", lines=3, placeholder="Describe the scene…")
	with gr.Accordion("Advanced", open=False):
	seed_in = gr.Number(label="Seed", value=DEFAULT_SEED, precision=0)
	steps_in = gr.Slider(label="Inference steps", minimum=10, maximum=50, value=DEFAULT_STEPS, step=1)
	guidance_in = gr.Slider(label="Guidance scale", minimum=1.0, maximum=10.0, value=DEFAULT_GUIDANCE, step=0.5)
	run_btn = gr.Button("Generate video", variant="primary")
	with gr.Column():
	video_out = gr.Video(label="Generated RGB video (1024×576, 25 frames @ 15fps)")
	depth_out = gr.Video(label="Predicted depth (turbo, from planar geometry)")
	normals_out = gr.Gallery(label="Predicted normals (sample frames)", columns=4, height=180)
	planes_out = gr.Model3D(label="World-space planar mesh (PLY)")

	run_btn.click(
	generate,
	inputs=[image_in, camera_in, geometry_in, prompt_in, seed_in, steps_in, guidance_in],
	outputs=[video_out, depth_out, normals_out, planes_out],
	api_name="generate",
	)

	if EXAMPLES:
	gr.Examples(
	examples=EXAMPLES,
	inputs=[image_in, camera_in, geometry_in, prompt_in],
	fn=generate,
	outputs=[video_out, depth_out, normals_out, planes_out],
	cache_examples=True,
	cache_mode="lazy",
	)

	if __name__ == "__main__":
	demo.launch(mcp_server=True, theme=gr.themes.Citrus())