Spaces:
Running on Zero
Running on Zero
| """The same codec contract is used by offline production and online training.""" | |
| from __future__ import annotations | |
| import hashlib | |
| from collections.abc import Mapping | |
| from dataclasses import asdict, dataclass, field | |
| from typing import Any, Protocol | |
| from solarwm.config.loader import canonical_json | |
| from solarwm.errors import DataContractError | |
| class TensorSpec: | |
| name: str | |
| shape: tuple[int | None, ...] | |
| dtype: str | |
| def __post_init__(self) -> None: | |
| if not self.name or not self.shape or not self.dtype: | |
| raise DataContractError("tensor spec requires name, shape, and dtype") | |
| if any(value is not None and value < 0 for value in self.shape): | |
| raise DataContractError(f"tensor spec {self.name!r} has a negative dimension") | |
| class EncoderContract: | |
| schema: str | |
| family: str | |
| format_version: str | |
| pixel_frames: int | |
| latent_frames: int | |
| height: int | |
| width: int | |
| camera_convention: str | |
| tensors: tuple[TensorSpec, ...] | |
| extras: Mapping[str, Any] = field(default_factory=dict) | |
| def __post_init__(self) -> None: | |
| if self.schema != "solarwm.encoder.v1": | |
| raise DataContractError("encoder schema must be solarwm.encoder.v1") | |
| if not self.family or not self.format_version or not self.camera_convention: | |
| raise DataContractError("encoder family/format/camera convention are required") | |
| if min(self.pixel_frames, self.latent_frames, self.height, self.width) < 1: | |
| raise DataContractError("encoder dimensions must be positive") | |
| names = [spec.name for spec in self.tensors] | |
| if not names or len(names) != len(set(names)): | |
| raise DataContractError("encoder tensor names must be non-empty and unique") | |
| def as_dict(self) -> dict[str, Any]: | |
| return { | |
| **asdict(self), | |
| "tensors": [asdict(spec) for spec in self.tensors], | |
| } | |
| def digest(self) -> str: | |
| return hashlib.blake2s(canonical_json(self.as_dict())).hexdigest() | |
| class OnlineCodec(Protocol): | |
| """Model-owned raw-to-tensor implementation shared with offline encoding.""" | |
| def contract(self) -> EncoderContract: ... | |
| def encode( | |
| self, | |
| *, | |
| sample_id: str, | |
| pixels: Any, | |
| caption: str, | |
| camera: Any, | |
| seed: int, | |
| ) -> Mapping[str, Any]: ... | |
| def _dtype_name(value: Any) -> str: | |
| dtype = getattr(value, "dtype", None) | |
| if dtype is None: | |
| raise DataContractError(f"tensor {type(value).__name__} has no dtype") | |
| rendered = str(dtype).lower() | |
| for prefix in ("torch.", "numpy."): | |
| if rendered.startswith(prefix): | |
| rendered = rendered[len(prefix) :] | |
| return rendered | |
| def validate_encoded_tensors(values: Mapping[str, Any], contract: EncoderContract) -> None: | |
| expected = {spec.name: spec for spec in contract.tensors} | |
| if set(values) != set(expected): | |
| raise DataContractError( | |
| f"encoded tensor fields differ: got {sorted(values)}, expected {sorted(expected)}" | |
| ) | |
| for name, spec in expected.items(): | |
| value = values[name] | |
| try: | |
| shape = tuple(int(dimension) for dimension in value.shape) | |
| except (AttributeError, TypeError, ValueError) as exc: | |
| raise DataContractError(f"encoded field {name!r} has no valid shape") from exc | |
| if len(shape) != len(spec.shape) or any( | |
| expected_dimension is not None and actual != expected_dimension | |
| for actual, expected_dimension in zip(shape, spec.shape, strict=True) | |
| ): | |
| raise DataContractError( | |
| f"encoded field {name!r} shape {shape} != contract {spec.shape}" | |
| ) | |
| dtype = _dtype_name(value) | |
| if dtype != spec.dtype.lower(): | |
| raise DataContractError(f"encoded field {name!r} dtype {dtype!r} != {spec.dtype!r}") | |
| tensor_isfinite = getattr(value, "isfinite", None) | |
| if callable(tensor_isfinite): | |
| finite = tensor_isfinite() | |
| else: | |
| try: | |
| import numpy as np | |
| finite = np.isfinite(value) | |
| except (TypeError, ValueError) as exc: | |
| raise DataContractError( | |
| f"encoded field {name!r} cannot be checked for finiteness" | |
| ) from exc | |
| finite_all = finite.all() | |
| if hasattr(finite_all, "item"): | |
| finite_all = finite_all.item() | |
| if not bool(finite_all): | |
| raise DataContractError(f"encoded field {name!r} contains non-finite values") | |
| class EncodedPayload: | |
| """Serialized sample ready for deterministic tar packing.""" | |
| sample_id: str | |
| key: str | |
| source_sample_id: str | |
| start_frame: int | |
| source_frame_indices: tuple[int, ...] | |
| encoder_contract_digest: str | |
| members: Mapping[str, bytes] | |
| metadata: Mapping[str, Any] = field(default_factory=dict) | |
| def __post_init__(self) -> None: | |
| if not self.sample_id or not self.key or not self.source_sample_id: | |
| raise DataContractError("encoded payload lacks sample/key/source identity") | |
| if self.start_frame < 0 or not self.source_frame_indices: | |
| raise DataContractError("encoded payload has an invalid frame window") | |
| if self.source_frame_indices[0] != self.start_frame: | |
| raise DataContractError("encoded payload start differs from first source frame") | |
| if len(self.encoder_contract_digest) != 64: | |
| raise DataContractError( | |
| "encoded payload has an invalid encoder contract content digest" | |
| ) | |
| if not self.members: | |
| raise DataContractError("encoded payload has no serialized members") | |
| for suffix, value in self.members.items(): | |
| if not suffix or "/" in suffix or suffix in {".", ".."} or not isinstance(value, bytes): | |
| raise DataContractError(f"invalid encoded member suffix {suffix!r}") | |