#!/usr/bin/env python3 """Numerical checks for fixed-reference NF score editing (CPU only).""" from contextlib import nullcontext, redirect_stdout from dataclasses import replace import io import json import os from pathlib import Path import sys import tempfile from types import SimpleNamespace import unittest from unittest.mock import patch import torch import nf_reference_score_editing as editing class TriangularFlow(torch.nn.Module): def forward(self, u, condition): scale, mean, slope = condition.unbind(1) x, y = u.unbind(1) z = torch.stack((x * (-scale).exp(), (y - mean) * (-slope * x).exp()), 1) return z, None, [], [-scale[:, None], (-slope * x)[:, None]] class ShiftFlow(torch.nn.Module): def __init__(self): super().__init__() self.queries = [] def forward(self, u, condition): self.queries.append((u.detach().clone(), condition.detach().clone(), u.is_leaf)) shift = condition.reshape(-1, *([1] * (u.ndim - 1))) return u - shift, None, [], [torch.zeros_like(u)] class ReferenceScoreTests(unittest.TestCase): @classmethod def setUpClass(cls): torch.set_num_threads(1) def test_full_density_and_nonconstant_logdet_score(self): model = TriangularFlow().eval().requires_grad_(False) u = torch.tensor([[0.4, -0.3], [-0.2, 0.7]]) condition = torch.tensor([[0.2, -0.1, 0.3], [-0.1, 0.8, -0.4]]) scale, mean, slope = condition.double().unbind(1) x, y = u.double().unbind(1) expected_score = torch.stack(( -x * (-2 * scale).exp() + slope * (y - mean).square() * (-2 * slope * x).exp() - slope, -(y - mean) * (-2 * slope * x).exp(), ), 1).float() expected_density = (torch.distributions.Normal(0, scale.exp()).log_prob(x) + torch.distributions.Normal(mean, (slope * x).exp()).log_prob(y)) for checkpointing in (False, True): score, density = editing.conditional_score(model, u, condition, checkpointing) torch.testing.assert_close(score, expected_score) torch.testing.assert_close(density[:, 0], expected_density) self.assertFalse(score.requires_grad) self.assertFalse(u.requires_grad) def test_fixed_source_noise_and_analytic_multiple_updates(self): model = ShiftFlow() source = torch.tensor([[0.7, -0.5]]) cs, ct = torch.tensor([[0.2]]), torch.tensor([[1.3]]) noise = torch.tensor([[1.2, -0.4], [-0.8, 0.6], [0.3, 0.9]]) original_noise = noise.clone() config = editing.EditingConfig(num_steps=4, step_size=0.2, noise_batch_size=2, gradient_checkpointing=False) edited, cached, result, trace = editing.edit_latent(model, source, cs, ct, noise, config, False) expected = source + (ct - cs) * (1 - (1 - config.step_size) ** config.num_steps) torch.testing.assert_close(edited, expected) torch.testing.assert_close(cached, cs - source - config.noise_std * noise.mean(0, keepdim=True)) torch.testing.assert_close(noise, original_noise) torch.testing.assert_close(source, torch.tensor([[0.7, -0.5]])) self.assertEqual(result['nf_score_calls'], 10) self.assertEqual(len(model.queries), 10) self.assertEqual(len(trace), 4) for index, (query, condition, is_leaf) in enumerate(model.queries): group, chunk = divmod(index, 2) state = source if group == 0 else source + (ct - cs) * (1 - 0.8 ** (group - 1)) expected_query = state + config.noise_std * noise[chunk * 2:(chunk + 1) * 2] torch.testing.assert_close(query, expected_query) torch.testing.assert_close(condition, (cs if group == 0 else ct).expand_as(condition)) self.assertTrue(is_leaf) self.assertFalse(edited.requires_grad) json.dumps(result, allow_nan=False) json.dumps(trace, allow_nan=False) def test_noise_batching_and_objective_gradient(self): model = TriangularFlow() source = torch.tensor([[0.1, -0.2]]) current = torch.tensor([[0.4, 0.3]]) noise = torch.tensor([[0.2, -0.5], [-0.4, 0.8], [0.9, 0.1], [-0.3, -0.6], [0.6, 0.4]]) cs, ct = torch.tensor([[0.2, -0.1, 0.3]]), torch.tensor([[-0.1, 0.8, -0.4]]) config = editing.EditingConfig(gradient_checkpointing=False) gs, ds = editing.mean_score(model, source, cs, noise, config) gt, dt = editing.mean_score(model, current, ct, noise, config) for batch_size in (2, 5): actual, density = editing.mean_score(model, current, ct, noise, replace(config, noise_batch_size=batch_size)) torch.testing.assert_close(actual, gt) torch.testing.assert_close(density, dt) direction = torch.tensor([[0.7, -0.4]]) objectives = [] for point in (current - 1e-3 * direction, current + 1e-3 * direction): density = editing.conditional_density(model, point + config.noise_std * noise, ct.expand(5, -1))[:, 0].mean() objectives.append((density - ds[0] - (gs.double() * (point - source).double()).sum()).item()) derivative = (objectives[1] - objectives[0]) / 2e-3 self.assertAlmostEqual(derivative, ((gt - gs) * direction).sum().item(), delta=1e-4) def test_identity_and_zero_steps_preserve_latent(self): source = torch.tensor([[0.8, -0.2]]) cs, ct = torch.tensor([[0.3]]), torch.tensor([[1.1]]) noise = torch.randn(3, 2, generator=torch.Generator().manual_seed(2)) for num_steps, target in ((0, ct), (4, cs)): config = editing.EditingConfig(num_steps=num_steps) edited, _, result, trace = editing.edit_latent(ShiftFlow(), source, cs, target, noise, config, False) torch.testing.assert_close(edited, source, rtol=0, atol=0) self.assertEqual(result['final_displacement_rms'], 0) self.assertEqual(len(trace), num_steps) def test_starflow_jacobian_and_checkpointed_input_gradient(self): from transformer_flow import Model torch.manual_seed(7) model = Model(in_channels=1, img_size=2, patch_size=1, channels=32, num_blocks=2, layers_per_block=[1, 1], head_dim=32, rope=True, sos=True, txt_size=2, txt_dim=3, cond_top_only=True, use_softplus=True) with torch.no_grad(): for block in model.blocks: block.proj_out.weight.normal_(std=0.1) model.eval().requires_grad_(False) u, condition = torch.randn(1, 1, 2, 2), torch.randn(1, 2, 3) score, density = editing.conditional_score(model, u, condition, False) checkpointed, _ = editing.conditional_score(model, u, condition, True) torch.testing.assert_close(score, checkpointed) jacobian = torch.autograd.functional.jacobian( lambda flat: model(flat.reshape_as(u), condition)[0].flatten(), u.flatten()) _, logdet = torch.linalg.slogdet(jacobian.double()) torch.testing.assert_close(density[0, 2], logdet, rtol=1e-5, atol=1e-6) direction = torch.randn_like(u) low = editing.conditional_density(model, u - 1e-3 * direction, condition)[0, 0] high = editing.conditional_density(model, u + 1e-3 * direction, condition)[0, 0] self.assertAlmostEqual(((high - low) / 2e-3).item(), (score * direction).sum().item(), delta=5e-4) self.assertTrue(all(parameter.grad is None for parameter in model.parameters())) self.assertFalse(model.training) def test_single_image_pipeline_and_vae_identity_decode(self): from PIL import Image import torchvision.transforms class FakeVAE(torch.nn.Module): def __init__(self): super().__init__() self.scaling_factor, self.downsample_factor = 1, 1 self.vae = SimpleNamespace(encode=lambda image: SimpleNamespace( latent_dist=SimpleNamespace(mode=lambda: image * 0.5))) def decode(self, latent): return latent for identity in (False, True): with self.subTest(identity=identity), tempfile.TemporaryDirectory() as temporary: root = Path(temporary) image_path, output = root / 'input.v1.png', root / 'output' Image.new('RGB', (4, 3), (200, 100, 80)).save(image_path) components = (SimpleNamespace(img_size=2, txt_size=4, noise_std=0.3, noise_type='gaussian', cond_noise_level=0), ShiftFlow(), FakeVAE(), None, torch.nn.Identity(), SimpleNamespace(distributed=False), torch.device('cpu')) mocks = { 'pie_bench_nf_components': SimpleNamespace(load_components=lambda cli: components), 'utils': SimpleNamespace(encode_text=lambda *args, **kwargs: torch.tensor( [[0.0], [0.0 if identity else 1.0]])), } argv = ['editing', '--image-path', str(image_path), '--source-prompt', 'source', '--target-prompt', 'source' if identity else 'target', '--output-dir', str(output), '--num-steps', '2', '--num-noise-samples', '3', '--step-size', '0.1'] with patch.dict(sys.modules, mocks), patch.object(sys, 'argv', argv), \ patch.dict(os.environ, {'WORLD_SIZE': '1'}), \ patch.object(torch.cuda, 'is_available', return_value=True), \ patch.object(torch.cuda, 'get_device_name', return_value='CPU test'), \ patch.object(torch, 'autocast', side_effect=lambda *a, **k: nullcontext()), \ redirect_stdout(io.StringIO()): editing.main() with self.assertRaises(FileExistsError): editing.main() tensors = torch.load(output / 'latents/single/input.v1.pt', weights_only=True) expected = tensors['source_latent'] + (0 if identity else 0.19) torch.testing.assert_close(tensors['edited_latent'], expected) self.assertEqual(tensors['noise'].shape, (3, 3, 2, 2)) config = json.loads((output / 'run_config.json').read_text()) self.assertEqual(config['editing']['noise_std'], 0.3) summary = json.loads((output / 'summary.json').read_text()) self.assertEqual(summary['count'], 1) self.assertEqual(summary['images'][0]['nf_score_calls'], 9) self.assertEqual(len((output / 'trajectory/single/input.v1.jsonl').read_text().splitlines()), 2) if identity: edited = (output / 'editing/single/input.v1.png').read_bytes() self.assertEqual(edited, (output / 'vae_only/single/input.v1.png').read_bytes()) self.assertNotEqual(edited, (output / 'source/single/input.v1.png').read_bytes()) if __name__ == '__main__': unittest.main()