#!/usr/bin/env python3 """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: # Keep a uniform reference dtype for concatenation and CPU gradient work. 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 # Host RVA0x1aa4b..0x1aaa1 derives these from native controls before model gates. 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: # PTX lines15152..15193: nonnegative semantic input switches the common # structure channel to 1; individual negatives fall back to gated structure. 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)