openftgs 0.1.0__py3-none-any.whl

openftgs/__init__.py

@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+"""An independent clean-room implementation of FreeTimeGS, not affiliated with its authors (D-62).
+
+openftgs implements FreeTimeGS (Wang et al., CVPR 2025). The distribution and the import
+package are named openftgs, and the command-line interface ``ftgs``.
+
+The public API is listed in spec §2.1; each module's ``__all__`` names its part of it, and
+everything else is private.
+"""
+
+__all__: list[str] = []

openftgs/_caffe_trunks.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+"""D-89's two Caffe trunk networks, `vgg16_oxford` and `alexnet_bvlc`: lpips's uncalibrated
+distance on Oxford's VGG16 and BVLC's AlexNet, built from the convolutions `openftgs._weights`
+converts from the publishers' Caffe files (#384).
+
+Spec: D-89, §2.3's pretrained-trunk row. A private top-level module that imports nothing from
+openftgs: `openftgs.losses` passes it the tensors and wraps it in `PerceptualLoss` (#385).
+"""
+
+import functools
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Literal
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+__all__ = ["CAFFE_TRUNKS", "CaffeTrunk"]
+
+
+# The construction below is adapted from #374's prototype `candidates.py` (the owner's
+# Apache-2.0 research code, SHA-256
+# 4386e7ce8bc97a2f94ca6b9e0eb11cc0d5559b20ccc08b72cb59c85c8e3486eb; source log 449), its
+# `_CaffeTrunk`, `VGG16CC` and `AlexNetCC`: the construction #374 tested, which D-89 names.
+# The layers are the publishers' (VGG16's configuration D, source log 420; BVLC's AlexNet
+# prototxt, source log 421).
+
+
+@dataclass(frozen=True)
+class _Conv:
+    """A convolution under its Caffe layer name, followed by its ReLU (Caffe's in-place one)."""
+
+    name: str
+    stride: int = 1
+    padding: int = 1
+    groups: int = 1
+
+
+@dataclass(frozen=True)
+class _Pool:
+    """Max pooling in Caffe's ceiling mode (D-89)."""
+
+    kernel: int
+    stride: int
+
+
+# A step of a trunk, in network order: a convolution and its ReLU, a pooling, Caffe's local
+# response normalization ("lrn"), or a feature D-89's distance reads ("tap").
+_Step = _Conv | _Pool | Literal["tap", "lrn"]
+
+_VGG16: tuple[_Step, ...] = (
+    *(_Conv("conv1_1"), _Conv("conv1_2"), "tap"),
+    *(_Pool(2, 2), _Conv("conv2_1"), _Conv("conv2_2"), "tap"),
+    *(_Pool(2, 2), _Conv("conv3_1"), _Conv("conv3_2"), _Conv("conv3_3"), "tap"),
+    *(_Pool(2, 2), _Conv("conv4_1"), _Conv("conv4_2"), _Conv("conv4_3"), "tap"),
+    *(_Pool(2, 2), _Conv("conv5_1"), _Conv("conv5_2"), _Conv("conv5_3"), "tap"),
+)
+# D-89: the taps after relu1 to relu5, before the normalization and pooling that follow relu1
+# and relu2; conv2, conv4 and conv5 in two groups.
+_ALEXNET: tuple[_Step, ...] = (
+    *(_Conv("conv1", stride=4, padding=0), "tap", "lrn", _Pool(3, 2)),
+    *(_Conv("conv2", padding=2, groups=2), "tap", "lrn", _Pool(3, 2)),
+    *(_Conv("conv3"), "tap"),
+    *(_Conv("conv4", groups=2), "tap"),
+    *(_Conv("conv5", groups=2), "tap"),
+)
+_STEPS: dict[str, tuple[_Step, ...]] = {"vgg16_oxford": _VGG16, "alexnet_bvlc": _ALEXNET}
+# Q-263: the smallest height and width each trunk's layers take: AlexNet's conv1 gives 4
+# from 23, which its poolings take to 2 and 1; VGG16's ceiling pooling takes 1 to 1.
+_MIN_SIDE = {"vgg16_oxford": 1, "alexnet_bvlc": 23}
+
+# D-89: each release's mean pixel, BGR on Caffe's 0 to 255 scale: Oxford's as published
+# (source log 420); BVLC's the per-channel means of Caffe's ilsvrc_2012_mean.npy (source log
+# 426, which records them to three decimals) to eight decimals (source log 449), as #374
+# tested them, which D-89's table rounds to three (Q-264).
+_MEAN_BGR = {
+    "vgg16_oxford": (103.939, 116.779, 123.68),
+    "alexnet_bvlc": (104.00698793, 116.66876762, 122.67891434),  # Q-264
+}
+
+# The two trunks, in D-89's table's order.
+CAFFE_TRUNKS: tuple[str, ...] = tuple(_STEPS)
+
+# D-89: lpips's `normalize_tensor` ε (source log 435).
+_EPS = 1e-10
+
+
+def _unit(features: Tensor) -> Tensor:
+    """Each feature vector (dim 1) divided by its norm plus ε, as lpips's `normalize_tensor`."""
+    return features / (features.pow(2).sum(dim=1, keepdim=True).sqrt() + _EPS)
+
+
+class CaffeTrunk(torch.nn.Module):
+    """The Caffe trunk `name` (one of `CAFFE_TRUNKS`) on `tensors`, its convolutions'
+    `<layer>.weight` and `<layer>.bias` under Caffe's layer names, as `openftgs._weights`'s
+    `trunk_tensors` returns them (D-89). It computes in their dtype, and holds them, with the
+    mean pixel, as buffers that `.to()` moves and converts but no `state_dict` carries
+    (Q-263; §2.3: never redistributed)."""
+
+    def __init__(self, name: str, tensors: Mapping[str, Tensor]) -> None:
+        super().__init__()
+        self._name = name
+        self._steps = _STEPS[name]
+        convs = [step for step in self._steps if isinstance(step, _Conv)]
+        for conv in convs:
+            for part in ("weight", "bias"):
+                tensor = tensors[f"{conv.name}.{part}"]
+                self.register_buffer(f"{conv.name}_{part}", tensor, persistent=False)
+        dtype = tensors[f"{convs[0].name}.weight"].dtype
+        mean = torch.tensor(_MEAN_BGR[name], dtype=dtype).view(1, 3, 1, 1)
+        self.register_buffer("mean_bgr", mean, persistent=False)
+
+    def features(self, image: Tensor) -> list[Tensor]:
+        """The taps of D-89's distance for a [B, 3, H, W] RGB image in [0, 1]: Caffe's input,
+        BGR × 255 minus the mean pixel, through the trunk's layers in network order."""
+        x = image.flip(1) * 255.0 - self.get_buffer("mean_bgr")
+        taps: list[Tensor] = []
+        for step in self._steps:
+            if step == "tap":
+                taps.append(x)
+            elif step == "lrn":
+                x = F.local_response_norm(x, 5, alpha=1e-4, beta=0.75, k=1.0)
+            elif isinstance(step, _Pool):
+                x = F.max_pool2d(x, step.kernel, step.stride, ceil_mode=True)
+            else:
+                weight = self.get_buffer(f"{step.name}_weight")
+                bias = self.get_buffer(f"{step.name}_bias")
+                x = F.conv2d(
+                    x, weight, bias, stride=step.stride, padding=step.padding, groups=step.groups
+                )
+                x = F.relu(x)
+        return taps
+
+    def forward(self, x: Tensor, y: Tensor) -> Tensor:
+        """D-89's distance between the [B, 3, H, W] RGB images x (the render) and y (the
+        reference), per image: lpips's uncalibrated distance, its `lpips=False` (source log
+        350), over the five taps. Images of another dtype than the weights', or AlexNet's
+        sides below 23 pixels, are a ValueError (Q-263)."""
+        # Q-263: checked on the dtype and shape alone, so nothing synchronizes.
+        dtype = self.get_buffer("mean_bgr").dtype
+        if x.dtype != dtype or y.dtype != dtype:
+            raise ValueError(
+                f"the {self._name} trunk computes in its weights' dtype {dtype}, so it takes "
+                f"{dtype} images, got {x.dtype} and {y.dtype}"
+            )
+        side = _MIN_SIDE[self._name]
+        if min(x.shape[-2:]) < side or min(y.shape[-2:]) < side:
+            raise ValueError(
+                f"the {self._name} trunk needs H and W >= {side}, got {tuple(x.shape[-2:])} "
+                f"and {tuple(y.shape[-2:])}"
+            )
+        # D-89: the reference without a gradient.
+        with torch.no_grad():
+            reference = [_unit(tap) for tap in self.features(y)]
+        # Per tap, the squared differences summed over channels and averaged over space.
+        layers = [
+            (_unit(tap) - unit_y).pow(2).sum(dim=1).mean(dim=(1, 2))
+            for tap, unit_y in zip(self.features(x), reference, strict=True)
+        ]
+        return functools.reduce(torch.add, layers)

