94 lines
3.6 KiB
Python
94 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Export the unchanged InternVLA Cosmos encoder as a fixed-shape graph."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import importlib.util
|
|
import json
|
|
import sys
|
|
import time
|
|
import types
|
|
from pathlib import Path
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as source:
|
|
for block in iter(lambda: source.read(8 << 20), b""):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--encoder", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--report", type=Path, required=True)
|
|
parser.add_argument("--batch-size", type=int, default=6)
|
|
args = parser.parse_args()
|
|
|
|
import torch
|
|
|
|
source_dir = Path(__file__).resolve().parents[3] / "vllm_omni" / "diffusion" / "models" / "internvla_a1"
|
|
package = types.ModuleType("internvla_cosmos_export")
|
|
package.__path__ = [str(source_dir)]
|
|
sys.modules[package.__name__] = package
|
|
spec = importlib.util.spec_from_file_location(f"{package.__name__}.model_cosmos", source_dir / "model_cosmos.py")
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError("cannot load pinned Cosmos source")
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
encoder_path = args.encoder.resolve(strict=True)
|
|
encoder_model = module.load_cosmos_component(encoder_path, component="encoder", device="cpu")
|
|
|
|
class LatentOnly(torch.nn.Module):
|
|
def __init__(self, inner: torch.nn.Module) -> None:
|
|
super().__init__()
|
|
self.inner = inner
|
|
|
|
def forward(self, pixels: torch.Tensor) -> torch.Tensor:
|
|
result = self.inner(pixels)
|
|
return result if isinstance(result, torch.Tensor) else result[0]
|
|
|
|
model = LatentOnly(encoder_model).eval()
|
|
image = torch.linspace(-1, 1, args.batch_size * 3 * 256 * 256, dtype=torch.float32).view(
|
|
args.batch_size, 3, 256, 256
|
|
)
|
|
started = time.perf_counter()
|
|
with torch.no_grad():
|
|
reference = model(image)
|
|
exported = torch.export.export(model, (image,), strict=False)
|
|
actual = exported.module()(image)
|
|
export_s = time.perf_counter() - started
|
|
reference_tensor = reference
|
|
actual_tensor = actual
|
|
if not torch.allclose(reference_tensor, actual_tensor, atol=1e-5, rtol=1e-5):
|
|
raise RuntimeError("exported CPU graph does not match source module")
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
torch.export.save(exported, args.output)
|
|
report = {
|
|
"status": "passed",
|
|
"scope": "fixed-shape Cosmos CI8x8 encoder graph; no InternVLA action-policy claim",
|
|
"encoder_sha256": sha256(encoder_path),
|
|
"source_sha256": {name: sha256(source_dir / name) for name in ("model_cosmos.py", "cosmos_ci_torch.py")},
|
|
"torch": torch.__version__,
|
|
"input_shape": list(image.shape),
|
|
"input_dtype": str(image.dtype),
|
|
"output_shape": list(reference_tensor.shape),
|
|
"max_abs_export_error": float((reference_tensor - actual_tensor).abs().max()),
|
|
"export_s": export_s,
|
|
"artifact_sha256": sha256(args.output),
|
|
"artifact_bytes": args.output.stat().st_size,
|
|
"artifact_path": str(args.output.resolve()),
|
|
}
|
|
args.report.parent.mkdir(parents=True, exist_ok=True)
|
|
args.report.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
|
|
print(json.dumps(report, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|