201 lines
7.6 KiB
Python
201 lines
7.6 KiB
Python
import unittest
|
|
|
|
import torch
|
|
|
|
from h3lora.adaln import (
|
|
AdalnContext, fit_basis, mark_stack_adaln, port_adaln_pairs,
|
|
port_attached_patches, sidecar_basis, upstream_adaln_fix_warning,
|
|
)
|
|
|
|
|
|
class AdalnSidecarTests(unittest.TestCase):
|
|
def test_fit_basis_recovers_affine(self):
|
|
g, k, d = 32, 4, 16
|
|
table = torch.randn(g, k)
|
|
v = torch.randn(d, k)
|
|
c = torch.randn(d)
|
|
grid = c.unsqueeze(0) + table @ v.T
|
|
v2, c2, residual = fit_basis(table, grid)
|
|
self.assertLess(residual, 1e-5)
|
|
self.assertTrue(torch.allclose(v2, v, atol=1e-4))
|
|
self.assertTrue(torch.allclose(c2, c, atol=1e-4))
|
|
|
|
def test_sidecar_v_c_preferred(self):
|
|
v = torch.randn(16, 4)
|
|
c = torch.randn(16)
|
|
got = sidecar_basis({"adaln_basis": v, "adaln_mean": c}, torch.randn(8, 4))
|
|
self.assertIsNotNone(got)
|
|
self.assertEqual(got[3], "checkpoint sidecars")
|
|
self.assertEqual(got[2], 0.0)
|
|
self.assertTrue(torch.equal(got[0], v.float()))
|
|
|
|
def test_context_uses_sidecars_without_grid_file(self):
|
|
g, k, d = 16, 3, 8
|
|
table = torch.randn(g, k)
|
|
v = torch.randn(d, k)
|
|
c = torch.randn(d)
|
|
ctx = AdalnContext(k, table, grid_path="", sidecars={"adaln_basis": v, "adaln_mean": c})
|
|
basis = ctx.basis()
|
|
self.assertIsNotNone(basis)
|
|
self.assertEqual(ctx.source, "checkpoint sidecars")
|
|
a = torch.randn(2, d)
|
|
b = torch.randn(6, 2)
|
|
sd = {
|
|
"blocks.0.adaln_proj.linear.lora_A.weight": a,
|
|
"blocks.0.adaln_proj.linear.lora_B.weight": b,
|
|
}
|
|
out, stats = port_adaln_pairs(sd, ctx)
|
|
self.assertEqual(stats["ported"], 1)
|
|
self.assertEqual(out["blocks.0.adaln_proj.linear.lora_A.weight"].shape[1], k)
|
|
|
|
def test_fit_basis_normalizes_mixed_devices(self):
|
|
if not torch.cuda.is_available():
|
|
self.skipTest("CUDA is not available")
|
|
table = torch.randn(16, 3, device="cuda")
|
|
grid = torch.randn(16, 8)
|
|
_v, _c, residual = fit_basis(table, grid)
|
|
self.assertTrue(torch.isfinite(torch.tensor(residual)))
|
|
|
|
def test_default_suffix_is_ported(self):
|
|
table = torch.randn(16, 3)
|
|
ctx = AdalnContext(
|
|
3,
|
|
table,
|
|
sidecars={"adaln_basis": torch.randn(8, 3), "adaln_mean": torch.randn(8)},
|
|
)
|
|
sd = {
|
|
"blocks.0.adaln_proj.linear.lora_A.default.weight": torch.randn(2, 8),
|
|
"blocks.0.adaln_proj.linear.lora_B.default.weight": torch.randn(6, 2),
|
|
}
|
|
out, stats = port_adaln_pairs(sd, ctx)
|
|
self.assertEqual(stats["ported"], 1)
|
|
self.assertIn("blocks.0.adaln_proj.linear.lora_A.weight", out)
|
|
self.assertIn("blocks.0.adaln_proj.linear.lora_B.weight", out)
|
|
|
|
def test_strip_drops_mismatched_pairs(self):
|
|
table = torch.randn(16, 3)
|
|
ctx = AdalnContext(
|
|
3,
|
|
table,
|
|
sidecars={"adaln_basis": torch.randn(8, 3), "adaln_mean": torch.randn(8)},
|
|
)
|
|
sd = {
|
|
"blocks.0.adaln_proj.linear.lora_A.weight": torch.randn(2, 8),
|
|
"blocks.0.adaln_proj.linear.lora_B.weight": torch.randn(6, 2),
|
|
"blocks.0.other.lora_A.weight": torch.randn(2, 4),
|
|
}
|
|
out, stats = port_adaln_pairs(sd, ctx, mode="strip")
|
|
self.assertEqual(stats["skipped"], 1)
|
|
self.assertNotIn("blocks.0.adaln_proj.linear.lora_A.weight", out)
|
|
self.assertIn("blocks.0.other.lora_A.weight", out)
|
|
|
|
|
|
class _Adapter:
|
|
def __init__(self, weights):
|
|
self.weights = weights
|
|
|
|
|
|
class _Patcher:
|
|
def __init__(self, patches, sd):
|
|
self.patches = {k: list(v) for k, v in patches.items()}
|
|
self._sd = sd
|
|
self.patches_uuid = None
|
|
|
|
def clone(self):
|
|
return _Patcher(self.patches, self._sd)
|
|
|
|
def model_state_dict(self):
|
|
return self._sd
|
|
|
|
def add_patches(self, patches, strength):
|
|
for key, value in patches.items():
|
|
self.patches.setdefault(key, []).append((strength, value, 1.0, None, None))
|
|
return list(patches)
|
|
|
|
|
|
class AttachedPatchTests(unittest.TestCase):
|
|
def _ctx(self, k=3, d=8, g=16):
|
|
table = torch.randn(g, k)
|
|
v = torch.randn(d, k)
|
|
c = torch.randn(d)
|
|
ctx = AdalnContext(k, table, sidecars={"adaln_basis": v, "adaln_mean": c})
|
|
return ctx, v, c
|
|
|
|
def test_ports_dense_lora_already_on_model(self):
|
|
ctx, v, c = self._ctx()
|
|
r, out = 2, 6
|
|
a = torch.randn(r, v.shape[0])
|
|
b = torch.randn(out, r)
|
|
key = "diffusion_model.blocks.0.adaln_proj.linear.weight"
|
|
bias_key = "diffusion_model.blocks.0.adaln_proj.linear.bias"
|
|
sd = {key: torch.randn(out, v.shape[1]), bias_key: torch.randn(out)}
|
|
patcher = _Patcher(
|
|
{key: [(0.5, _Adapter((b, a, None, None, None, None)), 1.0, None, None)]},
|
|
sd,
|
|
)
|
|
out_p, stats = port_attached_patches(patcher, ctx)
|
|
self.assertEqual(stats["ported"], 1)
|
|
self.assertEqual(stats["keys"], 1)
|
|
rebuilt = out_p.patches[key][-1][1]
|
|
self.assertEqual(tuple(rebuilt.weights[1].shape), (r, v.shape[1]))
|
|
self.assertIn(bias_key, out_p.patches)
|
|
expected_a = a.float() @ v.float()
|
|
self.assertTrue(torch.allclose(rebuilt.weights[1].float(), expected_a, atol=1e-4))
|
|
|
|
def test_strip_removes_mismatched_attached(self):
|
|
ctx, v, _c = self._ctx()
|
|
key = "diffusion_model.blocks.0.adaln_proj.linear.weight"
|
|
sd = {key: torch.randn(6, v.shape[1])}
|
|
patcher = _Patcher(
|
|
{key: [(1.0, _Adapter((torch.randn(6, 2), torch.randn(2, v.shape[0]),
|
|
None, None, None, None)), 1.0, None, None)]},
|
|
sd,
|
|
)
|
|
out_p, stats = port_attached_patches(patcher, ctx, mode="strip")
|
|
self.assertEqual(stats["stripped"], 1)
|
|
self.assertNotIn(key, out_p.patches)
|
|
|
|
def test_skips_locon_mid(self):
|
|
ctx, v, _c = self._ctx()
|
|
key = "diffusion_model.blocks.0.adaln_proj.linear.weight"
|
|
sd = {key: torch.randn(6, v.shape[1])}
|
|
locon = _Adapter((torch.randn(6, 2), torch.randn(2, v.shape[0]),
|
|
None, torch.randn(2, 2), None, None))
|
|
patcher = _Patcher({key: [(1.0, locon, 1.0, None, None)]}, sd)
|
|
_out_p, stats = port_attached_patches(patcher, ctx)
|
|
self.assertEqual(stats["ported"], 0)
|
|
self.assertEqual(stats["unportable"], 1)
|
|
|
|
|
|
class UpstreamFixWarningTests(unittest.TestCase):
|
|
def _patcher(self, n, marked=False):
|
|
patches = {}
|
|
for i in range(n):
|
|
wkey = f"diffusion_model.blocks.{i}.adaln_proj.linear.weight"
|
|
bkey = f"diffusion_model.blocks.{i}.adaln_proj.linear.bias"
|
|
w = torch.randn(6, 3)
|
|
b = torch.randn(6)
|
|
patches[wkey] = [(1.0, ("diff", (w,)), 1.0, None, None)]
|
|
patches[bkey] = [(1.0, ("diff", (b,)), 1.0, None, None)]
|
|
patcher = _Patcher(patches, {})
|
|
patcher.model_options = {}
|
|
if marked:
|
|
mark_stack_adaln(patcher)
|
|
return patcher
|
|
|
|
def test_warns_on_multiplied_out_adaln_diffs(self):
|
|
note = upstream_adaln_fix_warning(self._patcher(10))
|
|
self.assertIn("post-hoc AdaLN fix", note)
|
|
self.assertIn("Disable the extra AdaLN LoRA Fix node", note)
|
|
self.assertIn("adaln_port=off", note)
|
|
|
|
def test_silent_when_too_few_keys(self):
|
|
self.assertEqual(upstream_adaln_fix_warning(self._patcher(2)), "")
|
|
|
|
def test_silent_when_this_stack_already_marked(self):
|
|
self.assertEqual(upstream_adaln_fix_warning(self._patcher(10, marked=True)), "")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|