diff --git a/.env.example b/.env.example index 1b4e9090..4dfd9c5b 100644 --- a/.env.example +++ b/.env.example @@ -100,7 +100,7 @@ VOLC_ACCESSKEY= VOLC_SECRETKEY= # Set to "true" for local video gen (needs GPU + diffusers). VIDEO_GEN_LOCAL_ENABLED= -# Local model: wan2.1-1.3b, wan2.1-14b, hunyuan-1.5, ltx2-local, cogvideo-5b. +# Local model: wan2.2-ti2v-5b, wan2.1-1.3b, wan2.1-14b, hunyuan-1.5, ltx2-local, cogvideo-5b. VIDEO_GEN_LOCAL_MODEL= # Modal self-hosted LTX-2 endpoint (optional). MODAL_LTX2_ENDPOINT_URL= diff --git a/.gitignore b/.gitignore index 85473523..62e3307e 100644 --- a/.gitignore +++ b/.gitignore @@ -101,3 +101,6 @@ venv/ # Workspace-local third-party runtimes (for example the portable Blender LTS # used by blender_world). These are checksum-verified but never committed. .runtime/ + +# Local Wan / diffusers GPU environment +.venv-wan/ diff --git a/README.md b/README.md index d8007737..c6b2c132 100644 --- a/README.md +++ b/README.md @@ -272,7 +272,7 @@ make install-gpu # Then add to .env: VIDEO_GEN_LOCAL_ENABLED=true -VIDEO_GEN_LOCAL_MODEL=wan2.1-1.3b # or wan2.1-14b, hunyuan-1.5, ltx2-local, cogvideo-5b +VIDEO_GEN_LOCAL_MODEL=wan2.2-ti2v-5b # or wan2.1-1.3b, wan2.1-14b, hunyuan-1.5, ltx2-local, cogvideo-5b ``` diff --git a/README_zh-CN.md b/README_zh-CN.md index 9c6d27e9..423ab86b 100644 --- a/README_zh-CN.md +++ b/README_zh-CN.md @@ -200,7 +200,7 @@ make install-gpu # 然后添加到 .env: VIDEO_GEN_LOCAL_ENABLED=true -VIDEO_GEN_LOCAL_MODEL=wan2.1-1.3b # 或 wan2.1-14b, hunyuan-1.5, ltx2-local, cogvideo-5b +VIDEO_GEN_LOCAL_MODEL=wan2.2-ti2v-5b # 或 wan2.1-1.3b, wan2.1-14b, hunyuan-1.5, ltx2-local, cogvideo-5b ``` diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index 9c03f9a3..1a59d4a5 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -73,7 +73,7 @@ TENCENT_TOKENHUB_API_KEY= # Tencent Hunyuan cloud video via TokenHub API # LOCAL (no keys needed — just GPU + install) VIDEO_GEN_LOCAL_ENABLED= # Set to "true" for local video gen -VIDEO_GEN_LOCAL_MODEL= # wan2.1-1.3b, wan2.1-14b, hunyuan-1.5, ltx2-local, cogvideo-5b +VIDEO_GEN_LOCAL_MODEL= # wan2.2-ti2v-5b, wan2.1-1.3b, wan2.1-14b, hunyuan-1.5, ltx2-local, cogvideo-5b # COMFYUI (optional overrides; localhost:8188 is the default) COMFYUI_SERVER_URL= # Local ComfyUI server for shared workflows @@ -1261,8 +1261,9 @@ pip install diffusers transformers accelerate torch pillow requests VIDEO_GEN_LOCAL_ENABLED=true # 3. Choose a model based on your GPU VRAM -VIDEO_GEN_LOCAL_MODEL=wan2.1-1.3b # 6GB+ VRAM (entry-level) -VIDEO_GEN_LOCAL_MODEL=wan2.1-14b # 24GB+ VRAM (best local quality) +VIDEO_GEN_LOCAL_MODEL=wan2.2-ti2v-5b # 12GB+ VRAM (default; 720p @ 24fps, T2V+I2V+V2V) +VIDEO_GEN_LOCAL_MODEL=wan2.1-1.3b # 6GB+ VRAM (entry-level, text-to-video only) +VIDEO_GEN_LOCAL_MODEL=wan2.1-14b # 24GB+ VRAM (best Wan 2.1 quality) VIDEO_GEN_LOCAL_MODEL=hunyuan-1.5 # 12GB+ VRAM VIDEO_GEN_LOCAL_MODEL=ltx2-local # 8GB+ VRAM (fastest) VIDEO_GEN_LOCAL_MODEL=cogvideo-5b # 10GB+ VRAM @@ -1273,14 +1274,76 @@ VIDEO_GEN_LOCAL_MODEL=cogvideo-2b # 6GB+ VRAM (lightest) | Model | VRAM | Quality | Speed | Best for | |-------|------|---------|-------|----------| -| **WAN 2.1 (1.3B)** | 6GB | Good | Fast | Entry-level GPU, quick iteration | -| **WAN 2.1 (14B)** | 24GB | Excellent | Slow | Best quality-to-VRAM ratio | +| **WAN 2.2 TI2V (5B)** | 12GB | Excellent | Medium | Default. 720p @ 24fps, one checkpoint for T2V/I2V/V2V | +| **WAN 2.1 (1.3B)** | 6GB | Good | Fast | Entry-level GPU, quick iteration (text-to-video only) | +| **WAN 2.1 (14B)** | 24GB | Excellent | Slow | Best Wan 2.1 quality | | **Hunyuan 1.5** | 12GB | Very good | Medium | Mid-range GPUs | | **LTX-2** | 8GB | Good | Fastest | Quick drafts, lowest latency | | **CogVideo (5B)** | 10GB | Good | Medium | Balanced option | | **CogVideo (2B)** | 6GB | Fair | Fast | Low-VRAM experimentation | -**All local models support:** Image-to-video, text-to-video, offline generation, seeded reproducibility. +**All local models support:** text-to-video, offline generation, seeded reproducibility. + +#### WAN 2.2 — the full operation matrix + +`wan_video` drives one checkpoint across every Wan task. The `operation` input selects +the diffusers pipeline; the tool refuses an operation the chosen variant has no weights for. + +| `operation` | What it does | Extra inputs | +|-------------|--------------|--------------| +| `text_to_video` | Prompt to clip | — | +| `image_to_video` | Animate a still | `reference_image_path` / `reference_image_url` | +| `video_to_video` | Restyle an existing clip | `source_video_path`, `strength` | +| `first_last_frame` | Interpolate between two stills | `reference_image_*` + `last_image_*` | +| `text_to_image` | Single frame, saved as PNG | — | + +Geometry is snapped to what the VAE can encode, so you can ask for any size: frame counts +land on `4k + 1`, and width/height land on a multiple of 32 for the TI2V line (its 16x VAE +plus 2x patching is why 720p is **1280x704**, not 1280x720) or 16 elsewhere. + +#### Clips longer than five seconds + +Wan is trained for roughly a five-second pass. Ask for more with `duration_seconds` and the +tool chains segments: each continuation is seeded with the previous segment's last frame, +regenerates it as its own frame 0, and that duplicate is dropped at the seam. + +```python +wan_video.execute({ + "prompt": "A red convertible on a coastal road at golden hour", + "duration_seconds": 30, # -> 720 frames @ 24fps -> 6 chained segments + "width": 704, "height": 480, + "num_inference_steps": 20, + "offload_mode": "sequential", + "segment_prompts": [...], # optional shot list, one entry per segment +}) +``` + +Chaining drifts: every hop conditions on already-generated pixels, so contrast and colour +creep. `correct_drift` (on by default) re-grades each segment to the opening segment's +channel statistics. Pass `segment_prompts` to steer the clip through a shot list instead of +repeating one prompt six times. + +#### Fitting a 12GB card + +`precision` (`auto`/`bf16`/`int8`/`int4`) and `offload_mode` (`auto`/`model`/`sequential`) +both default to `auto` and are picked from visible VRAM. + +A bf16 5B transformer is ~10GB resident. With `offload_mode="model"` that leaves almost +nothing for activations on a 12GB card — measured peak was 10.85GB at just 512x320. +`offload_mode="sequential"` streams one submodule at a time: the same card then peaks at +**1.97GB** and renders 704x480 comfortably, at the cost of ~4s per denoising step of PCIe +traffic. On a 12GB GPU `auto` selects it for you. + +> **NVIDIA driver note.** If the loaded kernel module and the installed userspace libraries +> are different versions, `nvmlInit` fails. CUDA compute still runs, but PyTorch calls NVML +> while composing an out-of-memory report — so real OOMs surface as an internal assert, and +> bitsandbytes (`int8`/`int4`) cannot load at all. Compare `cat /proc/driver/nvidia/version` +> with `nvidia-smi`; if they disagree, reboot. `wan_video` detects this and says so. + + +Image-to-video is *not* universal — Wan 2.1 never shipped 1.3B I2V weights and CogVideo is +text-only. Each variant's real capability list is `operations` in `tools/video/_shared.py`, +and `wan_video` refuses an operation the selected checkpoint cannot do. --- diff --git a/scripts/wan_disk_cleanup.sh b/scripts/wan_disk_cleanup.sh new file mode 100755 index 00000000..f05503d5 --- /dev/null +++ b/scripts/wan_disk_cleanup.sh @@ -0,0 +1,115 @@ +#!/usr/bin/env bash +# +# Reclaim disk used by local Wan video generation. +# +# Local generation is free of API cost but not of disk: the Wan 2.2 TI2V-5B +# checkpoint alone is ~32GB, the CUDA-enabled virtualenv is ~7GB, and uv keeps +# every wheel it has ever downloaded. +# +# Nothing is deleted unless you pass a target. Without --yes it only reports. +# +# ./scripts/wan_disk_cleanup.sh # show what is on disk +# ./scripts/wan_disk_cleanup.sh --models --yes # delete Wan model weights +# ./scripts/wan_disk_cleanup.sh --all --yes # weights + venv + uv cache +# +# Deleted weights are re-downloaded automatically on the next run, so this is +# safe to run whenever you need the space back. + +set -euo pipefail + +HF_HOME_DIR="${HF_HOME:-$HOME/.cache/huggingface}" +HUB_DIR="$HF_HOME_DIR/hub" +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +VENV_DIR="$REPO_ROOT/.venv-wan" +UV_CACHE_DIR="${UV_CACHE_DIR:-$HOME/.cache/uv}" + +# Only these are touched by --models. Other repos in the shared HF cache +# (OCR models, LLMs, anything another project put there) are left alone. +WAN_PATTERNS=( + "models--Wan-AI--*" + "models--FastVideo--FastWan*" + "models--lightx2v--Wan*" + "models--Comfy-Org--Wan*" + "models--QuantStack--Wan*" +) + +DO_MODELS=0; DO_VENV=0; DO_UV=0; CONFIRM=0 + +for arg in "$@"; do + case "$arg" in + --models) DO_MODELS=1 ;; + --venv) DO_VENV=1 ;; + --uv-cache) DO_UV=1 ;; + --all) DO_MODELS=1; DO_VENV=1; DO_UV=1 ;; + --yes|-y) CONFIRM=1 ;; + -h|--help) sed -n '2,20p' "$0" | sed 's/^# \?//'; exit 0 ;; + *) echo "Unknown option: $arg (try --help)" >&2; exit 2 ;; + esac +done + +human () { du -sh "$1" 2>/dev/null | cut -f1; } + +# Expand the Wan patterns into real paths. +wan_paths () { + local pattern path + for pattern in "${WAN_PATTERNS[@]}"; do + for path in "$HUB_DIR"/$pattern; do + [ -e "$path" ] && printf '%s\n' "$path" + done + done +} + +echo "=== Current usage ===" +if [ -d "$HUB_DIR" ]; then + total_wan=0 + while IFS= read -r path; do + printf ' %-8s %s\n' "$(human "$path")" "$(basename "$path")" + total_wan=$(( total_wan + $(du -sk "$path" 2>/dev/null | cut -f1) )) + done < <(wan_paths) + [ "$total_wan" -gt 0 ] \ + && printf ' %-8s Wan model weights (total)\n' "$(( total_wan / 1024 / 1024 ))G" \ + || echo " (no Wan weights cached)" + printf ' %-8s entire HF cache at %s\n' "$(human "$HF_HOME_DIR")" "$HF_HOME_DIR" +fi +[ -d "$VENV_DIR" ] && printf ' %-8s %s\n' "$(human "$VENV_DIR")" "$VENV_DIR" +[ -d "$UV_CACHE_DIR" ] && printf ' %-8s %s\n' "$(human "$UV_CACHE_DIR")" "$UV_CACHE_DIR" +echo + +if [ $((DO_MODELS + DO_VENV + DO_UV)) -eq 0 ]; then + echo "Nothing selected. Pass --models, --venv, --uv-cache or --all (add --yes to delete)." + exit 0 +fi + +targets=() +[ "$DO_MODELS" -eq 1 ] && while IFS= read -r path; do targets+=("$path"); done < <(wan_paths) +[ "$DO_VENV" -eq 1 ] && [ -d "$VENV_DIR" ] && targets+=("$VENV_DIR") +[ "$DO_UV" -eq 1 ] && [ -d "$UV_CACHE_DIR" ] && targets+=("$UV_CACHE_DIR") + +if [ "${#targets[@]}" -eq 0 ]; then + echo "Nothing to delete." + exit 0 +fi + +echo "=== Would delete ===" +for path in "${targets[@]}"; do printf ' %-8s %s\n' "$(human "$path")" "$path"; done +echo + +if [ "$CONFIRM" -ne 1 ]; then + echo "Dry run. Re-run with --yes to actually delete." + exit 0 +fi + +for path in "${targets[@]}"; do + echo "removing $path" + rm -rf -- "$path" +done + +echo +echo "Done. Free space now:" +df -h "$HOME" | tail -1 +echo +echo "To rebuild: model weights re-download on the next wan_video run;" +echo "recreate the venv with" +echo " uv venv --python 3.12 .venv-wan" +echo " uv pip install --python .venv-wan/bin/python torch torchvision --index-url https://download.pytorch.org/whl/cu128" +echo " uv pip install --python .venv-wan/bin/python 'diffusers>=0.36' transformers accelerate safetensors sentencepiece imageio imageio-ffmpeg ftfy protobuf" diff --git a/tests/contracts/test_wan_video.py b/tests/contracts/test_wan_video.py new file mode 100644 index 00000000..b932ebc5 --- /dev/null +++ b/tests/contracts/test_wan_video.py @@ -0,0 +1,509 @@ +"""Contract tests for the local Wan 2.x video tool. + +These run without a GPU, without model weights and without network: they cover +the tool contract, the variant catalog's internal consistency, and the frame / +segment arithmetic that decides how a long clip is chained together. +""" + +import pytest + +from tools.base_tool import ( + BaseTool, + ToolRuntime, + ToolStability, + ToolStatus, + ToolTier, +) +from tools.tool_registry import ToolRegistry +from tools.video._shared import WAN_VARIANTS +from tools.video._wan_engine import ( + _PIPELINE_FOR_OPERATION, + _oom_aware_error, + generate_wan, + plan_segments, + resolve_offload_mode, + resolve_precision, + snap_dimension, + snap_frames, +) +from tools.video.wan_video import DEFAULT_VARIANT, WanVideo + +# Wan's VAE compresses 4x in time. Spatial alignment is per-variant: the VAE's +# spatial compression times the transformer patch size, which is 32 for the +# 16x-VAE TI2V line and 16 for every 8x-VAE Wan line. +TEMPORAL_SCALE = 4 + + +# ------------------------------------------------------------------ +# Contract compliance +# ------------------------------------------------------------------ + + +class TestContract: + def test_inherits_base_tool(self): + assert issubclass(WanVideo, BaseTool) + + def test_identity(self): + tool = WanVideo() + assert tool.name == "wan_video" + assert tool.provider == "wan" + assert tool.capability == "video_generation" + assert tool.tier == ToolTier.GENERATE + assert tool.stability == ToolStability.EXPERIMENTAL + assert tool.runtime == ToolRuntime.LOCAL_GPU + + def test_input_schema(self): + schema = WanVideo().input_schema + assert schema["type"] == "object" + assert "prompt" in schema["required"] + assert schema["properties"]["model_variant"]["default"] == DEFAULT_VARIANT + + def test_local_generation_is_free(self): + assert WanVideo().estimate_cost({"prompt": "x"}) == 0.0 + + def test_status_requires_opt_in(self, monkeypatch): + monkeypatch.delenv("VIDEO_GEN_LOCAL_ENABLED", raising=False) + assert WanVideo().get_status() == ToolStatus.UNAVAILABLE + + def test_registry_discovers_tool(self): + registry = ToolRegistry() + registry.discover() + assert "wan_video" in registry.list_all() + assert registry.get("wan_video").provider == "wan" + + def test_schema_operations_match_engine(self): + declared = set(WanVideo().input_schema["properties"]["operation"]["enum"]) + assert declared == set(_PIPELINE_FOR_OPERATION) + + +# ------------------------------------------------------------------ +# Variant catalog +# ------------------------------------------------------------------ + + +@pytest.mark.parametrize("key", sorted(WAN_VARIANTS), ids=lambda k: k) +class TestVariantCatalog: + def test_required_fields(self, key): + meta = WAN_VARIANTS[key] + for field in ( + "name", + "hf_id", + "params_b", + "operations", + "default_width", + "default_height", + "default_num_frames", + "spatial_alignment", + "temporal_scale", + "default_steps", + "default_guidance", + "fps", + "license", + ): + assert field in meta, f"{key} is missing {field}" + + def test_operations_are_known_and_non_empty(self, key): + operations = WAN_VARIANTS[key]["operations"] + assert operations + assert set(operations) <= set(_PIPELINE_FOR_OPERATION) + + def test_default_geometry_is_vae_aligned(self, key): + meta = WAN_VARIANTS[key] + alignment = meta["spatial_alignment"] + assert meta["default_width"] % alignment == 0 + assert meta["default_height"] % alignment == 0 + + def test_declares_vae_scales(self, key): + meta = WAN_VARIANTS[key] + assert meta["spatial_alignment"] in (16, 32) + assert meta["temporal_scale"] == TEMPORAL_SCALE + + def test_default_frame_count_is_vae_aligned(self, key): + frames = WAN_VARIANTS[key]["default_num_frames"] + assert (frames - 1) % TEMPORAL_SCALE == 0 + + def test_image_operations_have_a_checkpoint(self, key): + """A variant may only claim an operation it has weights for.""" + meta = WAN_VARIANTS[key] + for operation in meta["operations"]: + assert meta.get(f"hf_{operation}_id") or meta["hf_id"] + + def test_legacy_flags_agree_with_operations(self, key): + meta = WAN_VARIANTS[key] + assert meta["t2v"] is ("text_to_video" in meta["operations"]) + assert meta["i2v"] is ("image_to_video" in meta["operations"]) + + +def test_wan21_1_3b_does_not_claim_image_to_video(): + """Wan 2.1 never shipped 1.3B I2V weights; claiming it pulled a 14B model.""" + meta = WAN_VARIANTS["wan2.1-1.3b"] + assert "image_to_video" not in meta["operations"] + assert "1.3B" in meta["hf_id"] or "1.3B" in meta["name"] + + +def test_default_variant_covers_the_full_matrix(): + """The default must handle every operation the tool advertises.""" + assert set(WAN_VARIANTS[DEFAULT_VARIANT]["operations"]) == set(_PIPELINE_FOR_OPERATION) + + +# ------------------------------------------------------------------ +# Frame and segment arithmetic +# ------------------------------------------------------------------ + + +class TestGeometry: + @pytest.mark.parametrize("value", [1, 5, 25, 80, 81, 100, 121, 720]) + def test_snap_frames_lands_on_a_latent_boundary(self, value): + snapped = snap_frames(value) + assert snapped == 1 or (snapped - 1) % TEMPORAL_SCALE == 0 + assert snapped >= 1 + + def test_snap_frames_is_idempotent(self): + assert snap_frames(snap_frames(97)) == snap_frames(97) + + def test_snap_frames_rejects_zero(self): + with pytest.raises(ValueError): + snap_frames(0) + + @pytest.mark.parametrize("value", [17, 300, 480, 704, 1280]) + def test_snap_dimension_aligns(self, value): + assert snap_dimension(value, 32) % 32 == 0 + + def test_snap_dimension_never_returns_zero(self): + assert snap_dimension(3, 32) == 32 + + +class TestSegmentPlanning: + @pytest.mark.parametrize("segment", [25, 49, 121]) + def test_plan_covers_every_requested_length(self, segment): + """Segments must always reach the target; the tail is trimmed, never short. + + Swept rather than spot-checked: an off-by-one in the carry (using the + segment length instead of segment-1) still lands on the right answer for + most inputs and only shows up at particular totals. + """ + for total in range(1, 40 * segment): + count = plan_segments(total, segment) + covered = segment + (count - 1) * (segment - 1) + assert covered >= total, f"{total} frames left short by {total - covered}" + # and must never buy a segment it does not need + if count > 1: + previous = segment + (count - 2) * (segment - 1) + assert previous < total + + def test_carry_accounts_for_the_duplicated_anchor_frame(self): + """Each continuation reproduces the anchor, so it adds segment-1 frames.""" + segment = 121 + # 1210 is an exact multiple of the segment length but not of the carry, + # which is where an off-by-one in the carry first becomes visible. + total = segment + 1210 + count = plan_segments(total, segment) + assert segment + (count - 1) * (segment - 1) >= total + + def test_single_segment_when_it_fits(self): + assert plan_segments(121, 121) == 1 + assert plan_segments(50, 121) == 1 + + def test_thirty_seconds_at_24fps(self): + """The headline case: 30s of 24fps video off a 121-frame model.""" + count = plan_segments(30 * 24, 121) + assert count == 6 + assert 121 + (count - 1) * 120 >= 720 + + +class TestPrecisionSelection: + def test_explicit_precision_is_respected(self): + assert resolve_precision("int4", 5.0) == "int4" + assert resolve_precision("bf16", 28.0) == "bf16" + + def test_auto_degrades_as_the_model_grows(self, monkeypatch): + monkeypatch.setattr("tools.video._wan_engine._vram_gb", lambda: 12.0) + assert resolve_precision("auto", 1.3) == "bf16" + assert resolve_precision("auto", 28.0) == "int4" + + def test_auto_on_a_big_card_stays_bf16(self, monkeypatch): + monkeypatch.setattr("tools.video._wan_engine._vram_gb", lambda: 80.0) + assert resolve_precision("auto", 28.0) == "bf16" + + def test_cpu_never_asks_for_a_cuda_only_quantizer(self, monkeypatch): + monkeypatch.setattr("tools.video._wan_engine._vram_gb", lambda: 0.0) + assert resolve_precision("auto", 28.0) == "bf16" + + +# ------------------------------------------------------------------ +# Planning surface on the tool +# ------------------------------------------------------------------ + + +class TestPlan: + def test_duration_drives_segment_count(self): + plan = WanVideo()._plan({"prompt": "x", "duration_seconds": 30}) + assert plan["fps"] == 24 + assert plan["total_frames"] == 720 + assert plan["segments"] == 6 + + def test_native_length_needs_one_segment(self): + plan = WanVideo()._plan({"prompt": "x"}) + assert plan["segments"] == 1 + + def test_explicit_num_frames_overrides_duration(self): + plan = WanVideo()._plan({"prompt": "x", "duration_seconds": 30, "num_frames": 49}) + assert plan["total_frames"] == 49 + assert plan["segments"] == 1 + + def test_text_to_image_is_a_single_frame(self): + plan = WanVideo()._plan({"prompt": "x", "operation": "text_to_image", "duration_seconds": 30}) + assert plan["total_frames"] == 1 + assert plan["segments"] == 1 + + def test_longer_request_estimates_longer_runtime(self): + tool = WanVideo() + short = tool.estimate_runtime({"prompt": "x", "duration_seconds": 5}) + long = tool.estimate_runtime({"prompt": "x", "duration_seconds": 30}) + assert long > short + + +# ------------------------------------------------------------------ +# Validation happens before any weights are touched +# ------------------------------------------------------------------ + + +def _generate(**inputs): + return generate_wan( + tool_name="wan_video", + variants=WAN_VARIANTS, + default_variant=DEFAULT_VARIANT, + inputs={"prompt": "a cat", **inputs}, + ) + + +class TestValidation: + def test_unknown_variant_is_rejected(self): + result = _generate(model_variant="wan9000") + assert not result.success + assert "Unknown model_variant" in result.error + + def test_unknown_operation_is_rejected(self): + result = _generate(operation="mind_to_video") + assert not result.success + assert "Unknown operation" in result.error + + def test_operation_unsupported_by_variant_is_rejected(self): + """wan2.1-1.3b has no I2V weights — this must fail fast, not load 14B.""" + result = _generate(model_variant="wan2.1-1.3b", operation="image_to_video") + assert not result.success + assert "does not support" in result.error + + def test_missing_prompt_is_rejected(self): + result = generate_wan( + tool_name="wan_video", + variants=WAN_VARIANTS, + default_variant=DEFAULT_VARIANT, + inputs={"prompt": ""}, + ) + assert not result.success + assert "prompt is required" in result.error + + def test_unavailable_tool_reports_install_instructions(self, monkeypatch): + monkeypatch.delenv("VIDEO_GEN_LOCAL_ENABLED", raising=False) + result = WanVideo().execute({"prompt": "a cat"}) + assert not result.success + assert "VIDEO_GEN_LOCAL_ENABLED" in result.error + + +def test_ti2v_is_the_only_32_aligned_line(): + """The 16x VAE plus 2x patching is what forces 1280x704 instead of 1280x720.""" + assert WAN_VARIANTS["wan2.2-ti2v-5b"]["spatial_alignment"] == 32 + assert WAN_VARIANTS["wan2.2-ti2v-5b"]["default_height"] == 704 + for key, meta in WAN_VARIANTS.items(): + if key != "wan2.2-ti2v-5b": + assert meta["spatial_alignment"] == 16 + + +# ------------------------------------------------------------------ +# Offload strategy and allocator diagnostics +# ------------------------------------------------------------------ + + +class TestOffloadSelection: + def test_explicit_mode_is_respected(self): + assert resolve_offload_mode("model", "bf16", 5.0) == "model" + assert resolve_offload_mode("sequential", "bf16", 1.3) == "sequential" + + def test_bf16_5b_on_a_12gb_card_goes_sequential(self, monkeypatch): + """10GB of resident bf16 weights leaves no room for activations.""" + monkeypatch.setattr("tools.video._wan_engine._vram_gb", lambda: 11.6) + assert resolve_offload_mode("auto", "bf16", 5.0) == "sequential" + + def test_quantized_weights_fit_resident_on_the_same_card(self, monkeypatch): + monkeypatch.setattr("tools.video._wan_engine._vram_gb", lambda: 11.6) + assert resolve_offload_mode("auto", "int4", 5.0) == "model" + + def test_big_card_keeps_weights_resident(self, monkeypatch): + monkeypatch.setattr("tools.video._wan_engine._vram_gb", lambda: 80.0) + assert resolve_offload_mode("auto", "bf16", 28.0) == "model" + + def test_cpu_only_does_not_ask_for_sequential(self, monkeypatch): + monkeypatch.setattr("tools.video._wan_engine._vram_gb", lambda: 0.0) + assert resolve_offload_mode("auto", "bf16", 28.0) == "model" + + +class TestAllocatorDiagnostics: + def test_nvml_assert_is_reported_as_a_driver_mismatch(self): + """PyTorch calls NVML while reporting OOM, so a driver mismatch masks it.""" + result = _oom_aware_error( + RuntimeError( + "NVML_SUCCESS == DriverAPI::get()->nvmlInit_v2_() INTERNAL ASSERT FAILED" + ), + "wan2.2-ti2v-5b", "bf16", "sequential", 704, 480, 121, + ) + assert not result.success + assert "driver" in result.error.lower() + assert "reboot" in result.error.lower() + + def test_real_oom_suggests_sequential_offload_first(self): + result = _oom_aware_error( + RuntimeError("CUDA out of memory. Tried to allocate 2.00 GiB"), + "wan2.2-ti2v-5b", "bf16", "model", 1280, 704, 121, + ) + assert not result.success + assert "sequential" in result.error + + def test_oom_already_sequential_does_not_suggest_it_again(self): + result = _oom_aware_error( + RuntimeError("CUDA out of memory"), + "wan2.2-ti2v-5b", "bf16", "sequential", 1280, 704, 121, + ) + assert 'offload_mode="sequential"' not in result.error + assert "resolution" in result.error + + def test_unrelated_failure_is_passed_through(self): + result = _oom_aware_error( + ValueError("something else entirely"), + "wan2.2-ti2v-5b", "bf16", "model", 704, 480, 49, + ) + assert "something else entirely" in result.error + + +class TestRuntimeEstimate: + def test_scales_with_segment_count(self): + """Each chained segment pays the full per-step cost again.""" + tool = WanVideo() + base = dict(prompt="x", width=704, height=480, num_inference_steps=20, + offload_mode="sequential") + one = tool.estimate_runtime({**base, "num_frames": 121}) + six = tool.estimate_runtime({**base, "duration_seconds": 30}) + assert 5.0 < (six / one) < 7.0 + + def test_matches_measured_704x480_run(self): + """Fitted against a real run: 49 frames, 20 steps, sequential -> 248s.""" + estimate = WanVideo().estimate_runtime({ + "prompt": "x", "num_frames": 49, "width": 704, "height": 480, + "num_inference_steps": 20, "offload_mode": "sequential", + }) + assert 0.7 < estimate / 248.0 < 1.4 + + def test_sequential_costs_more_than_resident(self): + tool = WanVideo() + base = dict(prompt="x", num_frames=49, width=704, height=480, num_inference_steps=20) + assert ( + tool.estimate_runtime({**base, "offload_mode": "sequential"}) + > tool.estimate_runtime({**base, "offload_mode": "model"}) + ) + + +# ------------------------------------------------------------------ +# Single-pass operations must not be handed a long-form frame count +# ------------------------------------------------------------------ + + +class TestNonChainableOperations: + """Only text_to_video and image_to_video can continue from a last frame. + + Everything else is single-pass, so a long ``duration_seconds`` must be + clamped to the model's native clip length instead of being sent to the GPU + as one enormous denoise. + """ + + @staticmethod + def _run(monkeypatch, **inputs): + """Drive generate_wan with the GPU replaced by a recorder.""" + from PIL import Image + + from tools.video import _wan_engine + + calls = {} + + class FakePipeline: + vae_scale_factor_temporal = 4 + vae_scale_factor_spatial = 16 + transformer = None + + monkeypatch.setattr(_wan_engine, "load_wan_pipeline", lambda **kw: FakePipeline()) + monkeypatch.setattr(_wan_engine, "pipeline_scales", lambda pipeline: (4, 32)) + + def fake_single(pipeline, **kwargs): + calls["num_frames"] = kwargs["num_frames"] + return [Image.new("RGB", (16, 16))] * kwargs["num_frames"] + + def fake_chained(pipeline, **kwargs): + calls["segment_count"] = kwargs["segment_count"] + total = kwargs["total_frames"] + return [Image.new("RGB", (16, 16))] * total, [] + + monkeypatch.setattr(_wan_engine, "_generate_single", fake_single) + monkeypatch.setattr(_wan_engine, "_generate_chained", fake_chained) + monkeypatch.setattr(_wan_engine, "export_frames", lambda *a, **k: None) + monkeypatch.setattr("tools.video._shared.probe_output", lambda path: {}) + + result = _wan_engine.generate_wan( + tool_name="wan_video", + variants=WAN_VARIANTS, + default_variant=DEFAULT_VARIANT, + inputs={"prompt": "x", **inputs}, + ) + return result, calls + + NATIVE = WAN_VARIANTS[DEFAULT_VARIANT]["default_num_frames"] + + @pytest.mark.parametrize("operation", ["video_to_video", "first_last_frame"]) + def test_long_request_is_clamped_to_one_pass(self, monkeypatch, operation, tmp_path): + result, calls = self._run( + monkeypatch, + operation=operation, + duration_seconds=30, + source_video_path=str(tmp_path / "in.mp4"), + output_path=str(tmp_path / "out.mp4"), + ) + assert result.success, result.error + assert result.data["segments"] == 1 + assert calls["num_frames"] <= self.NATIVE + # and the caller is told it was clamped rather than silently short-changed + assert result.data["clamped_from_frames"] == 30 * 24 + + def test_text_to_video_still_chains(self, monkeypatch, tmp_path): + result, calls = self._run( + monkeypatch, + operation="text_to_video", + duration_seconds=30, + output_path=str(tmp_path / "out.mp4"), + ) + assert result.success, result.error + assert calls["segment_count"] == 6 + assert result.data["clamped_from_frames"] is None + + def test_short_request_is_not_marked_as_clamped(self, monkeypatch, tmp_path): + result, _ = self._run( + monkeypatch, + operation="first_last_frame", + num_frames=49, + output_path=str(tmp_path / "out.mp4"), + ) + assert result.success, result.error + assert result.data["clamped_from_frames"] is None + + def test_chainable_operations_are_exactly_the_video_continuations(self): + chainable = {"text_to_video", "image_to_video"} + assert chainable < set(_PIPELINE_FOR_OPERATION) + assert "video_to_video" not in chainable + assert "first_last_frame" not in chainable diff --git a/tools/video/_shared.py b/tools/video/_shared.py index d4d679c8..8777dca4 100644 --- a/tools/video/_shared.py +++ b/tools/video/_shared.py @@ -32,37 +32,136 @@ HEYGEN_PROVIDERS = { "ltx_distilled": {"name": "LTX Distilled", "quality": "low", "speed": "fastest"}, } +# Wan variants. ``spatial_alignment`` is the VAE spatial compression times the +# transformer patch size — the TI2V line's 16x VAE plus 2x patching means 720p +# is 1280x704, while the 8x-VAE lines take a true 1280x720. +# Wan variants. ``operations`` is the authoritative capability list — the engine +# refuses anything absent from it rather than loading a checkpoint that cannot +# do the job. ``params_b`` is the transformer size used to auto-pick precision. WAN_VARIANTS = { - "wan2.1-1.3b": { - "name": "Wan 2.1 (1.3B)", - "hf_id": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - "hf_i2v_id": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", + "wan2.2-ti2v-5b": { + "name": "Wan 2.2 TI2V (5B)", + "hf_id": "Wan-AI/Wan2.2-TI2V-5B-Diffusers", "pipeline_class": "WanPipeline", - "vram_mb": 8000, + "params_b": 5.0, + "vram_mb": 12000, "quality": "high", "speed": "medium", + "operations": [ + "text_to_video", + "image_to_video", + "video_to_video", + "first_last_frame", + "text_to_image", + ], "t2v": True, "i2v": True, "license": "Apache-2.0", + # The TI2V line's 720p geometry is 1280x704, not 1280x720: its VAE + # compresses 16x spatially and the transformer patches 2x on top, so + # both dimensions must be multiples of 32 — and 720 is not. + "default_width": 1280, + "default_height": 704, + "default_num_frames": 121, + "default_steps": 50, + "default_guidance": 5.0, + "spatial_alignment": 32, + "temporal_scale": 4, + "fps": 24, + }, + "wan2.2-t2v-a14b": { + "name": "Wan 2.2 T2V MoE (A14B)", + "hf_id": "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + "pipeline_class": "WanPipeline", + "params_b": 28.0, # two 14B experts, both resident + "vram_mb": 40000, + "quality": "highest", + "speed": "slow", + "operations": ["text_to_video", "text_to_image"], + "t2v": True, + "i2v": False, + "license": "Apache-2.0", + "default_width": 1280, + "default_height": 720, + "default_num_frames": 81, + "default_steps": 40, + "default_guidance": 4.0, + "spatial_alignment": 16, + "temporal_scale": 4, + "fps": 16, + }, + "wan2.2-i2v-a14b": { + "name": "Wan 2.2 I2V MoE (A14B)", + "hf_id": "Wan-AI/Wan2.2-I2V-A14B-Diffusers", + "pipeline_class": "WanImageToVideoPipeline", + "params_b": 28.0, + "vram_mb": 40000, + "quality": "highest", + "speed": "slow", + "operations": ["image_to_video", "first_last_frame"], + "t2v": False, + "i2v": True, + "license": "Apache-2.0", + "default_width": 1280, + "default_height": 720, + "default_num_frames": 81, + "default_steps": 40, + "default_guidance": 3.5, + "spatial_alignment": 16, + "temporal_scale": 4, + "fps": 16, + }, + "wan2.1-1.3b": { + "name": "Wan 2.1 T2V (1.3B)", + "hf_id": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + "pipeline_class": "WanPipeline", + "params_b": 1.3, + "vram_mb": 8000, + "quality": "medium", + "speed": "fast", + # Wan 2.1 never shipped a 1.3B image-to-video checkpoint; the only I2V + # weights in that generation are 14B. Claiming i2v here silently pulled + # a 14B model onto an 8GB budget. + "operations": ["text_to_video", "video_to_video", "text_to_image"], + "t2v": True, + "i2v": False, + "license": "Apache-2.0", "default_width": 832, "default_height": 480, "default_num_frames": 81, + "default_steps": 50, + "default_guidance": 5.0, + "spatial_alignment": 16, + "temporal_scale": 4, "fps": 16, }, "wan2.1-14b": { "name": "Wan 2.1 (14B)", "hf_id": "Wan-AI/Wan2.1-T2V-14B-Diffusers", - "hf_i2v_id": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", + "hf_image_to_video_id": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", + "hf_first_last_frame_id": "Wan-AI/Wan2.1-FLF2V-14B-720P-diffusers", "pipeline_class": "WanPipeline", + "params_b": 14.0, "vram_mb": 24000, "quality": "highest", "speed": "slow", + "operations": [ + "text_to_video", + "image_to_video", + "video_to_video", + "first_last_frame", + "text_to_image", + ], "t2v": True, "i2v": True, "license": "Apache-2.0", "default_width": 1280, "default_height": 720, "default_num_frames": 81, + "default_steps": 50, + "default_guidance": 5.0, + "spatial_alignment": 16, + "temporal_scale": 4, "fps": 16, }, } diff --git a/tools/video/_wan_engine.py b/tools/video/_wan_engine.py new file mode 100644 index 00000000..3bde06be --- /dev/null +++ b/tools/video/_wan_engine.py @@ -0,0 +1,947 @@ +"""Local Wan 2.x generation engine. + +Covers the full Wan operation matrix on a single consumer GPU: + + text_to_video WanPipeline + image_to_video WanImageToVideoPipeline + video_to_video WanVideoToVideoPipeline + first_last_frame WanImageToVideoPipeline (``last_image``) + text_to_image WanPipeline with a single frame + +The entry point is :func:`generate_wan`. When the requested duration exceeds +what a Wan model produces in one pass (~5s), it chains image-to-video segments +off each previous segment's last frame to reach the target length. + +Wan's VAE constrains the request geometry: frame counts must land on +``temporal * k + 1`` and both spatial dimensions must be a multiple of the +spatial compression factor. Callers pass whatever they like; this module snaps +the values and reports what it actually used. +""" + +from __future__ import annotations + +import gc +import math +from pathlib import Path +from typing import Any + +# Wan 2.1 uses a 8x8x4 VAE, Wan 2.2's TI2V line uses the 16x16x4 high-compression +# VAE. Both are read off the loaded pipeline; these are only the fallbacks used +# for planning before a pipeline exists. +_DEFAULT_TEMPORAL_SCALE = 4 +_DEFAULT_SPATIAL_SCALE = 16 + +# Wan's own reference implementation ships this Chinese negative prompt and the +# model is visibly tuned for it; it suppresses the static/oversaturated failure +# mode far better than an English equivalent. +DEFAULT_NEGATIVE_PROMPT = ( + "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量," + "低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的," + "毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" +) + +_PIPELINE_FOR_OPERATION = { + "text_to_video": "WanPipeline", + "text_to_image": "WanPipeline", + "image_to_video": "WanImageToVideoPipeline", + "first_last_frame": "WanImageToVideoPipeline", + "video_to_video": "WanVideoToVideoPipeline", +} + +# Keyed by (model_id, pipeline class, precision, offload enabled, offload mode). +# One loaded pipeline is 10-28GB of host RAM, so at most one is kept alive. +_PIPELINE_CACHE: dict[tuple[str, str, str, bool, str], Any] = {} + + +# -------------------------------------------------------------------------- +# geometry +# -------------------------------------------------------------------------- + + +def snap_frames(num_frames: int, temporal_scale: int = _DEFAULT_TEMPORAL_SCALE) -> int: + """Snap ``num_frames`` down onto the nearest valid ``temporal_scale * k + 1``. + + Wan's causal VAE encodes the first frame on its own and then groups every + following ``temporal_scale`` frames, so only these counts round-trip. We + round to nearest and clamp to a single latent group. + """ + if num_frames < 1: + raise ValueError("num_frames must be >= 1") + if num_frames == 1: + return 1 # single-image mode + k = max(1, round((num_frames - 1) / temporal_scale)) + return temporal_scale * k + 1 + + +def snap_dimension(value: int, spatial_scale: int = _DEFAULT_SPATIAL_SCALE) -> int: + """Snap a width/height onto a multiple of the VAE spatial compression.""" + if value < spatial_scale: + return spatial_scale + return int(round(value / spatial_scale)) * spatial_scale + + +def plan_segments(total_frames: int, segment_frames: int) -> int: + """Number of chained segments needed to cover ``total_frames``. + + The first segment contributes all of its frames. Every continuation is + seeded with the previous segment's last frame, so it reproduces that frame + and contributes only ``segment_frames - 1`` new ones. + """ + if total_frames <= segment_frames: + return 1 + carry = segment_frames - 1 + return 1 + math.ceil((total_frames - segment_frames) / carry) + + +# -------------------------------------------------------------------------- +# loading +# -------------------------------------------------------------------------- + + +def _vram_gb() -> float: + try: + import torch + + if not torch.cuda.is_available(): + return 0.0 + return torch.cuda.get_device_properties(0).total_memory / 1024**3 + except Exception: + return 0.0 + + +def resolve_precision(precision: str, model_params_b: float) -> str: + """Pick a weight precision that fits the visible GPU. + + ``model_params_b`` is the transformer size in billions of parameters; bf16 + needs roughly 2GB per billion resident on the device, and cpu-offload only + keeps one module at a time, so the transformer alone sets the floor. + """ + if precision != "auto": + return precision + vram = _vram_gb() + if vram <= 0: + return "bf16" # CPU/MPS: quantization backends are CUDA-only + bf16_need = model_params_b * 2.0 + 2.5 # weights + activation headroom + if vram >= bf16_need: + return "bf16" + if vram >= model_params_b * 1.0 + 2.5: + return "int8" + return "int4" + + +def resolve_offload_mode(offload_mode: str, precision: str, model_params_b: float) -> str: + """Choose how aggressively to offload weights to host RAM. + + ``model`` keeps a whole component resident and is fastest. It only works if + the largest component plus its activations fit; a bf16 5B transformer is + 10GB, so on a 12GB card that leaves too little for anything above about + 512x320. ``sequential`` streams one submodule at a time, costing PCIe + bandwidth per step but freeing almost the entire card for activations. + """ + if offload_mode != "auto": + return offload_mode + vram = _vram_gb() + if vram <= 0: + return "model" + bytes_per_param = {"int4": 0.55, "int8": 1.05}.get(precision, 2.0) + resident = model_params_b * bytes_per_param + # Need room for the weights and a working set of activations on top. + return "model" if vram >= resident + 3.0 else "sequential" + + +def _quant_config(precision: str): + from diffusers import BitsAndBytesConfig + + import torch + + if precision == "int8": + return BitsAndBytesConfig(load_in_8bit=True) + return BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + ) + + +def free_wan_pipelines() -> None: + """Drop every cached pipeline and release GPU memory.""" + _PIPELINE_CACHE.clear() + gc.collect() + try: + import torch + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + except Exception: + pass + + +def load_wan_pipeline( + *, + model_id: str, + operation: str, + precision: str = "bf16", + enable_offload: bool = True, + offload_mode: str = "model", +): + """Load (and cache) the Wan pipeline class matching ``operation``. + + Wan repos declare a single ``_class_name`` in ``model_index.json`` but the + same weights back several tasks, so the pipeline class is chosen from the + operation rather than the repo metadata. Components absent from the repo + (Wan 2.2 TI2V has no CLIP image encoder) are optional on the diffusers + classes and resolve to ``None``. + """ + import diffusers + import torch + + from tools.video._shared import get_torch_device + + pipeline_name = _PIPELINE_FOR_OPERATION.get(operation) + if pipeline_name is None: + raise ValueError( + f"Unsupported operation: {operation}. " + f"Expected one of {', '.join(sorted(_PIPELINE_FOR_OPERATION))}" + ) + + key = (model_id, pipeline_name, precision, enable_offload, offload_mode) + if key in _PIPELINE_CACHE: + return _PIPELINE_CACHE[key] + + # Only one Wan pipeline fits in host RAM at a time. + free_wan_pipelines() + + device = get_torch_device() + if device == "cuda" and torch.cuda.is_bf16_supported(): + dtype = torch.bfloat16 + elif device == "cpu": + dtype = torch.float32 + else: + dtype = torch.float16 + + pipeline_class = getattr(diffusers, pipeline_name) + kwargs: dict[str, Any] = {"torch_dtype": dtype} + + if precision in {"int4", "int8"} and device == "cuda": + # Quantize the transformer(s) only. The T5 text encoder runs once per + # prompt and the VAE once per clip, so offloading those is cheaper than + # paying the dequantization cost on every denoising step. + from diffusers import WanTransformer3DModel + + quant = _quant_config(precision) + transformer = WanTransformer3DModel.from_pretrained( + model_id, subfolder="transformer", quantization_config=quant, torch_dtype=dtype + ) + kwargs["transformer"] = transformer + if _has_subfolder(model_id, "transformer_2"): + kwargs["transformer_2"] = WanTransformer3DModel.from_pretrained( + model_id, subfolder="transformer_2", quantization_config=quant, torch_dtype=dtype + ) + + pipeline = pipeline_class.from_pretrained(model_id, **kwargs) + + # bitsandbytes modules cannot be moved after load; the accelerate hooks + # dispatch rather than relocate them, so offload still applies. + if enable_offload and device == "cuda": + if offload_mode == "sequential": + # Keeps only the executing submodule on the GPU. A 5B transformer + # in bf16 is 10GB, which leaves nothing for activations on a 12GB + # card; sequential offload trades PCIe bandwidth for that headroom + # and is what makes 480p/720p reachable here at all. + pipeline.enable_sequential_cpu_offload() + else: + pipeline.enable_model_cpu_offload() + elif precision in {"int4", "int8"}: + pass # already placed by the quantization loader + else: + pipeline = pipeline.to(device) + + vae = getattr(pipeline, "vae", None) + if vae is not None: + # Tiling keeps VAE decode of a 720p/121-frame clip off the OOM line; + # it is the single biggest peak-memory win for long clips. + if hasattr(vae, "enable_tiling"): + vae.enable_tiling() + if hasattr(vae, "enable_slicing"): + vae.enable_slicing() + + pipeline._is_sequentially_offloaded = bool( + enable_offload and device == "cuda" and offload_mode == "sequential" + ) + _PIPELINE_CACHE[key] = pipeline + return pipeline + + +def _has_subfolder(model_id: str, subfolder: str) -> bool: + local = Path(model_id) + if local.is_dir(): + return (local / subfolder).is_dir() + try: + from huggingface_hub import model_info + + return any(s.rfilename.startswith(f"{subfolder}/") for s in model_info(model_id).siblings) + except Exception: + return False + + +def pipeline_scales(pipeline) -> tuple[int, int]: + """Read (temporal, spatial) VAE compression factors off a loaded pipeline.""" + temporal = getattr(pipeline, "vae_scale_factor_temporal", None) or _DEFAULT_TEMPORAL_SCALE + spatial = getattr(pipeline, "vae_scale_factor_spatial", None) or _DEFAULT_SPATIAL_SCALE + # Transformer patch size doubles the effective spatial alignment requirement. + transformer = getattr(pipeline, "transformer", None) + patch = getattr(getattr(transformer, "config", None), "patch_size", None) + if patch: + spatial = spatial * int(patch[-1]) + return int(temporal), int(spatial) + + +# -------------------------------------------------------------------------- +# media io +# -------------------------------------------------------------------------- + + +def load_image(path: str | None, url: str | None, width: int, height: int): + """Load a reference still from disk or URL, resized to the target geometry.""" + from io import BytesIO + + from PIL import Image + + if path: + image = Image.open(path).convert("RGB") + elif url: + import requests + + response = requests.get(url, timeout=60) + response.raise_for_status() + image = Image.open(BytesIO(response.content)).convert("RGB") + else: + raise ValueError("This operation requires reference_image_path or reference_image_url") + return image.resize((width, height), Image.LANCZOS) + + +def load_video_frames(path: str, width: int, height: int, max_frames: int) -> list: + """Read up to ``max_frames`` frames from a video file as resized PIL images.""" + import imageio.v3 as iio + from PIL import Image + + frames = [] + for index, frame in enumerate(iio.imiter(path)): + if index >= max_frames: + break + frames.append(Image.fromarray(frame).convert("RGB").resize((width, height), Image.LANCZOS)) + if not frames: + raise ValueError(f"No frames could be read from {path}") + return frames + + +def export_frames(frames: list, output_path: Path, fps: int) -> None: + """Encode PIL frames to an H.264 MP4 that plays everywhere.""" + import imageio + import numpy as np + + output_path.parent.mkdir(parents=True, exist_ok=True) + writer = imageio.get_writer( + str(output_path), + fps=fps, + codec="libx264", + quality=8, + macro_block_size=None, + ffmpeg_params=["-pix_fmt", "yuv420p"], + ) + try: + for frame in frames: + writer.append_data(np.asarray(frame)) + finally: + writer.close() + + +def _to_pil_list(output) -> list: + """Normalise a diffusers video output into a flat list of PIL images.""" + from PIL import Image + import numpy as np + + frames = output.frames[0] if hasattr(output, "frames") else output.images + result = [] + for frame in frames: + if isinstance(frame, Image.Image): + result.append(frame) + else: + array = np.asarray(frame) + if array.dtype != np.uint8: + array = (np.clip(array, 0.0, 1.0) * 255).round().astype(np.uint8) + result.append(Image.fromarray(array)) + return result + + +# -------------------------------------------------------------------------- +# drift control +# -------------------------------------------------------------------------- + + +def channel_stats(frames: list) -> tuple: + """Per-channel mean and std over a list of frames, for drift correction.""" + import numpy as np + + sample = np.stack([np.asarray(f, dtype=np.float32) for f in frames[:: max(1, len(frames) // 8)]]) + return sample.mean(axis=(0, 1, 2)), sample.std(axis=(0, 1, 2)) + 1e-6 + + +def match_channel_stats(frames: list, reference: tuple) -> list: + """Re-grade ``frames`` so their colour statistics match ``reference``. + + Chained image-to-video drifts in contrast and colour temperature every hop, + because each segment is conditioned on an already-generated frame. Matching + each segment back to the opening segment's statistics keeps a 30s chain from + visibly fading or warming over its length. + """ + import numpy as np + from PIL import Image + + ref_mean, ref_std = reference + cur_mean, cur_std = channel_stats(frames) + scale = ref_std / cur_std + corrected = [] + for frame in frames: + array = np.asarray(frame, dtype=np.float32) + array = (array - cur_mean) * scale + ref_mean + corrected.append(Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))) + return corrected + + +# -------------------------------------------------------------------------- +# generation +# -------------------------------------------------------------------------- + + +def _generator(seed: int | None): + import torch + + if seed is None: + return None + # A CPU generator stays reproducible regardless of how the pipeline is + # offloaded across devices. + return torch.Generator(device="cpu").manual_seed(seed) + + +def _run_pass( + pipeline, + *, + operation: str, + prompt: str, + negative_prompt: str, + width: int, + height: int, + num_frames: int, + steps: int, + guidance_scale: float, + guidance_scale_2: float | None, + seed: int | None, + image=None, + last_image=None, + video=None, + strength: float = 0.8, +) -> list: + """Run one Wan pipeline call and return its frames as PIL images.""" + args: dict[str, Any] = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "height": height, + "width": width, + "num_inference_steps": steps, + "guidance_scale": guidance_scale, + "generator": _generator(seed), + "output_type": "pil", + } + # WanVideoToVideoPipeline derives its length from the input clip and takes + # no num_frames; every other Wan pipeline requires it. + if operation == "video_to_video": + args["video"] = video + args["strength"] = strength + else: + args["num_frames"] = num_frames + + # guidance_scale_2 only exists on the MoE (A14B) pipelines that carry a + # second expert; passing it to the 5B dense model raises. + if guidance_scale_2 is not None and getattr(pipeline, "transformer_2", None) is not None: + args["guidance_scale_2"] = guidance_scale_2 + + if image is not None: + args["image"] = image + if last_image is not None: + args["last_image"] = last_image + + return _to_pil_list(pipeline(**args)) + + +def generate_wan( + *, + tool_name: str, + variants: dict[str, dict[str, Any]], + default_variant: str, + inputs: dict[str, Any], +): + """Generate with a local Wan model, chaining segments for long durations.""" + from tools.base_tool import ToolResult + + variant_key = inputs.get("model_variant") or default_variant + if variant_key not in variants: + return ToolResult( + success=False, + error=f"Unknown model_variant: {variant_key}. Available: {', '.join(sorted(variants))}", + ) + meta = variants[variant_key] + operation = inputs.get("operation", "text_to_video") + + if operation not in _PIPELINE_FOR_OPERATION: + return ToolResult( + success=False, + error=( + f"Unknown operation: {operation}. " + f"Expected one of {', '.join(sorted(_PIPELINE_FOR_OPERATION))}" + ), + ) + supported = meta.get("operations", []) + if operation not in supported: + return ToolResult( + success=False, + error=( + f"{meta['name']} does not support {operation}. " + f"It supports: {', '.join(supported)}." + ), + ) + + prompt = inputs.get("prompt") + if not prompt: + return ToolResult(success=False, error="prompt is required") + + model_id = meta.get(f"hf_{operation}_id") or meta["hf_id"] + precision = resolve_precision(inputs.get("precision", "auto"), meta["params_b"]) + enable_offload = inputs.get("enable_model_offload", True) + offload_mode = resolve_offload_mode( + inputs.get("offload_mode", "auto"), precision, meta["params_b"] + ) + + try: + pipeline = load_wan_pipeline( + model_id=model_id, + operation=operation, + precision=precision, + enable_offload=enable_offload, + offload_mode=offload_mode, + ) + except Exception as exc: + return ToolResult(success=False, error=f"Failed to load {model_id}: {exc}") + + temporal_scale, spatial_scale = pipeline_scales(pipeline) + width = snap_dimension(int(inputs.get("width") or meta["default_width"]), spatial_scale) + height = snap_dimension(int(inputs.get("height") or meta["default_height"]), spatial_scale) + fps = int(inputs.get("fps") or meta["fps"]) + + segment_frames = snap_frames( + int(inputs.get("segment_frames") or meta["default_num_frames"]), temporal_scale + ) + + # Requested length: explicit num_frames wins, else duration_seconds, else + # the model's native single-pass clip length. + if inputs.get("num_frames"): + total_frames = int(inputs["num_frames"]) + elif inputs.get("duration_seconds"): + total_frames = max(1, round(float(inputs["duration_seconds"]) * fps)) + else: + total_frames = segment_frames + + if operation == "text_to_image": + total_frames = segment_frames = 1 + + steps = int(inputs.get("num_inference_steps") or meta["default_steps"]) + guidance_scale = float(inputs.get("guidance_scale") or meta["default_guidance"]) + guidance_scale_2 = inputs.get("guidance_scale_2") + negative_prompt = inputs.get("negative_prompt", DEFAULT_NEGATIVE_PROMPT) + seed = inputs.get("seed") + strength = float(inputs.get("strength", 0.8)) + + output_path = Path( + inputs.get("output_path") + or f"{tool_name}_{variant_key}_{operation}.{'png' if operation == 'text_to_image' else 'mp4'}" + ) + + # Only text_to_video and image_to_video can be continued from a last frame. + # The rest are single-pass, so a long duration request is clamped to what the + # model can actually denoise at once rather than being sent to certain OOM. + chainable = operation in {"text_to_video", "image_to_video"} + clamped_from = None + if chainable: + segment_count = plan_segments(total_frames, segment_frames) + else: + segment_count = 1 + if total_frames > segment_frames: + clamped_from = total_frames + total_frames = segment_frames + progress = inputs.get("progress_callback") + + try: + if segment_count > 1: + frames, segment_log = _generate_chained( + pipeline, + operation=operation, + inputs=inputs, + prompt=prompt, + negative_prompt=negative_prompt, + width=width, + height=height, + segment_frames=segment_frames, + total_frames=total_frames, + segment_count=segment_count, + steps=steps, + guidance_scale=guidance_scale, + guidance_scale_2=guidance_scale_2, + seed=seed, + correct_drift=inputs.get("correct_drift", True), + progress=progress, + ) + else: + frames = _generate_single( + pipeline, + operation=operation, + inputs=inputs, + prompt=prompt, + negative_prompt=negative_prompt, + width=width, + height=height, + num_frames=snap_frames(total_frames, temporal_scale) + if total_frames > 1 + else 1, + steps=steps, + guidance_scale=guidance_scale, + guidance_scale_2=guidance_scale_2, + seed=seed, + strength=strength, + ) + segment_log = [{"index": 0, "frames": len(frames), "operation": operation}] + except Exception as exc: + return _oom_aware_error(exc, variant_key, precision, offload_mode, width, height, total_frames) + + if operation == "text_to_image": + output_path.parent.mkdir(parents=True, exist_ok=True) + frames[0].save(output_path) + artifact_kind = "image" + else: + frames = frames[:total_frames] + export_frames(frames, output_path, fps) + artifact_kind = "video" + + from tools.video._shared import probe_output + + data = { + "provider": tool_name, + "model_variant": variant_key, + "provider_name": meta["name"], + "mode": "local", + "operation": operation, + "prompt": prompt, + "negative_prompt": negative_prompt, + "model_id": model_id, + "precision": precision, + "offload_mode": offload_mode if enable_offload else "none", + "width": width, + "height": height, + "fps": fps, + "num_frames": len(frames), + "duration_seconds": round(len(frames) / fps, 2), + "segments": segment_count, + "segment_frames": segment_frames, + "clamped_from_frames": clamped_from, + "segment_log": segment_log, + "num_inference_steps": steps, + "guidance_scale": guidance_scale, + "output": str(output_path), + "format": "png" if artifact_kind == "image" else "mp4", + "license": meta["license"], + } + if artifact_kind == "video": + data.update(probe_output(output_path)) + + return ToolResult( + success=True, + data=data, + artifacts=[str(output_path)], + seed=seed, + model=model_id, + ) + + +def derive_pipeline(base, target_operation: str): + """Build a sibling Wan pipeline that shares ``base``'s loaded weights. + + Chained generation starts in text-to-video and continues in image-to-video. + Both tasks run the same checkpoint, so re-reading tens of GB from disk just + to change pipeline class would dominate the runtime. Instead the modules + are handed to the other class directly; components that class declares but + this checkpoint lacks (Wan 2.2 TI2V has no CLIP image encoder) are optional + and resolve to ``None``. Falls back to a plain reload if anything about the + checkpoint makes the shortcut invalid. + """ + import inspect + + import diffusers + + target_name = _PIPELINE_FOR_OPERATION[target_operation] + if type(base).__name__ == target_name: + return base + + target_class = getattr(diffusers, target_name) + components = dict(base.components) + config = dict(base.config) + + kwargs: dict[str, Any] = {} + for name in list(inspect.signature(target_class.__init__).parameters)[1:]: + if name in components: + kwargs[name] = components[name] + elif name in config: + kwargs[name] = config[name] + else: + kwargs[name] = None + + # Offload hooks belong to whichever pipeline installed them; move them over + # rather than leaving two pipelines fighting over the same modules. + offloaded = getattr(base, "_all_hooks", None) + sequential = getattr(base, "_is_sequentially_offloaded", False) + if offloaded: + base.remove_all_hooks() + + derived = target_class(**kwargs) + if sequential: + derived.enable_sequential_cpu_offload() + elif offloaded: + derived.enable_model_cpu_offload() + + vae = getattr(derived, "vae", None) + if vae is not None: + if hasattr(vae, "enable_tiling"): + vae.enable_tiling() + if hasattr(vae, "enable_slicing"): + vae.enable_slicing() + return derived + + +def _generate_single( + pipeline, + *, + operation: str, + inputs: dict[str, Any], + prompt: str, + negative_prompt: str, + width: int, + height: int, + num_frames: int, + steps: int, + guidance_scale: float, + guidance_scale_2: float | None, + seed: int | None, + strength: float, +) -> list: + image = last_image = video = None + + if operation in {"image_to_video", "first_last_frame"}: + image = load_image( + inputs.get("reference_image_path"), inputs.get("reference_image_url"), width, height + ) + if operation == "first_last_frame": + last_image = load_image( + inputs.get("last_image_path"), inputs.get("last_image_url"), width, height + ) + if operation == "video_to_video": + source = inputs.get("source_video_path") + if not source: + raise ValueError("video_to_video requires source_video_path") + video = load_video_frames(source, width, height, num_frames) + # This pipeline takes its length from the clip rather than num_frames, + # so the clip itself must land on a latent boundary. + usable = snap_frames(len(video)) + if usable > len(video): + usable -= _DEFAULT_TEMPORAL_SCALE + video = video[: max(1, usable)] + + return _run_pass( + pipeline, + operation=operation, + prompt=prompt, + negative_prompt=negative_prompt, + width=width, + height=height, + num_frames=num_frames, + steps=steps, + guidance_scale=guidance_scale, + guidance_scale_2=guidance_scale_2, + seed=seed, + image=image, + last_image=last_image, + video=video, + strength=strength, + ) + + +def _generate_chained( + pipeline, + *, + operation: str, + inputs: dict[str, Any], + prompt: str, + negative_prompt: str, + width: int, + height: int, + segment_frames: int, + total_frames: int, + segment_count: int, + steps: int, + guidance_scale: float, + guidance_scale_2: float | None, + seed: int | None, + correct_drift: bool, + progress, +) -> tuple[list, list[dict[str, Any]]]: + """Produce a long clip by chaining image-to-video off each segment's last frame. + + Wan models are trained for roughly five seconds per pass; anything longer is + built by continuation. Each segment after the first is seeded with the + previous segment's final frame, which it regenerates as its own frame 0 — + that duplicate is dropped when the segments are joined so the seam lands on + a single shared frame. + + ``segment_prompts`` lets a caller drive the clip through a shot list; with a + single prompt every segment is asked for the same thing, which reads as one + continuous take. + """ + segment_prompts = inputs.get("segment_prompts") or [] + frames: list = [] + log: list[dict[str, Any]] = [] + reference_stats = None + anchor = None + + for index in range(segment_count): + segment_prompt = ( + segment_prompts[index] if index < len(segment_prompts) else prompt + ) + segment_seed = None if seed is None else seed + index + + if index == 0: + segment_operation = operation + if operation == "image_to_video": + anchor = load_image( + inputs.get("reference_image_path"), + inputs.get("reference_image_url"), + width, + height, + ) + else: + segment_operation = "image_to_video" + + if index > 0: + # No-op once the chain is already running on the I2V pipeline. + pipeline = derive_pipeline(pipeline, "image_to_video") + + if progress: + progress(index + 1, segment_count, segment_prompt) + + produced = _run_pass( + pipeline, + operation=segment_operation, + prompt=segment_prompt, + negative_prompt=negative_prompt, + width=width, + height=height, + num_frames=segment_frames, + steps=steps, + guidance_scale=guidance_scale, + guidance_scale_2=guidance_scale_2, + seed=segment_seed, + image=anchor if segment_operation == "image_to_video" else None, + ) + + if index == 0: + reference_stats = channel_stats(produced) + kept = produced + else: + if correct_drift and reference_stats is not None: + produced = match_channel_stats(produced, reference_stats) + # Frame 0 is the anchor the model was handed; keep the original. + kept = produced[1:] + + frames.extend(kept) + anchor = frames[-1] + log.append( + { + "index": index, + "operation": segment_operation, + "prompt": segment_prompt, + "seed": segment_seed, + "frames_generated": len(produced), + "frames_kept": len(kept), + "cumulative_frames": len(frames), + } + ) + + if len(frames) >= total_frames: + break + + return frames, log + + +# When the installed NVIDIA userspace libraries do not match the loaded kernel +# module, nvmlInit fails. PyTorch calls it while composing an out-of-memory +# report, so on such a machine every OOM — and several unrelated allocator +# paths — surface as this assert instead of the real error. +_NVML_ASSERT_MARKER = "nvmlInit" + + +def _oom_aware_error( + exc: Exception, + variant: str, + precision: str, + offload_mode: str, + width: int, + height: int, + frames: int, +): + """Turn an allocator failure into advice the caller can act on.""" + from tools.base_tool import ToolResult + + free_wan_pipelines() + message = str(exc) + + if _NVML_ASSERT_MARKER in message: + return ToolResult( + success=False, + error=( + "The NVIDIA driver on this machine is in a mismatched state: the loaded " + "kernel module and the installed userspace libraries are different versions, " + "so nvmlInit fails. CUDA compute still runs, but PyTorch calls NVML while " + "reporting an out-of-memory condition, so real OOMs surface as an internal " + "assert and bitsandbytes quantization cannot load at all. Reboot to load the " + "matching kernel module (check with 'cat /proc/driver/nvidia/version' against " + "'nvidia-smi'). Until then, stay on precision=\"bf16\" and reduce resolution " + "or segment_frames so generation never reaches the OOM path." + ), + ) + + if "out of memory" not in message.lower(): + return ToolResult(success=False, error=f"Wan generation failed: {exc}") + + smaller = f"{max(320, (width * 2) // 3)}x{max(320, (height * 2) // 3)}" + hints = [] + if offload_mode != "sequential": + hints.append('set offload_mode="sequential" to free the card for activations') + next_precision = {"bf16": "int8", "int8": "int4"}.get(precision) + if next_precision: + hints.append(f'set precision="{next_precision}"') + hints.append(f"lower the resolution (try around {smaller})") + hints.append("reduce segment_frames") + return ToolResult( + success=False, + error=( + f"Ran out of GPU memory generating {width}x{height} x{frames} frames with " + f"{variant} at {precision} (offload_mode={offload_mode}). " + "Try one of: " + "; ".join(hints) + "." + ), + ) diff --git a/tools/video/video_selector.py b/tools/video/video_selector.py index f2f71cda..0235d09e 100644 --- a/tools/video/video_selector.py +++ b/tools/video/video_selector.py @@ -386,6 +386,9 @@ class VideoSelector(BaseTool): env_hint = os.environ.get("VIDEO_GEN_LOCAL_MODEL", "").lower() env_map = { + "wan2.2-ti2v-5b": "wan", + "wan2.2-t2v-a14b": "wan", + "wan2.2-i2v-a14b": "wan", "wan2.1-1.3b": "wan", "wan2.1-14b": "wan", "hunyuan-1.5": "hunyuan", diff --git a/tools/video/wan_video.py b/tools/video/wan_video.py index 74b0398b..b00815a8 100644 --- a/tools/video/wan_video.py +++ b/tools/video/wan_video.py @@ -1,4 +1,10 @@ -"""Wan local video generation.""" +"""Wan local video generation. + +Wraps the Wan 2.x family running on the machine's own GPU: text-to-video, +image-to-video, video-to-video, first/last-frame interpolation, and single-frame +text-to-image. Clips longer than a model's native ~5s pass are produced by +chaining image-to-video segments — see ``tools/video/_wan_engine.py``. +""" from __future__ import annotations @@ -17,12 +23,32 @@ from tools.base_tool import ( ToolStatus, ToolTier, ) -from tools.video._shared import WAN_VARIANTS, estimate_local_runtime, generate_local_video, local_generation_status, local_install_instructions +from tools.video._shared import WAN_VARIANTS, local_generation_status, local_install_instructions +from tools.video._wan_engine import ( + DEFAULT_NEGATIVE_PROMPT, + generate_wan, + plan_segments, + snap_frames, +) + +DEFAULT_VARIANT = "wan2.2-ti2v-5b" + +# Runtime model fitted to measurements on an Ada-class laptop GPU (11.6GB) with +# the 5B TI2V checkpoint in bf16 at 704x480: +# +# 49 frames, 20 steps -> 11.9 s/step 121 frames, 20 steps -> 23.3 s/step +# +# which separates into a per-step cost that does not depend on the clip (under +# sequential offload every step streams the whole transformer across PCIe) and a +# per-step cost proportional to the number of pixels being denoised. +_SEQUENTIAL_STREAM_SECONDS_PER_STEP = 4.1 +_SECONDS_PER_STEP_PER_MPX_FRAME = 0.47 +_MODEL_LOAD_SECONDS = 15.0 class WanVideo(BaseTool): name = "wan_video" - version = "0.1.0" + version = "0.2.0" tier = ToolTier.GENERATE capability = "video_generation" provider = "wan" @@ -36,18 +62,32 @@ class WanVideo(BaseTool): fallback_tools = ["hunyuan_video", "ltx_video_local", "cogvideo_video", "image_selector"] agent_skills = ["ltx2"] - capabilities = ["text_to_video", "image_to_video", "model_selection"] + capabilities = [ + "text_to_video", + "image_to_video", + "video_to_video", + "first_last_frame", + "text_to_image", + "long_form_chaining", + "model_selection", + ] supports = { "reference_image": True, "offline": True, "native_audio": False, "local_gpu": True, + "long_form": True, } best_for = [ "best quality-to-VRAM ratio for local generation", - "local pipelines that still want image-to-video support", + "long local clips (20s+) assembled by image-to-video chaining", + "local pipelines that need image-to-video and video-to-video", + ] + not_good_for = [ + "CPU-only machines", + "instant iteration on low-end hardware", + "single-pass shots longer than ~5s (Wan's trained clip length)", ] - not_good_for = ["CPU-only machines", "instant iteration on low-end hardware"] provider_matrix = {key: {"tool": "wan_video", **value, "mode": "local_gpu"} for key, value in WAN_VARIANTS.items()} input_schema = { @@ -55,25 +95,80 @@ class WanVideo(BaseTool): "required": ["prompt"], "properties": { "prompt": {"type": "string"}, - "operation": {"type": "string", "enum": ["text_to_video", "image_to_video"], "default": "text_to_video"}, - "model_variant": {"type": "string", "enum": sorted(WAN_VARIANTS), "default": "wan2.1-1.3b"}, + "negative_prompt": {"type": "string", "default": DEFAULT_NEGATIVE_PROMPT}, + "operation": { + "type": "string", + "enum": [ + "text_to_video", + "image_to_video", + "video_to_video", + "first_last_frame", + "text_to_image", + ], + "default": "text_to_video", + }, + "model_variant": {"type": "string", "enum": sorted(WAN_VARIANTS), "default": DEFAULT_VARIANT}, "reference_image_url": {"type": "string"}, "reference_image_path": {"type": "string"}, + "last_image_url": {"type": "string", "description": "Target final frame for first_last_frame"}, + "last_image_path": {"type": "string", "description": "Target final frame for first_last_frame"}, + "source_video_path": {"type": "string", "description": "Input clip for video_to_video"}, + "strength": {"type": "number", "default": 0.8, "description": "video_to_video denoise strength"}, "width": {"type": "integer"}, "height": {"type": "integer"}, - "num_frames": {"type": "integer"}, + "fps": {"type": "integer"}, + "duration_seconds": { + "type": "number", + "description": "Target length. Beyond one segment the clip is chained image-to-video.", + }, + "num_frames": {"type": "integer", "description": "Exact frame count; overrides duration_seconds"}, + "segment_frames": { + "type": "integer", + "description": "Frames per chained segment. Defaults to the model's native clip length.", + }, + "segment_prompts": { + "type": "array", + "items": {"type": "string"}, + "description": "Per-segment prompts for long clips — a shot list. Falls back to prompt.", + }, + "correct_drift": { + "type": "boolean", + "default": True, + "description": "Re-grade chained segments to the opening segment's colour statistics.", + }, "num_inference_steps": {"type": "integer"}, + "guidance_scale": {"type": "number"}, + "guidance_scale_2": {"type": "number", "description": "Second-expert guidance on MoE (A14B) variants"}, + "precision": { + "type": "string", + "enum": ["auto", "bf16", "int8", "int4"], + "default": "auto", + "description": "Weight precision. auto picks the best fit for the visible VRAM.", + }, "enable_model_offload": {"type": "boolean", "default": True}, + "offload_mode": { + "type": "string", + "enum": ["auto", "model", "sequential"], + "default": "auto", + "description": ( + "How weights are staged onto the GPU. 'model' keeps a whole component " + "resident (fastest); 'sequential' streams one submodule at a time, which " + "is what lets a 12GB card render above 512x320." + ), + }, "seed": {"type": "integer"}, "output_path": {"type": "string"}, }, } - resource_profile = ResourceProfile(cpu_cores=2, ram_mb=16000, vram_mb=8000, disk_mb=4000, network_required=False) + resource_profile = ResourceProfile(cpu_cores=2, ram_mb=16000, vram_mb=8000, disk_mb=40000, network_required=False) retry_policy = RetryPolicy(max_retries=1) - idempotency_key_fields = ["prompt", "model_variant", "operation", "seed"] + idempotency_key_fields = ["prompt", "model_variant", "operation", "seed", "duration_seconds"] side_effects = ["writes video file to output_path", "may download model weights"] - user_visible_verification = ["Watch generated clip for motion coherence and artifacts"] + user_visible_verification = [ + "Watch generated clip for motion coherence and artifacts", + "On chained clips, check the segment seams for a jump in colour or subject pose", + ] def get_status(self) -> ToolStatus: return local_generation_status() @@ -81,18 +176,69 @@ class WanVideo(BaseTool): def estimate_cost(self, inputs: dict[str, Any]) -> float: return 0.0 + def _plan(self, inputs: dict[str, Any]) -> dict[str, Any]: + """Resolve the frame/segment plan without loading any weights.""" + meta = WAN_VARIANTS.get(inputs.get("model_variant") or DEFAULT_VARIANT, WAN_VARIANTS[DEFAULT_VARIANT]) + fps = int(inputs.get("fps") or meta["fps"]) + segment_frames = snap_frames(int(inputs.get("segment_frames") or meta["default_num_frames"])) + if inputs.get("num_frames"): + total_frames = int(inputs["num_frames"]) + elif inputs.get("duration_seconds"): + total_frames = max(1, round(float(inputs["duration_seconds"]) * fps)) + else: + total_frames = segment_frames + if inputs.get("operation") == "text_to_image": + total_frames = segment_frames = 1 + # A clip shorter than one native segment is generated at its own length, + # not padded out to the model's default. + segment_frames = min(segment_frames, max(1, total_frames)) + + from tools.video._wan_engine import resolve_offload_mode, resolve_precision + + precision = resolve_precision(inputs.get("precision", "auto"), meta["params_b"]) + return { + "meta": meta, + "precision": precision, + "offload_mode": resolve_offload_mode( + inputs.get("offload_mode", "auto"), precision, meta["params_b"] + ), + "fps": fps, + "segment_frames": segment_frames, + "total_frames": total_frames, + "segments": plan_segments(total_frames, segment_frames), + "steps": int(inputs.get("num_inference_steps") or meta["default_steps"]), + "width": int(inputs.get("width") or meta["default_width"]), + "height": int(inputs.get("height") or meta["default_height"]), + } + def estimate_runtime(self, inputs: dict[str, Any]) -> float: - variant = WAN_VARIANTS.get(inputs.get("model_variant", "wan2.1-1.3b"), WAN_VARIANTS["wan2.1-1.3b"]) - return estimate_local_runtime(variant["speed"]) + """Estimated wall-clock seconds, including weight loading. + + Every chained segment pays the full per-step cost again, so a 30s clip + costs roughly six times a 5s one rather than a little more. + """ + plan = self._plan(inputs) + megapixels = (plan["width"] * plan["height"]) / 1_000_000 + stream = ( + _SEQUENTIAL_STREAM_SECONDS_PER_STEP + if plan["offload_mode"] == "sequential" + else 0.0 + ) + per_step = stream + plan["segment_frames"] * megapixels * _SECONDS_PER_STEP_PER_MPX_FRAME + return round(plan["segments"] * plan["steps"] * per_step + _MODEL_LOAD_SECONDS, 1) def execute(self, inputs: dict[str, Any]) -> ToolResult: if self.get_status() != ToolStatus.AVAILABLE: return ToolResult(success=False, error="Wan local video generation is unavailable. " + self.install_instructions) start = time.time() try: - result = generate_local_video(tool_name=self.name, variants=WAN_VARIANTS, default_variant="wan2.1-1.3b", inputs=inputs) + result = generate_wan( + tool_name=self.name, + variants=WAN_VARIANTS, + default_variant=DEFAULT_VARIANT, + inputs=inputs, + ) except Exception as exc: return ToolResult(success=False, error=f"Wan video generation failed: {exc}") result.duration_seconds = round(time.time() - start, 2) return result -