File size: 5,296 Bytes
f2c0505 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | from __future__ import annotations
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import site
ROOT = Path(__file__).resolve().parents[1]
PINS = json.loads((ROOT / "source_pins.json").read_text())
CONTENT = Path("/content")
def run(cmd, *, cwd=None):
print("+", " ".join(map(str, cmd)), flush=True)
subprocess.run(list(map(str, cmd)), cwd=cwd, check=True)
def clone_at(name: str, destination: Path):
spec = PINS[name]
sha = spec["commit"]
if destination.exists():
shutil.rmtree(destination)
run(["git", "init", "-q", destination])
run(["git", "-C", destination, "remote", "add", "origin", spec["url"]])
run(["git", "-C", destination, "fetch", "-q", "--depth=1", "origin", sha])
run(["git", "-C", destination, "checkout", "-q", "FETCH_HEAD"])
got = subprocess.check_output(["git", "-C", destination, "rev-parse", "HEAD"], text=True).strip()
if got != sha:
raise RuntimeError(f"{name} pin mismatch: expected {sha}, got {got}")
print(f"✓ {name} {got}")
def version_tuple(text: str):
base = text.split("+")[0].split("rc")[0]
vals = []
for x in base.split(".")[:3]:
try:
vals.append(int(x))
except Exception:
vals.append(0)
return tuple(vals + [0] * (3 - len(vals)))
def main():
import torch
print("Torch:", torch.__version__)
print("CUDA :", torch.version.cuda)
if version_tuple(torch.__version__) < (2, 10, 0):
raise RuntimeError(
"Sol-Attn's pinned sol-engine backend requires PyTorch >= 2.10. "
"Use a Colab runtime with Torch 2.10+ rather than silently replacing the CUDA stack here."
)
if not torch.cuda.is_available():
raise RuntimeError("CUDA GPU is required")
if torch.version.cuda is None or version_tuple(torch.version.cuda) < (12, 8, 0):
raise RuntimeError(
f"Pinned Sol-Attn requires CUDA >= 12.8; current Torch CUDA runtime is {torch.version.cuda!r}. "
"Use a newer Colab GPU runtime instead of silently replacing PyTorch/CUDA."
)
cc = torch.cuda.get_device_capability()
print("GPU :", torch.cuda.get_device_name())
print("CC :", cc)
if cc[0] < 8:
raise RuntimeError("Sol-Attn supports NVIDIA SM80+ in this package")
# Install runtime Python dependencies without replacing PyTorch.
deps = [
"triton>=3.6", "safetensors==0.7.0", "numpy>=1.26", "PyYAML", "easydict",
"huggingface_hub==0.36.0", "transformers==4.57.6", "accelerate==1.13.0",
"diffusers==0.36.0", "imageio", "imageio-ffmpeg", "decord", "opencv-python-headless",
"moviepy", "loguru", "einops", "tqdm", "psutil", "omegaconf", "addict", "ftfy",
"scipy", "sentencepiece", "protobuf", "packaging", "pillow"
]
run([sys.executable, "-m", "pip", "install", "-q", "-U", *deps])
wan = CONTENT / "Wan-Animate-2"
sana = CONTENT / "Sana-sol-engine"
para = CONTENT / "ParaAttention"
clone_at("Wan-Video/Wan-Animate-2", wan)
clone_at("NVlabs/Sana", sana)
clone_at("chengzeyi/ParaAttention", para)
# Verify the exact upstream Animate-2 distilled contract we patch into.
import yaml
source_cfg = yaml.safe_load((wan / "infer" / "wan_animate_2_distillation.yaml").read_text())
tr = source_cfg["model"]["transformer"]
expected = {
"in_dim": 36,
"dim": 5120,
"ffn_dim": 13824,
"num_heads": 40,
"num_layers": 40,
"log_scale": -1.3,
}
for key, value in expected.items():
if tr.get(key) != value:
raise RuntimeError(f"official Animate-2 source contract changed: {key}={tr.get(key)!r}, expected {value!r}")
if source_cfg["model"].get("flow_solver") != "euler" or float(source_cfg["test_cfg"].get("sample_shift")) != 5.0:
raise RuntimeError("official distilled Euler/shift-5 contract changed")
print("✓ official Animate-2 architecture + distilled Euler contract")
# These integrations are pure Python/Triton at runtime. Put the exact
# checked-out source directories on sys.path via one .pth file instead of
# invoking setuptools-scm/editable builds that can drift or pull deps.
paths = [
str(ROOT),
str(wan),
str(sana / "techniques" / "sparse_backends"),
str(para / "src"),
]
site_dir = Path(site.getsitepackages()[0])
pth = site_dir / "orbitquant_wan_a2_third_party.pth"
pth.write_text("\n".join(paths) + "\n")
for path in reversed(paths):
if path not in sys.path:
sys.path.insert(0, path)
# Import gates catch namespace/path mistakes immediately.
import orbitquant_wan_a2 # noqa: F401
from sol_attn import sol_attn # noqa: F401
import para_attn.primitives # noqa: F401
from wanxiang.models.wan_animate_2_model import WanAnimate2Transformer # noqa: F401
print("\nBOOTSTRAP PASS")
print(" official source :", wan)
print(" Sol-Attn :", sana / "techniques" / "sparse_backends")
print(" ParaAttention :", para)
print(" runtime package :", ROOT)
print(" path file :", pth)
print("\nNext: python", ROOT / "scripts" / "kernel_selftest.py")
if __name__ == "__main__":
main()
|