File size: 13,186 Bytes
2f3db25 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 | #!/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()
|