File size: 3,708 Bytes
2e605e2 | 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 | from __future__ import annotations
import argparse, json, os, sys, time, gc
from pathlib import Path
# Export with PyTorch SDPA instead of flash-attn custom calls.
os.environ.setdefault("ATTN_BACKEND", "sdpa")
import torch
MODEL_ID = "VAST-AI/AniGen"
OUT_DEFAULT = "/tmp/cf-ss-flow"
MODEL_ROOT_DEFAULT = "/tmp/anigen-model"
APP_ROOT_DEFAULT = "/home/user/app"
def ensure_model(root: Path):
from huggingface_hub import snapshot_download
snapshot_download(
MODEL_ID,
token=os.environ.get("HF_TOKEN"),
local_dir=root,
allow_patterns=[
"ckpts/anigen/ss_flow_solo/config.json",
"ckpts/anigen/ss_flow_solo/ckpts/**",
],
)
class SSFlowExport(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, x, x_skl, timestep, cond):
y, y_skl = self.model(x, x_skl, timestep, cond)
return y, y_skl
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--out", default=os.environ.get("CF_ONNX_OUT", OUT_DEFAULT))
ap.add_argument("--model-root", default=os.environ.get("ANIGEN_MODEL_ROOT", MODEL_ROOT_DEFAULT))
ap.add_argument("--app-root", default=os.environ.get("ANIGEN_APP_ROOT", APP_ROOT_DEFAULT))
args = ap.parse_args()
out = Path(args.out)
model_root = Path(args.model_root)
app_root = Path(args.app_root)
out_dir = out / "onnx/anigen/ss-flow-solo"
out_dir.mkdir(parents=True, exist_ok=True)
sys.path.insert(0, str(app_root))
ensure_model(model_root)
os.chdir(model_root)
from anigen.utils.model_utils import load_model_from_path
model, cfg = load_model_from_path(
str(model_root / "ckpts/anigen/ss_flow_solo"),
model_name_in_config="denoiser",
device="cuda",
use_ema=False,
)
model.eval()
wrapper = SSFlowExport(model).eval()
# Production profile is batch=1. SS latent is fixed 16^3; DINO ViT-L/14-reg
# on 518x518 yields 1369 patch + cls + 4 register = 1374 conditioning tokens.
x = torch.zeros((1, 8, 16, 16, 16), device="cuda", dtype=torch.float32)
x_skl = torch.zeros((1, 4, 16, 16, 16), device="cuda", dtype=torch.float32)
timestep = torch.tensor([500.0], device="cuda", dtype=torch.float32)
cond = torch.zeros((1, 1374, 1024), device="cuda", dtype=torch.float32)
path = out_dir / "model.onnx"
started = time.time()
with torch.inference_mode():
torch.onnx.export(
wrapper,
(x, x_skl, timestep, cond),
str(path),
input_names=["x", "x_skl", "timestep", "cond"],
output_names=["velocity", "velocity_skl"],
opset_version=23,
dynamo=True,
external_data=True,
)
export_s = time.time() - started
import onnx
onnx.checker.check_model(str(path))
meta = {
"component": "anigen-ss-flow-solo",
"source": MODEL_ID,
"checkpoint": "ckpts/anigen/ss_flow_solo",
"opset": 23,
"precision_source": "fp16 torso / fp32 io",
"static_profile": {
"x": [1, 8, 16, 16, 16],
"x_skl": [1, 4, 16, 16, 16],
"timestep": [1],
"cond": [1, 1374, 1024],
},
"export_seconds": round(export_s, 3),
"torch": torch.__version__,
"cuda": torch.version.cuda,
"gpu": torch.cuda.get_device_name(0),
}
(out_dir / "export_meta.json").write_text(json.dumps(meta, indent=2))
print("SS_FLOW_EXPORTED", json.dumps(meta), flush=True)
del wrapper, model, x, x_skl, timestep, cond
gc.collect(); torch.cuda.empty_cache()
if __name__ == "__main__":
main()
|