openftgs/_legacy.py

@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+"""The perceptual settings of revisions before 549o: LPIPS's two backbones as `loss.perc_net`,
+which `openftgs.config` refuses (#375, decided item 3) and which checkpoints saved before 549o
+still carry (Q-247).
+
+Such a checkpoint's model doesn't depend on the loss it trained with, so the readers that take
+only its model (`io.load_checkpoint`, `FreeTimeGaussians.load`, the CLI's model reader) validate
+its configuration with 549o's defaults in place of its `perc_net` and `lambda_perc`, in memory
+only; resuming it is refused (Q-247, Q-271). Plain JSON dicts in and out.
+
+Spec: §2.1 `PercNet` and `LossConfig`, §2.2's checkpoint contract, D-9; Q-247. A private
+top-level module that imports nothing from openftgs, so every band may import it, as the stage
+instrumentation (PRD #220).
+"""
+
+from collections.abc import Mapping
+from typing import Any
+
+__all__ = ["LPIPS_BACKBONES", "resume_refusal", "with_current_loss"]
+
+# #375 (decided item 3): LPIPS's backbones, the `perc_net` values revision 549o retired.
+LPIPS_BACKBONES = ("vgg", "alex")
+
+
+def _legacy_perc_net(config: Mapping[str, Any]) -> str | None:
+    """The `loss.perc_net` of a stored TrainConfig dict when it is one of LPIPS's backbones,
+    else None (another value, valid or not, or none)."""
+    loss = config.get("loss")
+    net = loss.get("perc_net") if isinstance(loss, Mapping) else None
+    return net if isinstance(net, str) and net in LPIPS_BACKBONES else None
+
+
+# Q-247: in memory only; the dict given is never changed, so `TrainerState.config` stays as saved.
+def with_current_loss(config: Mapping[str, Any]) -> dict[str, Any]:
+    """A copy of a stored TrainConfig dict that `config_from_dict` reads as 549o's: a legacy
+    `loss.perc_net` and its `loss.lambda_perc` left out, so they take 549o's defaults (ꟻLIP at
+    its calibrated weight, D-9); every other key, and any other dict, as it is."""
+    out = dict(config)
+    if _legacy_perc_net(config) is not None:
+        loss = config["loss"]
+        out["loss"] = {k: v for k, v in loss.items() if k not in ("perc_net", "lambda_perc")}
+    return out
+
+
+def resume_refusal(config: Mapping[str, Any], where: str) -> str | None:
+    """Q-247: why the run that `where` (a checkpoint, say) stores can't be resumed, when its
+    configuration names LPIPS's backbone; else None."""
+    net = _legacy_perc_net(config)
+    if net is None:
+        return None
+    return (
+        f"{where} was trained with loss.perc_net {net!r}, one of LPIPS's backbones, which "
+        f"revision 549o of the spec retired (D-9, D-89): its run can't be resumed, though ftgs "
+        f"eval, render and export still read its model (Q-247)"
+    )

