| |
| """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() |
|
|
