| |
| """Differentiable CPU reference for the statically recovered NR temporal envelope. |
| |
| Evidence: docs/nr-temporal-reconstruction.md and its JSON report. This is not a |
| complete NR model, a CUDA sampler, or a numerically validated NVIDIA replacement. |
| Images use [..., channels] here as a convenient reference layout, not the native |
| packed activation layout. Caller-supplied noise and warped history must already |
| match the required pixels: no RNG, motion convention or texture filter is guessed. |
| |
| ``smooth`` evaluates the recovered equations in the input floating dtype for |
| training/gradcheck. ``half_staging`` additionally reproduces the observed half |
| conversion boundaries, but does not reproduce PTX FMA, FTZ, transcendental or |
| texture-hardware behavior. Autograd through half casts is a surrogate gradient; |
| it is not the mathematical derivative of discrete quantization. |
| |
| Example, after supplying a separately validated network and sampler:: |
| |
| pre = build_preblock_channels(color, noise_rgb=noise, |
| warped_history_rgb=warped_pre, history_enabled=pre_history_enabled, |
| controls=controls, capabilities=capabilities) |
| projected = network(pre) # omitted from this module |
| post = reconstruct_post(projected, current_rgb=color, |
| warped_history_rgb=warped_post, history_enabled=post_history_enabled) |
| display = compose_display(post.temporal_rgb, color, intensity=0.8) |
| next_history = post.temporal_rgb # distinct from artistically blended display |
| |
| The caller owns actual RGBA16F state storage, texture sampling, all history gates, |
| reset/counter lifecycle, UI composition and any training detach policy. |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Literal |
|
|
| import torch |
|
|
| Arithmetic = Literal["smooth", "half_staging"] |
| Scalar = float | torch.Tensor |
|
|
|
|
| @dataclass(frozen=True) |
| class ArtistControls: |
| """Native frame controls; style_count must come from the selected model.""" |
|
|
| tone: Scalar |
| structure: Scalar |
| skin_structure: Scalar |
| auto_mask: bool |
| style_id: int |
| style_count: int |
|
|
|
|
| @dataclass(frozen=True) |
| class ConditioningCapabilities: |
| """Explicit model capability bytes; this reference does not infer them.""" |
|
|
| local_controls: bool |
| semantic_conditioning: bool |
| style: bool |
|
|
|
|
| @dataclass(frozen=True) |
| class PostResult: |
| candidate_rgb: torch.Tensor |
| temporal_rgb: torch.Tensor |
| history_weight: torch.Tensor |
| output_alpha: torch.Tensor |
|
|
| @property |
| def temporal_rgba(self) -> torch.Tensor: |
| return torch.cat((self.temporal_rgb, self.output_alpha), dim=-1) |
|
|
|
|
| def _arithmetic(mode: Arithmetic) -> None: |
| if mode not in ("smooth", "half_staging"): |
| raise ValueError("arithmetic must be 'smooth' or 'half_staging'") |
|
|
|
|
| def _bool(value: bool, name: str) -> None: |
| if type(value) is not bool: |
| raise ValueError(f"{name} must be an explicit bool") |
|
|
|
|
| def _image(value: torch.Tensor, name: str, channels: tuple[int, ...], |
| reference: torch.Tensor | None = None) -> torch.Tensor: |
| if not isinstance(value, torch.Tensor): |
| raise ValueError(f"{name} must be a tensor") |
| if value.device.type != "cpu" or value.dtype not in (torch.float32, torch.float64): |
| raise ValueError(f"{name} must be CPU float32 or float64") |
| if value.ndim < 2 or value.shape[-1] not in channels or value.numel() == 0: |
| raise ValueError(f"{name} must have nonempty [..., {channels}] image shape") |
| if reference is not None and (value.shape[:-1] != reference.shape[:-1] |
| or value.dtype != reference.dtype): |
| raise ValueError(f"{name} must match reference pixel shape and dtype") |
| if not bool(torch.isfinite(value).all()): |
| raise ValueError(f"{name} must be finite") |
| return value |
|
|
|
|
| def _scalar(value: Scalar, reference: torch.Tensor, name: str, |
| *, allow_nonfinite: bool = False) -> torch.Tensor: |
| if isinstance(value, torch.Tensor): |
| if value.ndim != 0 or value.device.type != "cpu" or value.dtype != reference.dtype: |
| raise ValueError(f"{name} must be a scalar CPU tensor of the image dtype") |
| result = value |
| else: |
| if isinstance(value, bool) or not isinstance(value, (int, float)): |
| raise ValueError(f"{name} must be a real scalar or scalar tensor") |
| result = reference.new_tensor(value) |
| if not allow_nonfinite and not bool(torch.isfinite(result)): |
| raise ValueError(f"{name} must be finite") |
| return result |
|
|
|
|
| def _half(value: torch.Tensor, arithmetic: Arithmetic) -> torch.Tensor: |
| |
| return (value.to(torch.float16).to(value.dtype) |
| if arithmetic == "half_staging" else value) |
|
|
|
|
| def normalize_color(color_rgb: torch.Tensor, *, arithmetic: Arithmetic = "smooth") -> torch.Tensor: |
| """PTX pre color normalization; host S=1/16, hence multiplier=1/8. |
| |
| Evidence: module_0 lines 15080-15089,15320-15358; host RVA0x60f29. |
| No gamma transform, scaling/resampling, or range assumption is added here. |
| """ |
| _arithmetic(arithmetic) |
| _image(color_rgb, "color_rgb", (3,)) |
| sampled = _half(color_rgb, arithmetic) |
| centered = _half(sampled - 0.5, arithmetic) |
| return _half(centered * 0.125, arithmetic) |
|
|
|
|
| def conditioning_channels(reference_rgb: torch.Tensor, *, controls: ArtistControls, |
| capabilities: ConditioningCapabilities, |
| control_mask: torch.Tensor | None = None, |
| arithmetic: Arithmetic = "smooth") -> torch.Tensor: |
| """Return channels 10..14: style, tone, structure, semantic A, semantic B. |
| |
| Mask pixels must already be sampled at the preblock coordinates. Green scales |
| tone, blue scales structure; red is used later by compose_display. Model |
| capabilities are separate from the user auto-mask option. Branch/sentinel |
| boundaries are discrete and are not smoothly differentiable. |
| """ |
| _arithmetic(arithmetic) |
| _image(reference_rgb, "reference_rgb", (3,)) |
| for name in ("local_controls", "semantic_conditioning", "style"): |
| _bool(getattr(capabilities, name), name) |
| _bool(controls.auto_mask, "auto_mask") |
| if (type(controls.style_id) is not int or type(controls.style_count) is not int |
| or not 0 <= controls.style_id <= 0xFFFFFFFF or controls.style_count < 0): |
| raise ValueError("style_id must be uint32 and style_count a nonnegative integer") |
| if control_mask is not None: |
| _image(control_mask, "control_mask", (3, 4), reference_rgb) |
| tone = _scalar(controls.tone, reference_rgb, "tone") |
| structure = _scalar(controls.structure, reference_rgb, "structure") |
| skin = _scalar(controls.skin_structure, reference_rgb, "skin_structure") |
| zero = reference_rgb.new_zeros(()) |
| one = reference_rgb.new_ones(()) |
| negative_one = -one |
|
|
| |
| if controls.auto_mask and control_mask is None: |
| semantic_a = torch.where(skin >= 0, skin, structure) |
| semantic_b = structure |
| else: |
| semantic_a = negative_one |
| semantic_b = negative_one |
|
|
| local_tone = tone if capabilities.local_controls else zero |
| local_structure = structure if capabilities.local_controls else zero |
| if capabilities.semantic_conditioning and control_mask is None: |
| |
| |
| any_semantic = torch.maximum(semantic_a, semantic_b) >= 0 |
| semantic_a = torch.where(any_semantic, |
| torch.where(semantic_a >= 0, semantic_a, local_structure), |
| negative_one) |
| semantic_b = torch.where(any_semantic, |
| torch.where(semantic_b >= 0, semantic_b, local_structure), |
| negative_one) |
| local_structure = torch.where(any_semantic, one, local_structure) |
| else: |
| sentinel = negative_one if capabilities.semantic_conditioning else zero |
| semantic_a = sentinel |
| semantic_b = sentinel |
| if control_mask is not None: |
| local_tone = local_tone * control_mask[..., 1] |
| local_structure = local_structure * control_mask[..., 2] |
|
|
| index = min(controls.style_id, controls.style_count - 1) if controls.style_count else 0 |
| style = reference_rgb.new_tensor(index / 128.0 if capabilities.style else 0.0) |
| pixel_shape = reference_rgb.shape[:-1] |
| values = (style, local_tone, local_structure, semantic_a, semantic_b) |
| return _half(torch.stack([v.expand(pixel_shape) for v in values], dim=-1), arithmetic) |
|
|
|
|
| def build_preblock_channels(current_rgb: torch.Tensor, *, noise_rgb: torch.Tensor, |
| warped_history_rgb: torch.Tensor | None, |
| history_enabled: bool, controls: ArtistControls, |
| capabilities: ConditioningCapabilities, |
| control_mask: torch.Tensor | None = None, |
| arithmetic: Arithmetic = "smooth") -> torch.Tensor: |
| """Construct the verified 16-channel preprojection input, not packed tensors. |
| |
| The caller supplies native-matched noise and independently warped history. |
| history_enabled must already incorporate applicable resource, reset, counter |
| and global gates. A disabled history path uses current RGB in channels7..9. |
| No native frame dimensions, padding or texture sampling are inferred. |
| """ |
| _arithmetic(arithmetic) |
| _bool(history_enabled, "history_enabled") |
| _image(current_rgb, "current_rgb", (3,)) |
| _image(noise_rgb, "noise_rgb", (3,), current_rgb) |
| if warped_history_rgb is not None: |
| _image(warped_history_rgb, "warped_history_rgb", (3,), current_rgb) |
| if history_enabled and warped_history_rgb is None: |
| raise ValueError("enabled pre history requires explicitly sampled warped_history_rgb") |
| current = normalize_color(current_rgb, arithmetic=arithmetic) |
| history = (normalize_color(warped_history_rgb, arithmetic=arithmetic) |
| if history_enabled else current) |
| artist = conditioning_channels(current_rgb, controls=controls, capabilities=capabilities, |
| control_mask=control_mask, arithmetic=arithmetic) |
| constant = torch.ones_like(current[..., :1]) |
| return torch.cat((_half(noise_rgb, arithmetic), constant, current, history, |
| artist, torch.zeros_like(constant)), dim=-1) |
|
|
|
|
| def reconstruct_post(projected_rgba: torch.Tensor, *, current_rgb: torch.Tensor | None, |
| warped_history_rgb: torch.Tensor | None, history_enabled: bool, |
| learned_blend_scale: Scalar | None = None, |
| final_color_mode: bool = True, |
| arithmetic: Arithmetic = "smooth") -> PostResult: |
| """Recover RGB residual and history gate from four already-projected values. |
| |
| These inputs are the final four projection values, not the NR input or hidden |
| activations. Caller applies the post history/resource eligibility conditions |
| and supplies its already-sampled history independently of the pre warp. |
| The sigmoid uses PyTorch, not the PTX approximate ex2/reciprocal sequence. |
| temporal_rgb remains unquantized by default; native RGBA16F state persistence |
| and any training detach operation are caller decisions, not hidden here. |
| """ |
| _arithmetic(arithmetic) |
| _bool(history_enabled, "history_enabled") |
| _bool(final_color_mode, "final_color_mode") |
| _image(projected_rgba, "projected_rgba", (4,)) |
| if current_rgb is not None: |
| _image(current_rgb, "current_rgb", (3,), projected_rgba) |
| if warped_history_rgb is not None: |
| _image(warped_history_rgb, "warped_history_rgb", (3,), projected_rgba) |
| active = history_enabled and final_color_mode |
| if active and warped_history_rgb is None: |
| raise ValueError("enabled final post history requires explicitly sampled warped_history_rgb") |
|
|
| projection = _half(projected_rgba, arithmetic) |
| internal_rgb = projection[..., :3] * 0.03125 |
| if current_rgb is not None: |
| internal_rgb = internal_rgb + (current_rgb * 0.125 - 0.0625) |
| candidate = ((internal_rgb * 8.0 + 0.5).clamp(0.0, 1.0) |
| if final_color_mode else internal_rgb) |
| if active: |
| scale = _scalar(1.0 if learned_blend_scale is None else learned_blend_scale, |
| projected_rgba, "learned_blend_scale", allow_nonfinite=True) |
| scale = _half(scale, arithmetic) |
| scale = torch.where(torch.isfinite(scale), scale, torch.zeros_like(scale)).clamp(0.0, 1.0) |
| weight = (torch.sigmoid(projection[..., 3:4]) * scale).clamp(0.0, 1.0) |
| temporal = candidate + weight * (warped_history_rgb - candidate) |
| else: |
| weight = torch.zeros_like(projection[..., :1]) |
| temporal = candidate |
| output_alpha = torch.full_like(weight, 1.0 if final_color_mode else 0.0) |
| return PostResult(candidate, temporal, weight, output_alpha) |
|
|
|
|
| def compose_display(temporal_rgb: torch.Tensor, source_rgb: torch.Tensor, *, |
| intensity: Scalar, control_mask: torch.Tensor | None = None) -> torch.Tensor: |
| """Separate final artistic blend; never overwrites the temporal/history result. |
| |
| Source and mask must already be sampled in output coordinates. Simple-blend |
| PTX saturates intensity (module_0:123734..123736); mask mode saturates the |
| product intensity*mask.red (195329..195340), then both saturate final RGB. |
| UI composition is outside this function's recovered contract. |
| """ |
| _image(temporal_rgb, "temporal_rgb", (3,)) |
| _image(source_rgb, "source_rgb", (3,), temporal_rgb) |
| strength = _scalar(intensity, temporal_rgb, "intensity") |
| if control_mask is None: |
| weight = strength.clamp(0.0, 1.0) |
| else: |
| _image(control_mask, "control_mask", (3, 4), temporal_rgb) |
| weight = (strength * control_mask[..., :1]).clamp(0.0, 1.0) |
| return (source_rgb + weight * (temporal_rgb - source_rgb)).clamp(0.0, 1.0) |
|
|
|
|
| def cubic_history_weights(fraction_xy: torch.Tensor) -> torch.Tensor: |
| """Five normalized cubic weights in left, up, center, down, right order. |
| |
| Fractions refer to the native history pixel-center lattice and are supplied |
| explicitly. This computes coefficients only; texture lookup coordinates, |
| interpolation, rect remapping and addressing are not implemented here. |
| Evidence: module_0:15450..15541. Fractions are clamped as in the PTX. |
| """ |
| _image(fraction_xy, "fraction_xy", (2,)) |
| t = fraction_xy.clamp(0.0, 1.0) |
| square = t * t |
| cube = t * square |
| w0 = -0.5 * (t + cube) + square |
| w1 = 1.5 * cube - 2.5 * square + 1.0 |
| w3 = 0.5 * (cube - square) |
| w2 = 1.0 - w0 - w1 - w3 |
| middle = w1 + w2 |
| weights = torch.stack((w0[..., 0] * middle[..., 1], |
| w0[..., 1] * middle[..., 0], |
| middle[..., 0] * middle[..., 1], |
| w3[..., 1] * middle[..., 0], |
| w3[..., 0] * middle[..., 1]), dim=-1) |
| return weights / weights.sum(dim=-1, keepdim=True) |
|
|
|
|
| def combine_cubic_history_samples(samples: tuple[torch.Tensor, ...], |
| fraction_xy: torch.Tensor) -> torch.Tensor: |
| """Combine five externally obtained RGB texture samples differentiably. |
| |
| samples order is left, up, center, down, right. In particular these are five |
| *texture lookup results*, not five raw integer texels. Passing raw neighbors |
| would omit the native middle-pair interpolation and is not equivalent. |
| """ |
| if not isinstance(samples, (tuple, list)) or len(samples) != 5: |
| raise ValueError("provide exactly five texture sample tensors") |
| _image(fraction_xy, "fraction_xy", (2,)) |
| for index, sample in enumerate(samples): |
| _image(sample, f"samples[{index}]", (3,), fraction_xy) |
| weights = cubic_history_weights(fraction_xy) |
| return (torch.stack(tuple(samples), dim=-2) * weights.unsqueeze(-1)).sum(dim=-2) |
|
|
