| |
| """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]])) |
| |
| 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_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)) |
| |
| 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() |
|
|
