File size: 3,881 Bytes
2cf4a54 | 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 | #!/usr/bin/env python3
"""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()
|