Wan 2.1 on MAX/Mojo

Wan 2.1 T2V-1.3B is a text-to-video diffusion model (a diffusion transformer paired with a 3D VAE and a text encoder) whose inference pipeline spans text encoding, iterative denoising, and video decoding on the GPU. Porting it off PyTorch means re-expressing every model component on a different ML stack: Modular's MAX graph and inference APIs, with custom Mojo kernels where the built-in operations fall short.

The port is measured against the original: the agent's pipeline must reproduce the PyTorch reference's output frames to high pixel fidelity, working entirely offline from the reference source, the pretrained weights, and the bundled MAX API documentation. PyTorch itself is not installed; the target stack is the only stack available.

Implement Wan 2.1 T2V-1.3B text-to-video generation on Modular's MAX/Mojo stack. The deliverable is the Python package at /app/wan21_max/, whose entrypoint module wan_pipeline.py must export generate_video(prompt, height, width, num_frames, num_steps, seed) returning a list of PIL Images, one per frame.

• The signature stays stable and everything derives from the arguments; sources are .py modules (plus .mojo kernels for custom ops) under /app/wan21_max/.
• Everything the pipeline needs at generation time must live in the package: /app/reference/ and the .py files under /app/weights/ are removed before scoring.
• The port must be built on the MAX SDK; torch, transformers, and diffusers are hard library constraints: they are not installed and must not be used.
• A short generation (5 frames, 4 steps) must complete within 10 minutes including any first-call compilation.

The package scaffold at /app/wan21_max/ and an /app/README.md spelling out the exact API, the workload family, what counts as correct output, and the library constraints. For porting: the Wan 2.1 PyTorch reference source at /app/reference/ (read-only), the pretrained T2V-1.3B weights in diffusers format at /app/weights/, and the Modular MAX API documentation at /app/max_docs/. For self-checking: sample workloads with baked reference frames under /app/examples/ and /app/verify_correctness.py, which generates each sample workload, compares output against the reference frames, and enforces the per-generation time budget. The machine has one H100 GPU, is offline, ships the MAX SDK and Mojo toolchain with warmed compile caches, and exposes a sandbox-timer CLI for the remaining budget.

A clean-room verifier runs in a fresh container on the captured workspace. It reconstructs the scored package from a pristine scaffold plus only the agent's wan21_max/ files into a root-owned, read-only directory, strips the reference source and stray Python files as disclosed, then runs a gate sequence: a source-level anti-cheat verdict (banned imports, verifier-internals references, shell or dynamic-import reach, and a required MAX SDK import), an import check, and the disclosed smoke gate. Only then does it generate frames for nine hidden workloads as a non-root user. Root then scores the saved frames against root-only reference frames; no candidate code runs during scoring.

• Per workload, validity gates require the exact frame count and size, no symlinked frames (which would otherwise echo back the references), and no blank frames; credit is then a continuous ramp on mean per-frame PSNR: 0 at 10 dB, 1.0 at the disclosed 25 dB bar.
• The reward is the geometric mean of per-workload credits, with a small 0.02 per-workload floor so a single hard workload cannot zero an otherwise strong port.
• Banned-library imports, a missing MAX import, an unimportable entrypoint, or a failed smoke test (crash, timeout, or blank frames) zero the result; workloads that miss the generation stage's start deadline degrade to partial credit rather than erroring the run.