#!/usr/bin/env python3 """CPU equation, branch and gradient checks; no NVIDIA inference-parity claim.""" import unittest import torch from torch.autograd import gradcheck from nr_temporal_reference import ( ArtistControls, ConditioningCapabilities, build_preblock_channels, combine_cubic_history_samples, compose_display, conditioning_channels, cubic_history_weights, normalize_color, reconstruct_post, ) CAPS = ConditioningCapabilities(local_controls=True, semantic_conditioning=True, style=True) def controls(**overrides): values = dict(tone=0.6, structure=0.4, skin_structure=0.2, auto_mask=False, style_id=2, style_count=3) return ArtistControls(**(values | overrides)) def tensor(value, *, grad=False): return torch.tensor(value, dtype=torch.float64, requires_grad=grad) class TemporalReferenceTests(unittest.TestCase): @classmethod def setUpClass(cls): torch.set_num_threads(2) def test_normalization_known_domain(self): self.assertTrue(torch.equal(normalize_color(tensor([[0, 0.5, 1]])), tensor([[-0.0625, 0, 0.0625]]))) def test_half_staging_preserves_separate_conversion_boundaries(self): color = tensor([[0.5001, 0.12345, 0.9999]]) expected = ((color.half() - torch.tensor(0.5, dtype=torch.float16)) * torch.tensor(0.125, dtype=torch.float16)).double() actual = normalize_color(color, arithmetic="half_staging") self.assertTrue(torch.equal(actual, expected)) self.assertFalse(torch.equal(actual, normalize_color(color))) def test_exact_channel_order_and_explicit_history(self): color = tensor([[0, 0.5, 1]]) history = tensor([[1, 0, 0.5]]) noise = tensor([[-2, 3, 4]]) pre = build_preblock_channels(color, noise_rgb=noise, warped_history_rgb=history, history_enabled=True, controls=controls(), capabilities=CAPS) expected = tensor([[-2, 3, 4, 1, -0.0625, 0, 0.0625, 0.0625, -0.0625, 0, 2/128, 0.6, 0.4, -1, -1, 0]]) self.assertTrue(torch.equal(pre, expected)) no_history = build_preblock_channels(color, noise_rgb=noise, warped_history_rgb=history, history_enabled=False, controls=controls(), capabilities=CAPS) self.assertTrue(torch.equal(no_history[..., 7:10], no_history[..., 4:7])) with self.assertRaisesRegex(ValueError, "explicitly sampled"): build_preblock_channels(color, noise_rgb=noise, warped_history_rgb=None, history_enabled=True, controls=controls(), capabilities=CAPS) def test_auto_semantics_and_negative_skin_fallback(self): color = tensor([[0.5, 0.5, 0.5]]) for skin, expected_skin in ((0.2, 0.2), (-1.0, 0.4)): channels = conditioning_channels(color, controls=controls(auto_mask=True, skin_structure=skin), capabilities=CAPS) torch.testing.assert_close(channels, tensor([[2/128, 0.6, 1, expected_skin, 0.4]])) # A nonnegative zero is an explicit semantic value, not the -1 sentinel. channels = conditioning_channels(color, controls=controls(auto_mask=True, skin_structure=0.0), capabilities=CAPS) self.assertEqual(channels[0, 2:].tolist(), [1.0, 0.0, 0.4]) def test_mask_rgb_roles_and_alpha_ignored(self): color = tensor([[0.5, 0.5, 0.5]]) mask = tensor([[0.1, 0.5, 0.25, 999]]) result = conditioning_channels(color, controls=controls(auto_mask=True), capabilities=CAPS, control_mask=mask) torch.testing.assert_close(result, tensor([[2/128, 0.3, 0.1, -1, -1]])) output = compose_display(tensor([[1, 1, 1]]), tensor([[0, 0, 0]]), intensity=0.5, control_mask=mask) torch.testing.assert_close(output, tensor([[0.05, 0.05, 0.05]])) def test_model_capabilities_and_style_bounds(self): color = tensor([[0.5, 0.5, 0.5]]) disabled = ConditioningCapabilities(False, False, False) self.assertTrue(torch.equal(conditioning_channels(color, controls=controls(auto_mask=True), capabilities=disabled), tensor([[0, 0, 0, 0, 0]]))) # Semantic values are host-derived before local-control capability gates. semantic_only = ConditioningCapabilities(False, True, False) result = conditioning_channels(color, controls=controls(auto_mask=True), capabilities=semantic_only) torch.testing.assert_close(result, tensor([[0, 0, 1, 0.2, 0.4]])) for style_id, count, expected in ((100, 3, 2/128), (2, 0, 0), (0, 3, 0)): channels = conditioning_channels(color, controls=controls(style_id=style_id, style_count=count), capabilities=CAPS) self.assertEqual(float(channels[0, 0]), expected) def test_post_residual_gate_and_distinct_display(self): color = tensor([[0.4, 0.5, 0.6]]) projected = tensor([[0.4, -0.4, 0, 0]]) history = tensor([[0.9, 0.9, 0.9]]) result = reconstruct_post(projected, current_rgb=color, warped_history_rgb=history, history_enabled=True, learned_blend_scale=0.5) torch.testing.assert_close(result.candidate_rgb, tensor([[0.5, 0.4, 0.6]])) torch.testing.assert_close(result.history_weight, tensor([[0.25]])) torch.testing.assert_close(result.temporal_rgb, tensor([[0.6, 0.525, 0.675]])) snapshot = result.temporal_rgb.clone() display = compose_display(result.temporal_rgb, color, intensity=0) self.assertTrue(torch.equal(display, color)) self.assertTrue(torch.equal(result.temporal_rgb, snapshot)) self.assertFalse(torch.equal(result.temporal_rgb, display)) self.assertEqual(result.temporal_rgba.shape[-1], 4) self.assertTrue(torch.equal(result.temporal_rgba[..., 3:], tensor([[1]]))) def test_blend_scale_nonfinite_clamping_and_default(self): projected = tensor([[0, 0, 0, 0]]) color = tensor([[0.2, 0.2, 0.2]]) history = tensor([[0.8, 0.8, 0.8]]) for scale, expected_weight in ((None, 0.5), (-1, 0), (2, 0.5), (float("nan"), 0), (float("inf"), 0), (-float("inf"), 0)): with self.subTest(scale=scale): result = reconstruct_post(projected, current_rgb=color, warped_history_rgb=history, history_enabled=True, learned_blend_scale=scale) self.assertEqual(float(result.history_weight), expected_weight) self.assertTrue(bool(torch.isfinite(result.temporal_rgb).all())) def test_no_history_and_internal_color_mode(self): projected = tensor([[0.4, -0.4, 0, 0]]) color = tensor([[0.4, 0.5, 0.6]]) plain = reconstruct_post(projected, current_rgb=color, warped_history_rgb=None, history_enabled=False) self.assertTrue(torch.equal(plain.temporal_rgb, plain.candidate_rgb)) self.assertEqual(float(plain.history_weight), 0) internal = reconstruct_post(projected, current_rgb=color, warped_history_rgb=None, history_enabled=True, final_color_mode=False) torch.testing.assert_close(internal.temporal_rgb, tensor([[0, -0.0125, 0.0125]])) self.assertEqual(float(internal.output_alpha), 0) without_source = reconstruct_post(tensor([[0, 0, 0, 0]]), current_rgb=None, warped_history_rgb=None, history_enabled=False) self.assertTrue(torch.equal(without_source.temporal_rgb, tensor([[0.5, 0.5, 0.5]]))) with self.assertRaisesRegex(ValueError, "explicitly sampled"): reconstruct_post(projected, current_rgb=color, warped_history_rgb=None, history_enabled=True) def test_display_intensity_boundary_and_mask_product_clamp(self): zero = tensor([[0, 0, 0]]) one = tensor([[1, 1, 1]]) for intensity, expected in ((-2, zero), (0, zero), (1, one), (4, one)): self.assertTrue(torch.equal(compose_display(one, zero, intensity=intensity), expected)) # Clamp the mask product, not intensity before multiplication. torch.testing.assert_close(compose_display(one, zero, intensity=2, control_mask=tensor([[0.25, 1, 1]])), tensor([[0.5, 0.5, 0.5]])) def test_cubic_known_weights_and_constant_preservation(self): fractions = tensor([[0, 0], [0.5, 0.5], [-1, 2]]) weights = cubic_history_weights(fractions) torch.testing.assert_close(weights[0], tensor([0, 0, 1, 0, 0])) torch.testing.assert_close(weights[1], tensor([-1/14, -1/14, 9/7, -1/14, -1/14])) torch.testing.assert_close(weights.sum(-1), tensor([1, 1, 1])) sample = tensor([[0.2, 0.4, 0.8], [0.6, 0.7, 0.3], [0.1, 0.5, 0.9]]) torch.testing.assert_close(combine_cubic_history_samples((sample,) * 5, fractions), sample) def test_validation_rejects_ambiguous_dtype_shape_and_flags(self): color = tensor([[0.5, 0.5, 0.5]]) for bad in (torch.zeros(1, 3, dtype=torch.int64), torch.zeros(1, 2), tensor([[float("nan"), 0, 0]])): with self.assertRaises(ValueError): normalize_color(bad) with self.assertRaises(ValueError): normalize_color(color, arithmetic="native_exact") with self.assertRaises(ValueError): conditioning_channels(color, controls=controls(style_id=-1), capabilities=CAPS) with self.assertRaises(ValueError): conditioning_channels(color, controls=controls(tone=tensor([1, 2])), capabilities=CAPS) with self.assertRaises(ValueError): conditioning_channels(color, controls=controls(), capabilities=CAPS, control_mask=torch.zeros(1, 3, dtype=torch.float32)) with self.assertRaises(ValueError): build_preblock_channels(color, noise_rgb=color, warped_history_rgb=None, history_enabled=0, controls=controls(), capabilities=CAPS) with self.assertRaises(ValueError): combine_cubic_history_samples((color,) * 4, tensor([[0.5, 0.5]])) def test_gradcheck_smooth_pre_and_controls(self): def function(color, history, noise, tone, structure, skin): return build_preblock_channels(color, noise_rgb=noise, warped_history_rgb=history, history_enabled=True, controls=controls(tone=tone, structure=structure, skin_structure=skin, auto_mask=True), capabilities=CAPS) args = (tensor([[0.3, 0.4, 0.7]], grad=True), tensor([[0.6, 0.2, 0.4]], grad=True), tensor([[0.1, -0.2, 0.4]], grad=True), tensor(0.7, grad=True), tensor(0.5, grad=True), tensor(0.3, grad=True)) self.assertTrue(gradcheck(function, args, eps=1e-6, atol=1e-5, rtol=1e-4)) def test_gradcheck_smooth_post_display_and_mask(self): def function(projected, color, history, scale, intensity, mask): post = reconstruct_post(projected, current_rgb=color, warped_history_rgb=history, history_enabled=True, learned_blend_scale=scale) display = compose_display(post.temporal_rgb, color, intensity=intensity, control_mask=mask) return torch.cat((post.temporal_rgb, display), dim=-1) args = (tensor([[0.1, -0.2, 0.3, 0.4]], grad=True), tensor([[0.3, 0.4, 0.5]], grad=True), tensor([[0.7, 0.6, 0.4]], grad=True), tensor(0.6, grad=True), tensor(0.7, grad=True), tensor([[0.8, 0.6, 0.4]], grad=True)) self.assertTrue(gradcheck(function, args, eps=1e-6, atol=1e-5, rtol=1e-4)) def test_gradcheck_explicit_cubic_samples(self): def function(fractions, *samples): return combine_cubic_history_samples(samples, fractions) args = (tensor([[0.3, 0.6]], grad=True),) + tuple( tensor([[0.1 + i * 0.1, 0.3, 0.6 - i * 0.05]], grad=True) for i in range(5)) self.assertTrue(gradcheck(function, args, eps=1e-6, atol=1e-5, rtol=1e-4)) def test_two_frame_gradient_reaches_previous_projection(self): first_projection = tensor([[0.1, 0.1, 0.1, 0]], grad=True) first = reconstruct_post(first_projection, current_rgb=tensor([[0.4, 0.4, 0.4]]), warped_history_rgb=None, history_enabled=False) second = reconstruct_post(tensor([[0.1, 0.1, 0.1, 0]]), current_rgb=tensor([[0.6, 0.6, 0.6]]), warped_history_rgb=first.temporal_rgb, history_enabled=True) second.temporal_rgb.sum().backward() torch.testing.assert_close(first_projection.grad, tensor([[0.125, 0.125, 0.125, 0]])) if __name__ == "__main__": unittest.main()