[wip] add lingbot world 2 impl + test.

2026-07-13 17:09:46 -07:00
parent b3f5f411e8
commit 6f089ba83a
20 changed files with 3187 additions and 1 deletions
@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Create a FastVideo-native LingBot World 2 causal-fast bundle."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+
+LATENTS_MEAN = [
+    -0.7571,
+    -0.7089,
+    -0.9113,
+    0.1075,
+    -0.1745,
+    0.9653,
+    -0.1517,
+    1.5508,
+    0.4134,
+    -0.0715,
+    0.5517,
+    -0.3632,
+    -0.1922,
+    -0.9497,
+    0.2503,
+    -0.2921,
+]
+LATENTS_STD = [
+    2.8184,
+    1.4541,
+    2.3275,
+    2.6558,
+    1.2196,
+    1.7708,
+    2.6052,
+    2.0743,
+    3.2687,
+    2.1526,
+    2.8652,
+    1.5579,
+    1.6382,
+    1.1253,
+    2.8251,
+    1.9160,
+]
+
+
+def _write_json(path: Path, payload: dict) -> None:
+    """Write deterministic JSON config files for a local bundle."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
+
+
+def _symlink_or_replace(source: Path, target: Path) -> None:
+    """Create a relative symlink, replacing stale symlinks only."""
+    target.parent.mkdir(parents=True, exist_ok=True)
+    if target.is_symlink() or target.exists():
+        if target.is_symlink() and target.resolve() == source.resolve():
+            return
+        raise FileExistsError(f"Refusing to replace existing path: {target}")
+    target.symlink_to(os.path.relpath(source, target.parent))
+
+
+def convert_vae(source_dir: Path, output_dir: Path) -> None:
+    """Link the exact LingBot World 2 Wan VAE checkpoint consumed by LingBotWorld2WanVAE."""
+    _symlink_or_replace(source_dir / "Wan2.1_VAE.pth", output_dir / "vae" / "Wan2.1_VAE.pth")
+
+
+def build_bundle(source_dir: Path, output_dir: Path) -> None:
+    """Create configs and symlinks for FastVideo-native LingBot World 2 loading."""
+    output_dir.mkdir(parents=True, exist_ok=True)
+    _write_json(
+        output_dir / "model_index.json",
+        {
+            "_class_name": "LingBotWorld2CausalFastPipeline",
+            "_diffusers_version": "0.30.0",
+            "scheduler": ["diffusers", "FlowUniPCMultistepScheduler"],
+            "text_encoder": ["transformers", "LingBotWorld2T5EncoderModel"],
+            "tokenizer": ["transformers", "AutoTokenizer"],
+            "transformer": ["diffusers", "LingBotWorld2CausalFastTransformer3DModel"],
+            "vae": ["diffusers", "LingBotWorld2WanVAE"],
+        },
+    )
+    _write_json(
+        output_dir / "transformer" / "config.json",
+        {
+            "_class_name": "LingBotWorld2CausalFastTransformer3DModel",
+            "_diffusers_version": "0.30.0",
+            "model_type": "i2v",
+            "patch_size": [1, 2, 2],
+            "text_len": 512,
+            "in_dim": 36,
+            "dim": 5120,
+            "ffn_dim": 13824,
+            "freq_dim": 256,
+            "text_dim": 4096,
+            "out_dim": 16,
+            "num_heads": 40,
+            "num_layers": 40,
+            "local_attn_size": 18,
+            "sink_size": 6,
+            "qk_norm": True,
+            "cross_attn_norm": True,
+            "eps": 1e-6,
+        },
+    )
+    _write_json(
+        output_dir / "text_encoder" / "config.json",
+        {
+            "architectures": ["LingBotWorld2T5EncoderModel"],
+            "vocab_size": 256384,
+            "dim": 4096,
+            "dim_attn": 4096,
+            "dim_ffn": 10240,
+            "num_heads": 64,
+            "num_layers": 24,
+            "num_buckets": 32,
+            "text_len": 512,
+            "hidden_size": 4096,
+            "dropout": 0.1,
+        },
+    )
+    _write_json(
+        output_dir / "vae" / "config.json",
+        {
+            "_class_name": "LingBotWorld2WanVAE",
+            "base_dim": 96,
+            "z_dim": 16,
+            "dim_mult": [1, 2, 4, 4],
+            "num_res_blocks": 2,
+            "attn_scales": [],
+            "temperal_downsample": [False, True, True],
+            "dropout": 0.0,
+            "latents_mean": LATENTS_MEAN,
+            "latents_std": LATENTS_STD,
+        },
+    )
+    _write_json(
+        output_dir / "scheduler" / "scheduler_config.json",
+        {
+            "_class_name": "FlowUniPCMultistepScheduler",
+            "num_train_timesteps": 1000,
+            "shift": 1.0,
+            "use_dynamic_shifting": False,
+            "prediction_type": "flow_prediction",
+            "solver_order": 2,
+            "solver_type": "bh2",
+            "final_sigmas_type": "zero",
+        },
+    )
+    _write_json(output_dir / "scheduler" / "config.json", json.loads((output_dir / "scheduler" / "scheduler_config.json").read_text()))
+
+    _symlink_or_replace(source_dir / "models_t5_umt5-xxl-enc-bf16.pth", output_dir / "text_encoder" / "pytorch_model.pt")
+    _symlink_or_replace(source_dir / "google" / "umt5-xxl", output_dir / "tokenizer")
+    for shard in sorted((source_dir / "transformers").glob("*.safetensors")):
+        _symlink_or_replace(shard, output_dir / "transformer" / shard.name)
+    index_path = source_dir / "transformers" / "diffusion_pytorch_model.safetensors.index.json"
+    if index_path.exists():
+        _symlink_or_replace(index_path, output_dir / "transformer" / "diffusion_pytorch_model.safetensors.index.json")
+    convert_vae(source_dir, output_dir)
+
+
+def main() -> None:
+    """Parse CLI arguments and build the local converted bundle."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--source-dir", required=True, type=Path)
+    parser.add_argument("--output-dir", required=True, type=Path)
+    args = parser.parse_args()
+    build_bundle(args.source_dir, args.output_dir)
+
+
+if __name__ == "__main__":
+    main()