File size: 2,248 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 | from __future__ import annotations
import argparse
import json
from pathlib import Path
import sys
ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
sys.path.insert(0, str(ROOT))
from orbitquant_wan_a2.source_quant import build_packed_from_official_source
def main():
ap = argparse.ArgumentParser(
description="Directly quantize the freshly-downloaded official Wan-Animate-2 distilled BF16 transformer into packed OrbitQuant W4."
)
ap.add_argument("--transformer", default="/content/Wan2_2_Animate_Repo/transformer")
ap.add_argument("--output", default="/content/OrbitQuant_Animate2_W4A4_PACKED")
ap.add_argument("--seed", type=int, default=20260702)
ap.add_argument("--device", default="auto")
ap.add_argument("--row-chunk", type=int, default=512)
ap.add_argument("--max-shard-gib", type=float, default=0.75)
ap.add_argument("--model-revision", default=None)
ap.add_argument("--overwrite", action="store_true")
args = ap.parse_args()
manifest = build_packed_from_official_source(
args.transformer,
args.output,
seed=args.seed,
bits=4,
device=args.device,
row_chunk=args.row_chunk,
max_shard_gib=args.max_shard_gib,
overwrite=args.overwrite,
model_revision=args.model_revision,
)
audits = [v["audit"] for v in manifest["targets"].values()]
if len(audits) != 480 or not all(
a["code_exact_fraction"] == 1.0
and a["scale_exact_fraction"] == 1.0
and a["fake_bf16_exact_fraction"] == 1.0
for a in audits
):
raise RuntimeError("direct-source packed W4 audit did not pass 480/480")
if manifest["passthrough_count"] != 823 or manifest["full_transformer_tensor_count"] != 1303:
raise RuntimeError("wrong final transformer inventory")
print("\nDIRECT SOURCE -> PACKED ORBITQUANT W4 COMPLETE")
print(" source tensors : 1303 / 1303")
print(" packed W4 : 480 / 480")
print(" passthrough : 823 / 823")
print(" row scales : BF16")
print(" packed GiB : %.3f" % (manifest["packed_bytes"] / 2**30))
print(" output :", Path(args.output))
if __name__ == "__main__":
main()
|