125 lines
6.1 KiB
Python
125 lines
6.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Persistent native Wan VAE weights with independently restartable IPC consumers.
|
|
|
|
Run donor in one process; run consumer repeatedly in fresh processes on the same
|
|
GPU. The donor must stay alive. The example uses a local, immutable Wan 2.1 1.3B
|
|
Diffusers VAE component and FP32. It shares weights, not mutable VAE caches.
|
|
"""
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
import signal
|
|
import threading
|
|
import time
|
|
|
|
import torch
|
|
|
|
from fastvideo.configs.pipelines import PipelineConfig
|
|
from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
|
|
import fastvideo.envs as envs
|
|
from fastvideo.fastvideo_args import FastVideoArgs
|
|
from fastvideo.models.loader.component_loader import VAELoader
|
|
from fastvideo.models.loader.utils import set_default_torch_dtype
|
|
from fastvideo.models.loader.weight_cache import WeightCacheClient, WeightCacheKey, serve_weight_cache
|
|
from fastvideo.models.wan.vae import AutoencoderKLWan
|
|
from fastvideo.models.wan.vae_config import WanVAEConfig
|
|
from fastvideo.utils import get_open_port
|
|
|
|
|
|
def fingerprint(model):
|
|
digest = hashlib.sha256()
|
|
for name, tensor in [*model.named_parameters(remove_duplicate=False), *model.named_buffers(remove_duplicate=False)]:
|
|
digest.update(name.encode())
|
|
digest.update(str((tensor.dtype, tuple(tensor.shape), tuple(tensor.stride()), tensor.storage_offset())).encode())
|
|
digest.update(tensor.detach().contiguous().reshape(-1).view(torch.uint8).cpu().numpy().tobytes())
|
|
return digest.hexdigest()
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("role", choices=["donor", "consumer", "disk"])
|
|
parser.add_argument("--model-path", type=Path, required=True)
|
|
parser.add_argument("--checkpoint-id", required=True, help="Immutable checkpoint revision or content digest")
|
|
parser.add_argument("--socket", type=Path)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--tensor-output", type=Path)
|
|
args = parser.parse_args()
|
|
if args.role != "disk" and args.socket is None:
|
|
parser.error("--socket is required for donor/consumer")
|
|
config_bytes = (args.model_path / "config.json").read_bytes()
|
|
architecture = json.loads(config_bytes)
|
|
if architecture.get("_class_name") != "AutoencoderKLWan":
|
|
raise ValueError("This adapter supports native Wan VAE only")
|
|
key = WeightCacheKey(args.checkpoint_id, "vae", hashlib.sha256(config_bytes).hexdigest(),
|
|
"torch.float32", "native-wan-fp32-v1")
|
|
envs.setdefault_external("MASTER_ADDR", "127.0.0.1")
|
|
envs.setdefault_external("MASTER_PORT", str(get_open_port()))
|
|
maybe_init_distributed_environment_and_model_parallel(1, 1)
|
|
device = torch.device("cuda", torch.cuda.current_device())
|
|
report = {"role": args.role, "checkpoint": args.checkpoint_id, "torch": torch.__version__}
|
|
def load_disk():
|
|
runtime = FastVideoArgs(model_path=str(args.model_path), vae_cpu_offload=False,
|
|
pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision="fp32"))
|
|
return VAELoader().load(str(args.model_path), runtime).eval()
|
|
native_config = WanVAEConfig()
|
|
native_config.update_model_arch(architecture)
|
|
def construct():
|
|
return VAELoader().build_model(AutoencoderKLWan, native_config, torch.device("meta"))
|
|
def verify(model):
|
|
report["weight_sha256_before"] = fingerprint(model)
|
|
x = torch.randn((1, 16, 2, 12, 12), generator=torch.Generator().manual_seed(1024)).to(device)
|
|
with torch.no_grad():
|
|
first = model.decode(x)
|
|
second = model.decode(x)
|
|
torch.cuda.synchronize(device)
|
|
assert first.shape == (1, 3, 5, 96, 96) and torch.isfinite(first).all()
|
|
assert torch.equal(first, second), "Request-local VAE caches were not reset"
|
|
report["weight_sha256_after"] = fingerprint(model)
|
|
assert report["weight_sha256_before"] == report["weight_sha256_after"]
|
|
if args.tensor_output:
|
|
args.tensor_output.parent.mkdir(parents=True, exist_ok=True)
|
|
torch.save(first.cpu(), args.tensor_output)
|
|
try:
|
|
if args.role == "donor":
|
|
model = load_disk()
|
|
report["weight_sha256_before"] = fingerprint(model)
|
|
stop = threading.Event()
|
|
for sig in (signal.SIGINT, signal.SIGTERM):
|
|
signal.signal(sig, lambda *_: stop.set())
|
|
print("Serving donor weights", flush=True)
|
|
serve_weight_cache(model, key, args.socket, stop=stop)
|
|
report["weight_sha256_after"] = fingerprint(model)
|
|
assert report["weight_sha256_before"] == report["weight_sha256_after"]
|
|
else:
|
|
start = time.perf_counter()
|
|
allocated = torch.cuda.memory_allocated(device)
|
|
if args.role == "consumer":
|
|
with set_default_torch_dtype(torch.float32):
|
|
with WeightCacheClient(args.socket, key, device).load_model(construct) as attachment:
|
|
report["load_seconds"] = time.perf_counter() - start
|
|
report["allocated_weight_delta_bytes"] = torch.cuda.memory_allocated(device) - allocated
|
|
report["donor_pid"] = attachment.donor_pid
|
|
report["ipc_handles"] = attachment.handles
|
|
report["weight_source"] = "cuda_ipc"
|
|
verify(attachment.model)
|
|
else:
|
|
model = load_disk()
|
|
torch.cuda.synchronize(device)
|
|
report["load_seconds"] = time.perf_counter() - start
|
|
report["allocated_weight_delta_bytes"] = torch.cuda.memory_allocated(device) - allocated
|
|
report["weight_source"] = "disk"
|
|
verify(model)
|
|
report["status"] = "ok"
|
|
except Exception as error:
|
|
report.update(status="error", error=repr(error))
|
|
raise
|
|
finally:
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(report, indent=2) + "\n")
|
|
cleanup_dist_env_and_memory()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|