52 lines
2.2 KiB
Python
52 lines
2.2 KiB
Python
import unittest
|
|
|
|
import torch
|
|
|
|
from h3_blackwell_runtime.latent_upscaler import H3LatentResizer3D, _checkpoint_config, upscale_h3_latent
|
|
from h3_blackwell_runtime.runtime import normalize_upscale
|
|
|
|
|
|
class LatentUpscalerContracts(unittest.TestCase):
|
|
def test_request_upscale_normalization(self):
|
|
self.assertIsNone(normalize_upscale(None))
|
|
self.assertIsNone(normalize_upscale("none"))
|
|
self.assertIsNone(normalize_upscale(1))
|
|
self.assertEqual(normalize_upscale("2"), 2.0)
|
|
with self.assertRaisesRegex(ValueError, "greater than 1.0"):
|
|
normalize_upscale(0.5)
|
|
with self.assertRaisesRegex(ValueError, "at most 4.0"):
|
|
normalize_upscale(4.1)
|
|
with self.assertRaisesRegex(ValueError, "must be a number"):
|
|
normalize_upscale({"scale": 2})
|
|
|
|
def test_3d_model_preserves_time_and_scales_space(self):
|
|
model = H3LatentResizer3D(in_blocks=1, out_blocks=1, channels=32, dropout=0, temporal_every=0).eval()
|
|
latent = torch.randn(1, 24, 3, 2, 4)
|
|
output = upscale_h3_latent(model, latent, scale=2)
|
|
self.assertEqual(output.shape, (1, 24, 3, 4, 8))
|
|
|
|
def test_invalid_latent_shape_is_rejected(self):
|
|
model = H3LatentResizer3D(in_blocks=1, out_blocks=1, channels=32, temporal_every=0).eval()
|
|
with self.assertRaisesRegex(ValueError, r"\[B,24,T,H,W\]"):
|
|
upscale_h3_latent(model, torch.randn(1, 16, 3, 2, 4))
|
|
|
|
def test_scale_one_preserves_latent_exactly(self):
|
|
model = H3LatentResizer3D(in_blocks=1, out_blocks=1, channels=32, temporal_every=0).eval()
|
|
latent = torch.randn(1, 24, 3, 2, 4)
|
|
self.assertTrue(torch.equal(upscale_h3_latent(model, latent, scale=1), latent))
|
|
|
|
def test_checkpoint_architecture_detection_matches_module_layout(self):
|
|
model = H3LatentResizer3D(in_blocks=3, out_blocks=2, channels=32, temporal_every=2, temporal_kernel=3)
|
|
config = _checkpoint_config(model.state_dict())
|
|
self.assertEqual(config, {
|
|
"in_channels": 24,
|
|
"in_blocks": 3,
|
|
"out_blocks": 2,
|
|
"channels": 32,
|
|
"temporal_every": 2,
|
|
"temporal_kernel": 3,
|
|
})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|