Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """Expand a validated DiT NPZ into browser-fetchable float32 tensor files.""" | |
| from __future__ import annotations | |
| import argparse | |
| import hashlib | |
| import json | |
| from pathlib import Path | |
| from typing import Any | |
| SHAPES = { | |
| "latent": (1, 8192, 16), | |
| "camera": (1, 1, 5), | |
| "t": (1,), | |
| "feature1": (1, 4101, 1280), | |
| "feature2": (1, 4101, 128), | |
| "pred_latent": (1, 8192, 16), | |
| "pred_camera": (1, 1, 5), | |
| } | |
| def parse_args() -> argparse.Namespace: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--fixture", type=Path, required=True, help="Validated DiT NPZ fixture.") | |
| parser.add_argument("--output-dir", type=Path, required=True, help="Browser fixture directory.") | |
| return parser.parse_args() | |
| def sha256(path: Path) -> str: | |
| digest = hashlib.sha256() | |
| with path.open("rb") as stream: | |
| while block := stream.read(8 * 1024 * 1024): | |
| digest.update(block) | |
| return digest.hexdigest() | |
| def main() -> None: | |
| try: | |
| import numpy as np | |
| except ImportError as exc: | |
| raise SystemExit(f"NumPy is required: {exc}") from exc | |
| args = parse_args() | |
| fixture = args.fixture.expanduser().resolve() | |
| output = args.output_dir.expanduser().resolve() | |
| if not fixture.is_file(): | |
| raise FileNotFoundError(fixture) | |
| output.mkdir(parents=True, exist_ok=True) | |
| files: dict[str, Any] = {} | |
| with np.load(fixture, allow_pickle=False) as values: | |
| missing = [name for name in SHAPES if name not in values.files] | |
| if missing: | |
| raise KeyError(f"{fixture} is missing {missing}") | |
| for name, shape in SHAPES.items(): | |
| array = np.ascontiguousarray(values[name], dtype="<f4") | |
| if tuple(array.shape) != shape: | |
| raise ValueError(f"{name} has shape {array.shape}; expected {shape}") | |
| if not np.isfinite(array).all(): | |
| raise ValueError(f"{name} contains NaN or infinity") | |
| path = output / f"{name}.f32" | |
| array.tofile(path) | |
| files[name] = { | |
| "path": path.name, | |
| "dtype": "float32-le", | |
| "shape": list(shape), | |
| "elements": int(array.size), | |
| "bytes": path.stat().st_size, | |
| "sha256": sha256(path), | |
| } | |
| metadata = str(values["metadata"]) if "metadata" in values.files else None | |
| manifest = { | |
| "format": "triposplat-dit-browser-fixture-v1", | |
| "source": str(fixture), | |
| "metadata": json.loads(metadata) if metadata else None, | |
| "files": files, | |
| } | |
| manifest_path = output / "manifest.json" | |
| manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8") | |
| print(f"Wrote {manifest_path}") | |
| if __name__ == "__main__": | |
| main() | |