| |
| """CPU checks for high-CFG reconstruction from independently decoded FP64 z.""" |
|
|
| from io import BytesIO |
| from pathlib import Path |
| import sys |
| import unittest |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
|
|
| import torch |
|
|
| from transformer_flow import Model |
| from nf_roundtrip import closed_loop_forward, closed_loop_reverse, error_metrics |
|
|
|
|
| class ClosedLoopTests(unittest.TestCase): |
| def make_model(self, sos=True, patch_size=1, softplus=True): |
| torch.manual_seed(42) |
| model = Model(in_channels=4, img_size=4, patch_size=patch_size, channels=64, |
| num_blocks=3, layers_per_block=[1, 0, 2], head_dim=64, |
| rope=True, sos=sos, txt_size=4, txt_dim=16, |
| cond_top_only=True, use_softplus=softplus, |
| use_final_norm=True, soft_clip=4, seq_order='L2R').eval() |
| with torch.no_grad(): |
| for block in model.blocks: |
| if hasattr(block, 'proj_out'): |
| block.proj_out.weight.normal_(std=0.08) |
| block.proj_out.bias.normal_(std=0.05) |
| return model |
|
|
| @torch.no_grad() |
| def test_saved_z_independent_decode_and_feedback(self): |
| for sos, patch_size, softplus in ((True, 1, True), (False, 1, False), (True, 2, True)): |
| model = self.make_model(sos, patch_size, softplus) |
| x, y = torch.randn(2, 4, 4, 4), torch.randn(2, 4, 16) |
| x, y = x.bfloat16().float(), y.bfloat16().float() |
| for guidance in (0, 3.6, 7.2): |
| with self.subTest(sos=sos, patch_size=patch_size, softplus=softplus, cfg=guidance): |
| condition = torch.cat([y, torch.zeros_like(y)]) if guidance else y |
| z, states, simulated = closed_loop_forward(model, x, condition, guidance) |
| self.assertEqual(z.dtype, torch.float64) |
| self.assertTrue(all(p.dtype == torch.float32 for p in model.parameters())) |
| buffer = BytesIO() |
| torch.save(z, buffer) |
| buffer.seek(0) |
| loaded_z = torch.load(buffer, weights_only=True) |
| sequence = closed_loop_reverse(model, loaded_z, condition, guidance, return_sequence=True) |
| self.assertEqual(len(sequence), len(model.blocks) + 1) |
| torch.testing.assert_close(sequence[1].float(), simulated, atol=0, rtol=0) |
| torch.testing.assert_close(sequence[1], model.unpatchify(states[-2]).double(), |
| atol=1e-13, rtol=1e-13) |
| torch.testing.assert_close(sequence[-1], x, atol=3e-5, rtol=3e-5) |
| repeated = closed_loop_reverse(model, loaded_z, condition, guidance) |
| torch.testing.assert_close(repeated, sequence[-1], atol=0, rtol=0) |
| perturbed = closed_loop_reverse(model, loaded_z + 0.1, condition, guidance) |
| self.assertGreater((perturbed - repeated).abs().mean().item(), 0.01) |
|
|
| def test_rejects_truncated_z_and_low_precision_weights(self): |
| model = self.make_model() |
| x, condition = torch.zeros(1, 4, 4, 4), torch.zeros(2, 4, 16) |
| with self.assertRaisesRegex(ValueError, 'original FP64 z'): |
| closed_loop_reverse(model, x, condition, 3.6) |
| with self.assertRaisesRegex(ValueError, 'FP32 model weights'): |
| closed_loop_forward(model.bfloat16(), x, condition, 3.6) |
|
|
| def test_metrics_preserve_fp64_errors(self): |
| metrics = error_metrics(torch.tensor([1 + 1e-12], dtype=torch.float64), |
| torch.ones(1, dtype=torch.float64)) |
| self.assertGreater(metrics['mae'], 1e-13) |
| self.assertLess(metrics['mae'], 2e-12) |
|
|
|
|
| if __name__ == '__main__': |
| torch.set_num_threads(2) |
| unittest.main() |
|
|