openftgs/_shards.py

@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Thread-per-device sharding (D-54): independent jobs spread over shards, one worker thread per
+shard, without changing results.
+
+Spec: D-54 ("Independent work is spread across GPUs without changing results: RoMa pairs,
+evaluation frames, render requests"), §2.4 "Several GPUs". A private top-level module that
+imports nothing from openftgs, so every band may import it (PRD #220, as `openftgs._stages`):
+`init.roma_init` (Q-106), `render.render_request` and `eval.evaluate_images` (Q-175, Q-186)
+share it.
+
+A shard is whatever one worker needs for its device (a dense matcher, a model copy and its
+evaluator), built by the caller in its own thread, one shard after another, before the workers
+start: so nothing is downloaded, built or compiled from several threads at once. Workers take
+jobs from one shared queue. A shard's exception stops that shard and is raised in the caller's
+thread once every shard has stopped.
+"""
+
+import queue
+import threading
+from collections.abc import Callable, Iterable, Sequence
+from types import TracebackType
+from typing import Generic, TypeVar
+
+import torch
+
+__all__ = ["Shards", "devices", "run"]
+
+_S = TypeVar("_S")  # a shard: what one worker thread needs
+_J = TypeVar("_J")  # a job
+_R = TypeVar("_R")  # a job's result
+
+
+def devices(entries: Sequence[str | torch.device]) -> list[torch.device]:
+    """Q-106, Q-122, Q-186: one shard per entry of a non-empty sequence of device names, so
+    a repeated entry is another shard on the same device. A string, an empty sequence or an
+    entry that isn't a device is a ValueError."""
+    if isinstance(entries, str):  # Q-122: a string is a sequence of characters
+        raise ValueError(f"devices must be a sequence of device names, got the string {entries!r}")
+    out = []
+    for entry in entries:
+        try:
+            out.append(torch.device(entry))
+        except (RuntimeError, TypeError) as error:
+            raise ValueError(f"devices holds {entry!r}, which isn't a device: {error}") from None
+    if not out:
+        raise ValueError("devices must name at least one device")
+    return out
+
+
+class Shards(Generic[_S, _J, _R]):
+    """One worker thread per shard, taking jobs from a shared queue, as a context manager: the
+    threads start on entry and are joined on exit. Results come back in the order they finish,
+    with their job; a shard's exception is raised by `next_result`, and the shard that raised
+    takes no more jobs."""
+
+    def __init__(
+        self,
+        shards: Sequence[_S],
+        jobs: Iterable[_J],
+        work: Callable[[_S, _J], _R],
+        *,
+        name: str,
+    ) -> None:
+        self._jobs: queue.SimpleQueue[_J | None] = queue.SimpleQueue()
+        for job in jobs:
+            self._jobs.put(job)
+        for _ in shards:
+            self._jobs.put(None)  # one stop per worker, after every job
+        self._results: queue.SimpleQueue[tuple[_J, _R | BaseException]] = queue.SimpleQueue()
+        self._stop = threading.Event()
+        self._work = work
+        self._threads = [
+            threading.Thread(target=self._serve, args=(shard,), name=f"{name} {i}")
+            for i, shard in enumerate(shards)
+        ]
+
+    def _serve(self, shard: _S) -> None:
+        while not self._stop.is_set():
+            job = self._jobs.get()
+            if job is None:
+                return
+            try:
+                result: _R | BaseException = self._work(shard, job)
+            except BaseException as error:  # noqa: BLE001 - raised again in the caller's thread
+                self._results.put((job, error))
+                return
+            self._results.put((job, result))
+
+    def __enter__(self) -> "Shards[_S, _J, _R]":
+        for thread in self._threads:
+            thread.start()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self._stop.set()
+        for thread in self._threads:
+            thread.join()
+
+    def next_result(self) -> tuple[_J, _R]:
+        """The next job to finish and its result; a shard's exception is raised here."""
+        job, result = self._results.get()
+        if isinstance(result, BaseException):
+            raise result
+        return job, result
+
+
+def run(
+    shards: Sequence[_S], jobs: Sequence[_J], work: Callable[[_S, _J], _R], *, name: str
+) -> list[_R]:
+    """Every job's result, in the jobs' order, done by `work` on the shards' worker threads
+    (named `name` and the shard's number). The first exception a shard raises is raised here,
+    once every shard has stopped."""
+    done: dict[int, _R] = {}
+
+    def numbered(shard: _S, job: tuple[int, _J]) -> _R:
+        return work(shard, job[1])
+
+    with Shards(shards, list(enumerate(jobs)), numbered, name=name) as running:
+        while len(done) < len(jobs):
+            (index, _), result = running.next_result()
+            done[index] = result
+    return [done[index] for index in range(len(jobs))]