Spaces:
Running on Zero
Running on Zero
File size: 6,045 Bytes
7b1830e | 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 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | """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
@dataclass(frozen=True)
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")
@dataclass(frozen=True)
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],
}
@property
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."""
@property
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")
@dataclass(frozen=True)
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}")
|